import { useState, useEffect, useCallback, useRef, useMemo, type JSX } from "react";
import { z } from "zod";
import { CapturedLogSchema, type CapturedLog } from "../proxy/schemas";
import { useStripConfig } from "../lib/useStripConfig";
import { OnboardingBanner } from "./OnboardingBanner";
import { ProxyViewer } from "./ProxyViewer";

type SSEUpdate =
  | {
      type: "init";
      logs: CapturedLog[];
    }
  | {
      type: "update";
      log: CapturedLog;
    };

const SSEUpdateSchema = z.union([
  z.object({
    type: z.literal("init"),
    logs: z.array(CapturedLogSchema),
  }),
  z.object({
    type: z.literal("update"),
    log: CapturedLogSchema,
  }),
]);

function extractSessions(logs: CapturedLog[]): string[] {
  const set = new Set<string>();
  for (const l of logs) {
    if (l.sessionId !== null && l.sessionId !== "") set.add(l.sessionId);
  }
  return [...set];
}

function extractModels(logs: CapturedLog[]): string[] {
  const set = new Set<string>();
  for (const l of logs) {
    if (l.model !== null && l.model !== "") set.add(l.model);
  }
  return [...set];
}

/**
 * Filter logs by selected session/model. Both `__all__` and missing values
 * pass through. Returns the input array reference when no log matches the
 * filter, so referential equality is preserved for empty results.
 */
function filterLogs(
  logs: CapturedLog[],
  selectedSession: string,
  selectedModel: string,
): CapturedLog[] {
  if (selectedSession === "__all__" && selectedModel === "__all__") return logs;
  return logs.filter((l) => {
    if (selectedSession !== "__all__" && l.sessionId !== selectedSession) return false;
    if (selectedModel !== "__all__" && l.model !== selectedModel) return false;
    return true;
  });
}

const DEBOUNCE_MS = 50;

function buildLogsStreamUrl(sessionId: string | undefined): string {
  if (sessionId === undefined) return "/api/logs/stream";
  const params = new URLSearchParams({ sessionId });
  return `/api/logs/stream?${params.toString()}`;
}

