/**
 * @file useRunStream.ts
 * @description Owns the Run page's live stream-json state. Subscribes to the
 * WebSocket event bus and folds every `run_stream` envelope for one run id into
 * an envelope array (`mergeEnvelope` and friends, moved here verbatim from
 * `pages/Run.tsx`), exposes the typewriter-smoothed view of that array, and
 * hands `run_status` / `run_input_ack` back to the caller — the page still owns
 * the `RunHandle` and the run-list refresh, so those arrive as callbacks.
 *
 * The stream-json envelope types live here too, since this hook is what
 * produces them; the page imports them for rendering.
 *
 * @author Nguyễn Ngọc Trí Vĩ <vinnt@smartgift.vn>
 */

import { useEffect, useMemo, useRef, useState } from "react";
import { flushSync } from "react-dom";
import { eventBus } from "../lib/eventBus";
import type {
  RunInputAckPayload,
  RunStatusPayload,
  RunStreamPayload,
  WSMessage,
} from "../lib/types";

// ── Stream-json envelope shapes (the bits we render) ──────────────────

export type ContentBlock =
  | { type: "text"; text: string }
  | { type: "thinking"; thinking?: string }
  | { type: "tool_use"; id: string; name: string; input: unknown }
  | { type: "tool_result"; tool_use_id: string; content: unknown; is_error?: boolean };

export interface AssistantMessage {
  type: "assistant";
  message?: {
    content?: ContentBlock[] | string;
    usage?: {
      input_tokens?: number;
      output_tokens?: number;
      cache_read_input_tokens?: number;
      cache_creation_input_tokens?: number;
    };
  };
}
export interface UserMessage {
  type: "user";
  message?: { content?: ContentBlock[] | string };
}
export interface SystemInit {
  type: "system";
  subtype: "init";
  session_id?: string;
  model?: string;
  cwd?: string;
  tools?: string[];
  permissionMode?: string;
}
export interface ResultEnvelope {
  type: "result";
  subtype?: string;
  is_error?: boolean;
  duration_ms?: number;
  duration_api_ms?: number;
  num_turns?: number;
  result?: string;
  session_id?: string;
  total_cost_usd?: number;
  usage?: { input_tokens?: number; output_tokens?: number };
}
export type Envelope =
  | AssistantMessage
  | UserMessage
  | SystemInit
  | ResultEnvelope
  | { type: string; [k: string]: unknown };

// ── Streaming envelope merge ───────────────────────────────────────────
//
// `claude --output-format stream-json --include-partial-messages` emits two
// kinds of assistant output:
//
//   1. `stream_event` envelopes carrying Anthropic Messages API streaming
//      events (`message_start`, `content_block_start`, `content_block_delta`,
//      `content_block_stop`, `message_delta`, `message_stop`).
//   2. Eventually, a single complete `assistant` envelope summarising the turn.
//
// To make the chat actually stream character-by-character we accumulate the
// `stream_event` deltas into a synthetic assistant envelope. When the real
// `assistant` envelope arrives, we replace the synthetic one with it (their
// content is identical at that point, but the final envelope has authoritative
// usage / metadata).

interface StreamEventEnvelope {
  type: "stream_event";
  event?: {
    type: string;
    index?: number;
    delta?: {
      type: string;
      text?: string;
      thinking?: string;
      partial_json?: string;
    };
    content_block?: {
      type: string;
      text?: string;
      thinking?: string;
      id?: string;
      name?: string;
      input?: unknown;
    };
    message?: { id?: string };
  };
}

type StreamingAssistantBlock = ContentBlock & {
  _partialJson?: string;
};

interface StreamingAssistantMessage {
  type: "assistant";
  _streamId?: string;
  message: {
    id?: string;
    content: StreamingAssistantBlock[];
    _streaming?: boolean;
  };
}

function findLastStreamingAssistant(prev: Envelope[]): number {
  for (let i = prev.length - 1; i >= 0; i--) {
    const env = prev[i] as { type?: string; message?: { _streaming?: boolean } };
    if (env?.type === "assistant" && env.message?._streaming) return i;
  }
  return -1;
}

function findAssistantByMessageId(prev: Envelope[], id: string | undefined): number {
  if (!id) return findLastStreamingAssistant(prev);
  for (let i = prev.length - 1; i >= 0; i--) {
    const env = prev[i] as { type?: string; message?: { id?: string } };
    if (env?.type === "assistant" && env.message?.id === id) return i;
  }
  return findLastStreamingAssistant(prev);
}

