import { useQuery, useQueryClient } from "@tanstack/react-query";
import { useEffect } from "react";
import { z } from "zod";
import { errorMessage } from "~/lib/error-message";

const logRow = z.object({ seq: z.number(), step: z.number(), text: z.string() });
const snapshot = z.object({ rows: z.array(logRow), done: z.boolean() });
const event = z.discriminatedUnion("type", [
  z.object({ type: z.literal("log"), seq: z.number(), step: z.number(), text: z.string() }),
  z.object({ type: z.literal("step"), index: z.number(), status: z.string(), conclusion: z.string().nullable() }),
  z.object({ type: z.literal("done"), conclusion: z.string().nullable() }),
]);

export type LogRow = z.infer<typeof logRow>;
type Snapshot = z.infer<typeof snapshot>;
type State = { rows: LogRow[]; live: boolean; error: string | null };

const NO_ROWS: LogRow[] = [];
/** A finished job's log is immutable; a running one is kept current by the event stream, not by refetching. */
const LOG_STALE_MS = Number.POSITIVE_INFINITY;
/** Safety net under the stream (a dropped connection, a missed event): re-read the job while it isn't done. */
const LIVE_RESYNC_MS = 30_000;

export const jobLogQueryKey = (jobId: string) => ["job-log", jobId] as const;

async function fetchSnapshot(jobId: string): Promise<Snapshot> {
  const response = await fetch(`/api/jobs/${jobId}/logs`);
  if (!response.ok) throw new Error(`Failed to load logs (${response.status})`);
  return snapshot.parse(await response.json());
}

/**
 * Loads a job's log and, until the job finishes (queued jobs included), follows the Runner Durable Object's
 * server-sent event stream: log lines append in place, and step/start/finish events refresh the page data.
 */
export function useJobLog(jobId: string, onStepChange: () => void): State {
  const queryClient = useQueryClient();
  const query = useQuery({
    queryKey: jobLogQueryKey(jobId),
    queryFn: () => fetchSnapshot(jobId),
    staleTime: LOG_STALE_MS,
  });
  const live = query.data ? !query.data.done : false;

  useEffect(() => {
    if (!live) return;
    const key = jobLogQueryKey(jobId);
    const after = queryClient.getQueryData<Snapshot>(key)?.rows.at(-1)?.seq ?? 0;
    const source = new EventSource(`/api/jobs/${jobId}/logs/stream?after=${after}`);
    source.onmessage = (message) => {
      const parsed = event.safeParse(JSON.parse(message.data));
      if (!parsed.success) return;
      const payload = parsed.data;
      if (payload.type === "log") {
        queryClient.setQueryData<Snapshot>(key, (current) => {
          if (!current || current.rows.some((row) => row.seq === payload.seq)) return current;
          return { ...current, rows: [...current.rows, { seq: payload.seq, step: payload.step, text: payload.text }] };
        });
      } else if (payload.type === "step") {
        onStepChange();
      } else {
        queryClient.setQueryData<Snapshot>(key, (current) => (current ? { ...current, done: true } : current));
        source.close();
        onStepChange();
      }
    };
    const resync = setInterval(() => {
      void queryClient.invalidateQueries({ queryKey: key });
      onStepChange();
    }, LIVE_RESYNC_MS);
    return () => {
      clearInterval(resync);
      source.close();
    };
  }, [jobId, live, onStepChange, queryClient]);

  return {
    rows: query.data?.rows ?? NO_ROWS,
    live,
    error: query.isError ? errorMessage(query.error) : null,
  };
}