import { memo, useState, useEffect, useMemo, type JSX } from "react";
import { ChevronDown, ChevronRight, Loader2 } from "lucide-react";
import { Badge } from "../ui/badge";
import { JsonViewer } from "../ui/json-viewer";
import { Tooltip, TooltipTrigger, TooltipContent, TooltipProvider } from "../ui/tooltip";
import type { StreamingChunk } from "../../proxy/schemas";

export type StreamingChunkSequenceProps = {
  logId: number;
  truncated?: boolean;
};

type ChunkState =
  | { status: "idle" }
  | { status: "loading" }
  | { status: "success"; chunks: StreamingChunk[] }
  | { status: "error"; message: string };

type ChunkGroup = {
  index: number;
  chunks: StreamingChunk[];
};

export const StreamingChunkSequence = memo(function StreamingChunkSequence({
  logId,
  truncated,
}: StreamingChunkSequenceProps): JSX.Element {
  const [containerExpanded, setContainerExpanded] = useState(false);
  const [chunkState, setChunkState] = useState<ChunkState>({ status: "idle" });
  const [expandedIndices, setExpandedIndices] = useState<Set<number>>(new Set());

  useEffect(() => {
    if (!containerExpanded || chunkState.status !== "idle") return;

    let cancelled = false;
    setChunkState({ status: "loading" });

    fetch(`/api/logs/${logId}/chunks`)
      .then((res) => {
        if (!res.ok) {
          return Promise.reject(new Error("Chunks not found"));
        }
        return res.json();
      })
      .then((data: { chunks: StreamingChunk[]; truncated?: boolean }) => {
        if (!cancelled) {
          setChunkState({ status: "success", chunks: data.chunks });
        }
      })
      .catch(() => {
        if (!cancelled) {
          setChunkState({ status: "error", message: "Chunk data unavailable" });
        }
      });

    return () => {
      cancelled = true;
    };
  }, [containerExpanded, logId]);

  const groups = useMemo<ChunkGroup[]>(() => {
    if (chunkState.status !== "success") return [];

    const map = new Map<number, StreamingChunk[]>();
    for (const chunk of chunkState.chunks) {
      const existing = map.get(chunk.index);
      if (existing) {
        existing.push(chunk);
      } else {
        map.set(chunk.index, [chunk]);
      }
    }

    return Array.from(map.entries())
      .map(([index, chunks]) => ({ index, chunks }))
      .sort((a, b) => a.index - b.index);
  }, [chunkState]);

  const toggleIndex = (index: number) => {
    setExpandedIndices((prev) => {
      const next = new Set(prev);
      if (next.has(index)) {
        next.delete(index);
      } else {
        next.add(index);
      }
      return next;
    });
  };

  function renderBody(): JSX.Element {
    if (chunkState.status === "idle" || chunkState.status === "loading") {
      return (
        <div className="flex items-center gap-2 py-2 text-xs text-muted-foreground">
          <Loader2 className="size-3 animate-spin" />
          <span>Loading chunks...</span>
        </div>
      );
    }

    if (chunkState.status === "error") {
      return <div className="py-2 text-xs text-muted-foreground italic">{chunkState.message}</div>;
    }

    return (
      <div className="p-2 space-y-1">
        <div className="text-[10px] text-muted-foreground font-mono mb-2">
          {groups.length} index group{groups.length !== 1 ? "s" : ""} available
        </div>
        {groups.map((group) => {
          const isExpanded = expandedIndices.has(group.index);
          return (
            <div key={group.index} className="rounded border border-border bg-background">
              <button
                type="button"
                className="flex items-center gap-2 w-full px-2 py-1.5 text-left hover:bg-muted/50 transition-colors cursor-pointer"
                onClick={() => toggleIndex(group.index)}
              >
                {isExpanded ? (
                  <ChevronDown className="size-3 text-muted-foreground" />
                ) : (
                  <ChevronRight className="size-3 text-muted-foreground" />
                )}
                <span className="text-[10px] text-muted-foreground font-mono">
                  [{group.index}] {group.chunks[0]?.type ?? ""}
                </span>
                <span className="text-[10px] text-muted-foreground font-mono">
                  {group.chunks.length} chunk{group.chunks.length !== 1 ? "s" : ""}
                </span>
              </button>
              {isExpanded && (
                <div className="px-2 pb-2 space-y-1">
                  {group.chunks.map((chunk) => (
                    <div key={chunk.index} className="rounded border border-border bg-muted/20 p-2">
                      <div className="flex items-center gap-2 mb-1">
                        <span className="text-[10px] text-muted-foreground font-mono">
                          +{chunk.timestamp}ms
                        </span>
                        <span className="text-[10px] text-muted-foreground font-mono">
                          {chunk.type}
                        </span>
                      </div>
                      <JsonViewer data={chunk} defaultExpandDepth={0} showCopy />
                    </div>
                  ))}
                </div>
              )}
            </div>
          );
        })}
      </div>
    );
  }

  return (
    <div className="space-y-1">
      <TooltipProvider>
        <Tooltip>
          <TooltipTrigger asChild>
            <button
              type="button"
              className="flex items-center gap-1.5 text-xs text-muted-foreground hover:text-foreground transition-colors cursor-pointer"
              onClick={() => setContainerExpanded((v) => !v)}
            >
              {containerExpanded ? (
                <ChevronDown className="size-3" />
              ) : (
                <ChevronRight className="size-3" />
              )}
              <span>Raw SSE Events</span>
              <Badge variant="outline" className="text-[9px] px-1 py-0 h-4 font-mono ml-1">
                {logId}
                {truncated === true ? "+" : ""}
              </Badge>
            </button>
          </TooltipTrigger>
          <TooltipContent>Server-Sent Events streaming chunks from the provider</TooltipContent>
        </Tooltip>
      </TooltipProvider>

      {containerExpanded === true ? (
        <div className="rounded-md border border-border bg-muted/20 overflow-auto max-h-64">
          {renderBody()}
        </div>
      ) : null}
    </div>
  );
});