function mutateAssistantAt(
  prev: Envelope[],
  idx: number,
  fn: (m: StreamingAssistantMessage["message"]) => StreamingAssistantMessage["message"]
): Envelope[] {
  if (idx < 0) return prev;
  const env = prev[idx] as StreamingAssistantMessage;
  const next = [...prev];
  next[idx] = {
    ...env,
    message: fn(env.message || ({ content: [] } as StreamingAssistantMessage["message"])),
  };
  return next;
}

function mergeEnvelope(prev: Envelope[], envelope: Envelope): Envelope[] {
  if (!envelope || typeof envelope !== "object") return prev;
  const env = envelope as { type?: string };

  if (env.type === "stream_event") {
    const sse = envelope as StreamEventEnvelope;
    const evt = sse.event;
    if (!evt) return prev;

    if (evt.type === "message_start") {
      const placeholder: StreamingAssistantMessage = {
        type: "assistant",
        message: {
          id: evt.message?.id,
          content: [],
          _streaming: true,
        },
      };
      // Keep the message_start envelope itself in the array - its
      // `event.message.usage` is the only place we get the initial input /
      // cache token counts during live streaming. Without it, the meter is
      // stuck at zero until the post-reload replay re-injects the same
      // envelopes from the server.
      return [...prev, envelope, placeholder as unknown as Envelope];
    }

    if (evt.type === "content_block_start") {
      const idx = findAssistantByMessageId(prev, evt.message?.id);
      if (idx < 0) return prev;
      const blockIdx = evt.index ?? 0;
      return mutateAssistantAt(prev, idx, (msg) => {
        const blocks = [...(msg.content || [])];
        blocks[blockIdx] = { ...(evt.content_block as ContentBlock) };
        return { ...msg, content: blocks };
      });
    }

    if (evt.type === "content_block_delta") {
      const idx = findAssistantByMessageId(prev, evt.message?.id);
      if (idx < 0) return prev;
      const blockIdx = evt.index ?? 0;
      return mutateAssistantAt(prev, idx, (msg) => {
        const blocks = [...(msg.content || [])];
        const block = (blocks[blockIdx] || {}) as StreamingAssistantBlock;
        const next = { ...block } as StreamingAssistantBlock;
        const delta = evt.delta;
        if (delta?.type === "text_delta") {
          (next as { text?: string }).text =
            ((next as { text?: string }).text || "") + (delta.text || "");
          if (!next.type) (next as { type: string }).type = "text";
        } else if (delta?.type === "thinking_delta") {
          (next as { thinking?: string }).thinking =
            ((next as { thinking?: string }).thinking || "") + (delta.thinking || "");
          if (!next.type) (next as { type: string }).type = "thinking";
        } else if (delta?.type === "input_json_delta") {
          // tool_use input streams as JSON-string fragments; accumulate, parse
          // best-effort whenever the buffer is valid JSON.
          next._partialJson = (next._partialJson || "") + (delta.partial_json || "");
          try {
            (next as { input?: unknown }).input = JSON.parse(next._partialJson);
          } catch {
            /* still incomplete JSON - leave previous parsed value */
          }
        }
        blocks[blockIdx] = next;
        return { ...msg, content: blocks };
      });
    }

    if (evt.type === "message_stop") {
      const idx = findAssistantByMessageId(prev, evt.message?.id);
      if (idx < 0) return prev;
      return mutateAssistantAt(prev, idx, (msg) => ({ ...msg, _streaming: false }));
    }

    if (evt.type === "message_delta") {
      // message_delta carries the canonical per-message usage update (the
      // running output_tokens for this turn). Keep the envelope so
      // computeTokens can read it; otherwise the meter sits at the
      // message_start placeholder value (output_tokens=4 etc) for the
      // entire response.
      return [...prev, envelope];
    }

    // content_block_start/stop and other stream_event subtypes are mutations
    // on the placeholder we already track - no usage info, no need to keep
    // the envelope itself.
    return prev;
  }

  if (env.type === "assistant") {
    // Claude emits the canonical `assistant` envelope BEFORE `message_stop`,
    // so the message is still streaming at this point. Two regressions came
    // out of replacing the placeholder wholesale here:
    //   1. The `_streaming` flag was dropped, making the typewriter snap to
    //      full text the moment this envelope arrived.
    //   2. The final envelope sometimes ships only the `text` content block
    //      (the `thinking` block we accumulated from `thinking_delta`s
    //      disappears), so the thinking section vanished as soon as the
    //      stream finished.
    // Fix: when the placeholder was streaming, keep our delta-accumulated
    // content (it's the authoritative record of every block) and only pull
    // metadata from the incoming envelope. `message_stop` clears `_streaming`
    // and the typewriter then reveals any unrevealed tail instantly.
    const finalMsg = envelope as { message?: { id?: string; _streaming?: boolean } };
    const idx = findAssistantByMessageId(prev, finalMsg.message?.id);
    if (idx >= 0) {
      const prevEnv = prev[idx] as StreamingAssistantMessage;
      const next = [...prev];
      if (prevEnv.message?._streaming) {
        const incoming = envelope as { message?: Record<string, unknown> };
        const incomingMsg = (incoming.message || {}) as Record<string, unknown>;
        const accumulatedContent = prevEnv.message?.content || [];
        const incomingContent = (incomingMsg as { content?: ContentBlock[] }).content;
        // If the canonical envelope happens to carry MORE blocks (e.g. it
        // includes a tool_use we hadn't seen as a stream_event yet), prefer
        // it. Otherwise keep our accumulated blocks so we don't lose a
        // thinking section the canonical envelope omitted.
        const content =
          Array.isArray(incomingContent) && incomingContent.length > accumulatedContent.length
            ? incomingContent
            : accumulatedContent;
        next[idx] = {
          ...envelope,
          message: { ...incomingMsg, content, _streaming: true },
        } as Envelope;
      } else {
        next[idx] = envelope;
      }
      return next;
    }
    return [...prev, envelope];
  }

  return [...prev, envelope];
}