export function ProxyViewerContainer({
  initialSessionId,
}: {
  /**
   * Initial session filter. When the route is `/session/$id`, pass the URL
   * param here so the page opens already scoped to that session. Default
   * `__all__` keeps the existing "all sessions" behavior on `/`.
   *
   * NOTE: the value is only used as the initial state on mount. To pick up
   * a new value (e.g. user navigates to a different `/session/$id`), the
   * caller should also pass a `key` prop matching the session id, which
   * forces a remount.
   */
  initialSessionId?: string;
} = {}): JSX.Element {
  // `allLogs` is the unfiltered set populated by the SSE. The single SSE
  // connection never re-opens on filter change — we always carry the full
  // set and derive the displayed view with `useMemo` below.
  const [allLogs, setAllLogs] = useState<CapturedLog[]>([]);
  const [selectedSession, setSelectedSession] = useState(initialSessionId ?? "__all__");
  const [selectedModel, setSelectedModel] = useState("__all__");
  const [viewMode, setViewMode] = useState<"simple" | "full">("simple");
  const [error, setError] = useState<string | null>(null);
  const eventSourceRef = useRef<EventSource | null>(null);
  const reconnectTimeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null);

  // O(1) log lookup by id
  const logIndexRef = useRef<Map<number, number>>(new Map());

  // Debounce buffer for SSE updates
  const pendingUpdatesRef = useRef<CapturedLog[]>([]);
  const flushTimerRef = useRef<ReturnType<typeof setTimeout> | null>(null);

  // Derived view: the logs the user actually sees. Recomputed only when the
  // underlying set or filter changes — never on every SSE message.
  const logs = useMemo(
    () => filterLogs(allLogs, selectedSession, selectedModel),
    [allLogs, selectedSession, selectedModel],
  );
  const sessions = useMemo(() => extractSessions(allLogs), [allLogs]);
  const models = useMemo(() => extractModels(allLogs), [allLogs]);

  const flushUpdates = useCallback(() => {
    flushTimerRef.current = null;
    const updates = pendingUpdatesRef.current;
    pendingUpdatesRef.current = [];
    if (updates.length === 0) return;

    setAllLogs((prev) => {
      let next = prev;
      for (const log of updates) {
        const idx = logIndexRef.current.get(log.id);
        if (idx !== undefined) {
          // Replace existing entry
          next = [...next.slice(0, idx), log, ...next.slice(idx + 1)];
        } else {
          logIndexRef.current.set(log.id, next.length);
          next = [...next, log];
        }
      }
      return next;
    });
  }, []);

  const scheduleUpdate = useCallback(
    (log: CapturedLog) => {
      pendingUpdatesRef.current.push(log);
      if (flushTimerRef.current === null) {
        flushTimerRef.current = setTimeout(flushUpdates, DEBOUNCE_MS);
      }
    },
    [flushUpdates],
  );

  const connectSSE = useCallback(() => {
    // Clean up existing connection
    if (eventSourceRef.current) {
      eventSourceRef.current.close();
    }

    const es = new EventSource(buildLogsStreamUrl(initialSessionId));
    eventSourceRef.current = es;

    es.onmessage = (event: MessageEvent) => {
      try {
        const rawData = String(event.data);
        const parsed: unknown = JSON.parse(rawData);
        const updateResult = SSEUpdateSchema.safeParse(parsed);
        if (!updateResult.success) {
          setError("Failed to parse SSE data");
          return;
        }
        const update = updateResult.data;
        if (update.type === "init") {
          // Flush any pending debounce
          if (flushTimerRef.current !== null) {
            clearTimeout(flushTimerRef.current);
            flushTimerRef.current = null;
          }
          pendingUpdatesRef.current = [];

          // Build fresh index
          const idx = new Map<number, number>();
          for (let i = 0; i < update.logs.length; i++) {
            const log = update.logs[i];
            if (log !== undefined) idx.set(log.id, i);
          }
          logIndexRef.current = idx;
          setAllLogs(update.logs);
          setError(null);
        } else if (update.type === "update") {
          scheduleUpdate(update.log);
        }
      } catch {
        setError("Failed to parse SSE data");
      }
    };

    es.onerror = () => {
      setError("SSE connection lost, reconnecting...");
      es.close();
      if (reconnectTimeoutRef.current !== null) {
        clearTimeout(reconnectTimeoutRef.current);
      }
      reconnectTimeoutRef.current = setTimeout(connectSSE, 3000);
    };
  }, [initialSessionId, scheduleUpdate]);

  useEffect(() => {
    connectSSE();
    return () => {
      if (eventSourceRef.current) {
        eventSourceRef.current.close();
        eventSourceRef.current = null;
      }
      if (reconnectTimeoutRef.current !== null) {
        clearTimeout(reconnectTimeoutRef.current);
        reconnectTimeoutRef.current = null;
      }
      if (flushTimerRef.current !== null) {
        clearTimeout(flushTimerRef.current);
        flushTimerRef.current = null;
      }
    };
  }, [connectSSE]);

  const handleClearAll = useCallback(() => {
    if (initialSessionId !== undefined && allLogs.length === 0) return;
    void (async () => {
      try {
        const body =
          initialSessionId === undefined
            ? undefined
            : JSON.stringify({ ids: allLogs.map((log) => log.id) });
        const res = await fetch("/api/logs", {
          method: "DELETE",
          ...(body === undefined
            ? {}
            : {
                headers: { "Content-Type": "application/json" },
                body,
              }),
        });
        if (!res.ok) {
          setError("Failed to clear logs");
          return;
        }
        logIndexRef.current.clear();
        setAllLogs([]);
        setError(null);
      } catch (err) {
        setError(err instanceof Error ? err.message : "Unknown error clearing logs");
      }
    })();
  }, [allLogs, initialSessionId]);

  const handleClearGroup = useCallback((ids: number[]) => {
    if (ids.length === 0) return;
    void (async () => {
      try {
        const res = await fetch("/api/logs", {
          method: "DELETE",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify({ ids }),
        });
        if (!res.ok) {
          setError("Failed to clear group");
          return;
        }
        const idSet = new Set(ids);
        setAllLogs((prev) => {
          const remaining = prev.filter((l) => !idSet.has(l.id));
          // Rebuild index
          const idx = new Map<number, number>();
          for (let i = 0; i < remaining.length; i++) {
            const log = remaining[i];
            if (log !== undefined) idx.set(log.id, i);
          }
          logIndexRef.current = idx;
          return remaining;
        });
        setError(null);
      } catch (err) {
        setError(err instanceof Error ? err.message : "Unknown error clearing group");
      }
    })();
  }, []);

  // Read the strip config once at the container so the virtualized list does
  // not need N independent SWR subscriptions per row.
  const { strip, slowResponseThresholdSeconds } = useStripConfig();

  return (
    <>
      {error !== null && (
        <div className="fixed top-4 right-4 bg-destructive text-destructive-foreground px-4 py-2 rounded-md text-sm z-50">
          {error}
        </div>
      )}
      <OnboardingBanner />
      <ProxyViewer
        logs={logs}
        sessions={sessions}
        models={models}
        selectedSession={selectedSession}
        selectedModel={selectedModel}
        onSessionChange={setSelectedSession}
        onModelChange={setSelectedModel}
        onClearAll={handleClearAll}
        onClearGroup={handleClearGroup}
        viewMode={viewMode}
        onViewModeChange={setViewMode}
        strip={strip}
        slowResponseThresholdSeconds={slowResponseThresholdSeconds}
        // Session filter is the URL's job when `initialSessionId` was given.
        hideSessionFilter={initialSessionId !== undefined}
        pinnedSessionId={initialSessionId}
      />
    </>
  );
}