/**
 * Smooth out claude's bursty stream by dripping text/thinking deltas a few
 * characters per frame. Without this, short responses (where claude emits
 * the entire reply in one or two `text_delta` chunks) appear all-at-once.
 * The hook returns a derived envelope list with each actively-streaming
 * text/thinking block clamped to a displayed length that grows toward the
 * server's target via requestAnimationFrame.
 */
function useTypewriterEnvelopes(envelopes: Envelope[]): Envelope[] {
  const lengthsRef = useRef<Map<string, number>>(new Map());
  const envRef = useRef<Envelope[]>(envelopes);
  envRef.current = envelopes;
  const [tick, setTick] = useState(0);
  const rafRef = useRef<number | null>(null);
  const tickFnRef = useRef<(() => void) | null>(null);

  if (!tickFnRef.current) {
    tickFnRef.current = function tickFn() {
      const envs = envRef.current;
      const lengths = lengthsRef.current;
      let needsAnother = false;
      let mutated = false;
      for (let ei = 0; ei < envs.length; ei++) {
        const env = envs[ei];
        if (!env || (env as { type?: string }).type !== "assistant") continue;
        const e = env as StreamingAssistantMessage;
        const streaming = !!e.message?._streaming;
        const blocks = e.message?.content || [];
        for (let bi = 0; bi < blocks.length; bi++) {
          const b = blocks[bi];
          if (!b) continue;
          let key: string;
          let target: string;
          if (b.type === "text") {
            key = `${ei}:${bi}:t`;
            target = (b as { text?: string }).text || "";
          } else if (b.type === "thinking") {
            key = `${ei}:${bi}:th`;
            target = (b as { thinking?: string }).thinking || "";
          } else {
            continue;
          }
          const cur = lengths.get(key) ?? 0;
          if (cur >= target.length) continue;
          if (streaming) {
            // Catch up to target in roughly 0.4s; bigger gaps drip faster.
            const remaining = target.length - cur;
            const step = Math.max(2, Math.ceil(remaining / 24));
            lengths.set(key, Math.min(target.length, cur + step));
            needsAnother = true;
            mutated = true;
          } else {
            // Block is no longer streaming → reveal the rest instantly.
            lengths.set(key, target.length);
            mutated = true;
          }
        }
      }
      if (mutated) setTick((t) => (t + 1) & 0xffff);
      rafRef.current = needsAnother
        ? requestAnimationFrame(tickFnRef.current as FrameRequestCallback)
        : null;
    };
  }

  // Single long-lived RAF loop. Reads envelopes via ref so new server data
  // is picked up without tearing down and rescheduling the loop on every
  // websocket message - a previous version restarted on each envelope
  // change which dropped frames between bursts and hid the streaming.
  useEffect(() => {
    rafRef.current = requestAnimationFrame(tickFnRef.current as FrameRequestCallback);
    return () => {
      if (rafRef.current != null) cancelAnimationFrame(rafRef.current);
      rafRef.current = null;
    };
  }, []);

  // Wake the loop when new envelopes arrive if it's parked (no pending work).
  useEffect(() => {
    if (rafRef.current == null && envelopes.length > 0) {
      rafRef.current = requestAnimationFrame(tickFnRef.current as FrameRequestCallback);
    }
  }, [envelopes]);

  // Reset lengths when envelopes shrink (e.g., the user starts a new run).
  useEffect(() => {
    if (envelopes.length === 0 && lengthsRef.current.size > 0) {
      lengthsRef.current.clear();
    }
  }, [envelopes.length]);

  return useMemo(() => {
    const lengths = lengthsRef.current;
    return envelopes.map((env, ei) => {
      if (!env || (env as { type?: string }).type !== "assistant") return env;
      const e = env as StreamingAssistantMessage;
      const blocks = e.message?.content || [];
      let changed = false;
      const nextBlocks = blocks.map((b, bi) => {
        if (b.type === "text") {
          const full = (b as { text?: string }).text || "";
          const len = lengths.get(`${ei}:${bi}:t`) ?? full.length;
          if (len < full.length) {
            changed = true;
            return { ...b, text: full.slice(0, len) };
          }
        } else if (b.type === "thinking") {
          const full = (b as { thinking?: string }).thinking || "";
          const len = lengths.get(`${ei}:${bi}:th`) ?? full.length;
          if (len < full.length) {
            changed = true;
            return { ...b, thinking: full.slice(0, len) };
          }
        }
        return b;
      });
      if (!changed) return env;
      return {
        ...e,
        message: { ...e.message, content: nextBlocks },
      } as unknown as Envelope;
    });
    // tick is intentionally a dep so this memo re-runs on each RAF step.
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [envelopes, tick]);
}

/**
 * Subscribe to the live stream of one run.
 *
 * `runId` is the id whose frames this hook cares about — `null` while no run is
 * attached. `onStatus` and `onInputAck` fire only for a payload matching
 * `runId` (mirroring the page's old `handle && p.id === handle.id` guard);
 * `onAnyStatus` fires for EVERY `run_status` frame regardless of id, because
 * the page's run-list refresh has always been id-agnostic.
 */
export function useRunStream(
  runId: string | null,
  opts: {
    onStatus: (p: RunStatusPayload) => void;
    onInputAck: () => void;
    onAnyStatus: () => void;
  }
): {
  envelopes: Envelope[];
  setEnvelopes: React.Dispatch<React.SetStateAction<Envelope[]>>;
  displayEnvelopes: Envelope[];
} {
  const [envelopes, setEnvelopes] = useState<Envelope[]>([]);
  const displayEnvelopes = useTypewriterEnvelopes(envelopes);

  // Latest callbacks in a ref so the subscription's lifetime depends on the
  // run id alone - re-subscribing whenever a caller passes a fresh closure
  // would tear down and rebuild the bus handler on every page render.
  const optsRef = useRef(opts);
  optsRef.current = opts;

  // WebSocket subscription - only act on messages for the current run.
  useEffect(() => {
    return eventBus.subscribe((msg: WSMessage) => {
      if (msg.type === "run_stream") {
        const p = msg.data as RunStreamPayload;
        if (runId && p.id === runId) {
          // React 18 auto-batches async setStates, which collapses bursts of
          // stream_event deltas (and the final `assistant` envelope that
          // follows them) into a single render - visually erasing the
          // streaming effect. flushSync forces a commit per envelope so the
          // user sees text_delta / thinking_delta chunks paint as they
          // arrive instead of all at once.
          flushSync(() => {
            setEnvelopes((prev) => mergeEnvelope(prev, p.envelope as Envelope));
          });
        }
      } else if (msg.type === "run_status") {
        const p = msg.data as RunStatusPayload;
        if (runId && p.id === runId) {
          optsRef.current.onStatus(p);
        }
        optsRef.current.onAnyStatus();
      } else if (msg.type === "run_input_ack") {
        const p = msg.data as RunInputAckPayload;
        if (runId && p.id === runId) {
          optsRef.current.onInputAck();
        }
      }
    });
  }, [runId]);

  return { envelopes, setEnvelopes, displayEnvelopes };
}
