// src/kb/KbResponseInterceptor.ts
import type { ContentBlock, ToolUseContentBlock, ContinuationFn } from "./types.ts";
import type { KbStore } from "./KbStore.ts";
import { parseSseStream, encodeEvent } from "../codec/sseStream.ts";
import { getLedger } from "../stats/ledger.ts";

function formatKbResult(chunks: { content: string }[]): string {
  if (chunks.length === 0) return "(empty — KB has no entries yet)";
  return chunks.map(c => `- ${c.content}`).join("\n");
}

const CONTENT_BLOCK_EVENTS = new Set([
  "content_block_start",
  "content_block_delta",
  "content_block_stop",
]);

function isKbToolName(name: string): boolean {
  return name === "ftRetrieve" || name === "kb";
}

export class KbResponseInterceptor {
  constructor(
    private readonly store: KbStore,
    private readonly sessionId = "__default__",
  ) {}

  wrapResponse(upstreamRes: Response, continuationFn: ContinuationFn): Response {
    if (!upstreamRes.body) return upstreamRes;
    const stream = this.createInterceptStream(upstreamRes.body, continuationFn);
    return new Response(stream, {
      status: upstreamRes.status,
      statusText: upstreamRes.statusText,
      headers: upstreamRes.headers,
    });
  }

  private createInterceptStream(
    upstream: ReadableStream<Uint8Array>,
    continuationFn: ContinuationFn,
  ): ReadableStream<Uint8Array> {
    const store = this.store;
    const sessionId = this.sessionId;
    const enc = new TextEncoder();
    const dec = new TextDecoder("utf-8", { fatal: false });

    return new ReadableStream<Uint8Array>({
      async start(controller) {
        let didError = false;
        /**
         * Lazy-stream strategy (discriminated by tool name)
         * ==================================================
         * message_start     → flush immediately (always)
         * text blocks       → flush each event as it arrives (TTFT fix)
         * kb tool_use       → buffer (hide from client; do server-side continuation)
         * non-kb tool_use   → flush live (stream events to client immediately)
         * message_delta     → flush if no kb tool_use seen; buffer otherwise
         *                     (don't send stop_reason:tool_use to client for kb tools)
         * message_stop      → classify and handle (see below)
         *
         * At message_stop:
         *   no tool_use          → flush remaining buffer (should be empty) + stop
         *   kb tool_use only     → do kb continuation:
         *                            - always strip continuation's message_start
         *                              (client already received ours)
         *                            - if text blocks were pre-flushed: rewrite
         *                              content block indices in continuation stream
         *                              (offset by flushedBlockCount)
         *   non-kb tool only     → flush remaining buffer (should be empty) + stop
         *   mixed kb+non-kb      → DROP buffered kb events, flush only message_stop
         *                          (client already received non-kb events live)
         */

        // Events for kb tool_use blocks and message_delta (held until message_stop)
        const bufferedEvents: string[] = [];
        // Full assistant content for continuation call
        const assistantContent: ContentBlock[] = [];
        let activeToolUse: { id: string; name: string; inputJson: string; isKb: boolean } | null = null;
        let activeTextIndex: number | null = null;
        // Raw block clones — mutated by deltas, pushed verbatim at content_block_stop.
        // Preserves all fields Anthropic emits (including future additions) so the
        // continuation request matches Anthropic's signed thinking blocks byte-for-byte.
        let activeThinkingRaw: Record<string, unknown> | null = null;
        let activeRedactedThinkingRaw: Record<string, unknown> | null = null;
        let activeOtherRaw: Record<string, unknown> | null = null;

        // Number of text content blocks fully flushed to client before any kb tool_use appeared
        let flushedBlockCount = 0;
        // Set true on first kb tool_use content_block_start; causes subsequent events to buffer
        let hasKbToolUse = false;
        // Invariant: assistantContent.length MUST equal contentBlockStartCount at message_stop.
        // Mismatch = a content block was dropped from reconstruction → index shift → 400.
        let contentBlockStartCount = 0;

        function flush(rawEvent: string): void {
          controller.enqueue(enc.encode(rawEvent));
        }

        function flushOrBuffer(rawEvent: string): void {
          if (hasKbToolUse) bufferedEvents.push(rawEvent);
          else flush(rawEvent);
        }

        try {
          for await (const { rawEvent, parsed } of parseSseStream(upstream, dec)) {
            if (!parsed) { flushOrBuffer(rawEvent); continue; }

            const type = parsed.type as string;

            if (type === "message_start") {
              // Always flush immediately — never buffer message_start.
              flush(rawEvent);

            } else if (type === "content_block_start") {
              contentBlockStartCount++;
              const cb = parsed.content_block as Record<string, unknown>;
              const idx = parsed.index as number;

              if (cb.type === "tool_use") {
                const toolName = cb.name as string;
                const isKb = isKbToolName(toolName);
                if (isKb) {
                  hasKbToolUse = true;
                  bufferedEvents.push(rawEvent);
                } else {
                  flush(rawEvent);
                }
                activeToolUse = { id: cb.id as string, name: toolName, inputJson: "", isKb };
                activeTextIndex = null;
                activeThinkingRaw = null;
                activeRedactedThinkingRaw = null;
                activeOtherRaw = null;
              } else if (cb.type === "text") {
                activeTextIndex = idx;
                activeToolUse = null;
                activeThinkingRaw = null;
                activeRedactedThinkingRaw = null;
                activeOtherRaw = null;
                assistantContent.push({ type: "text", text: "" });
                flushOrBuffer(rawEvent);
              } else if (cb.type === "thinking") {
                activeTextIndex = null;
                activeToolUse = null;
                // Clone the entire content_block — preserves every field Anthropic
                // emits (thinking, signature, plus any future-added metadata) so the
                // continuation message matches the signed original byte-for-byte.
                activeThinkingRaw = { ...cb };
                if (typeof activeThinkingRaw.thinking !== "string") activeThinkingRaw.thinking = "";
                if (typeof activeThinkingRaw.signature !== "string") activeThinkingRaw.signature = "";
                activeRedactedThinkingRaw = null;
                activeOtherRaw = null;
                flushOrBuffer(rawEvent);
              } else if (cb.type === "redacted_thinking") {
                activeTextIndex = null;
                activeToolUse = null;
                activeThinkingRaw = null;
                activeRedactedThinkingRaw = { ...cb };
                if (typeof activeRedactedThinkingRaw.data !== "string") activeRedactedThinkingRaw.data = "";
                activeOtherRaw = null;
                flushOrBuffer(rawEvent);
              } else {
                // Unknown content block type (future server_tool_use, web_search_tool_result, etc.).
                // Capture verbatim so it round-trips into the continuation message —
                // dropping it shifts later block indices and Anthropic 400s on signature mismatch.
                activeTextIndex = null;
                activeToolUse = null;
                activeThinkingRaw = null;
                activeRedactedThinkingRaw = null;
                activeOtherRaw = { ...cb };
                flushOrBuffer(rawEvent);
              }

            } else if (type === "content_block_delta") {
              const delta = parsed.delta as Record<string, unknown>;
              if (delta.type === "input_json_delta" && activeToolUse) {
                activeToolUse.inputJson += (delta.partial_json as string) ?? "";
                if (activeToolUse.isKb) {
                  bufferedEvents.push(rawEvent); // kb: buffer tool input
                } else {
                  flush(rawEvent); // non-kb: stream live
                }
              } else if (delta.type === "text_delta" && activeTextIndex !== null) {
                const last = assistantContent[assistantContent.length - 1];
                if (last?.type === "text") {
                  (last as { type: "text"; text: string }).text += (delta.text as string) ?? "";
                }
                flushOrBuffer(rawEvent);
              } else if (delta.type === "thinking_delta" && activeThinkingRaw) {
                activeThinkingRaw.thinking = String(activeThinkingRaw.thinking ?? "") + ((delta.thinking as string) ?? "");
                flushOrBuffer(rawEvent);
              } else if (delta.type === "signature_delta" && activeThinkingRaw) {
                activeThinkingRaw.signature = String(activeThinkingRaw.signature ?? "") + ((delta.signature as string) ?? "");
                flushOrBuffer(rawEvent);
              } else {
                flushOrBuffer(rawEvent);
              }

            } else if (type === "content_block_stop") {
              if (activeToolUse) {
                // Tool-use block ending: materialise its input
                let input: Record<string, unknown> = {};
                try { input = JSON.parse(activeToolUse.inputJson || "{}"); } catch { /* ignore */ }
                assistantContent.push({
                  type: "tool_use",
                  id: activeToolUse.id,
                  name: activeToolUse.name,
                  input,
                });
                const wasKb = activeToolUse.isKb;
                activeToolUse = null;
                if (wasKb) {
                  bufferedEvents.push(rawEvent); // kb: buffer the stop
                } else {
                  flush(rawEvent); // non-kb: stream live
                }
              } else if (activeThinkingRaw) {
                // Push the mutated clone as-is — every original field preserved.
                assistantContent.push({ ...activeThinkingRaw } as unknown as ContentBlock);
                activeThinkingRaw = null;
                if (hasKbToolUse) {
                  bufferedEvents.push(rawEvent);
                } else {
                  flush(rawEvent);
                  flushedBlockCount++;
                }
              } else if (activeRedactedThinkingRaw) {
                assistantContent.push({ ...activeRedactedThinkingRaw } as unknown as ContentBlock);
                activeRedactedThinkingRaw = null;
                if (hasKbToolUse) {
                  bufferedEvents.push(rawEvent);
                } else {
                  flush(rawEvent);
                  flushedBlockCount++;
                }
              } else if (activeTextIndex !== null) {
                // Text block ending
                activeTextIndex = null;
                if (hasKbToolUse) {
                  // kb tool_use appeared earlier in this response; buffer
                  bufferedEvents.push(rawEvent);
                } else {
                  // No kb tool so far: flush and count the completed block
                  flush(rawEvent);
                  // Invariant: increment ONLY when block was flushed to client AND no kb tool_use seen yet.
                  flushedBlockCount++;
                }
              } else if (activeOtherRaw) {
                // Unknown block ends: must push to assistantContent so the continuation
                // message has the same content[].length as Anthropic's original response.
                // Dropping it would index-shift later blocks and trip the thinking-block
                // signature check (400 "thinking blocks cannot be modified").
                assistantContent.push({ ...activeOtherRaw } as unknown as ContentBlock);
                activeOtherRaw = null;
                if (hasKbToolUse) {
                  bufferedEvents.push(rawEvent);
                } else {
                  flush(rawEvent);
                  flushedBlockCount++;
                }
              }

            } else if (type === "message_stop") {
              const kbBlocks = assistantContent.filter(
                (b): b is ToolUseContentBlock => b.type === "tool_use" && isKbToolName(b.name),
              );
              const nonKbToolBlocks = assistantContent.filter(
                b => b.type === "tool_use" && !isKbToolName((b as ToolUseContentBlock).name),
              );

              if (kbBlocks.length > 0 && nonKbToolBlocks.length > 0) {
                // ── Mixed kb+non-kb path ─────────────────────────────────
                // Non-kb events were already flushed live; drop buffered kb events.
                console.warn("[kb-interceptor] mixed kb+non-kb tool_use in one response; dropping kb events");
                flush(rawEvent);

              } else if (kbBlocks.length > 0 && nonKbToolBlocks.length === 0) {
                // ── KB continuation path ─────────────────────────────────
                if (assistantContent.length !== contentBlockStartCount) {
                  console.warn(
                    `[kb-interceptor] reconstruction invariant violated: assistantContent.length=${assistantContent.length} != contentBlockStartCount=${contentBlockStartCount}. ` +
                    `Continuation may trip Anthropic's signed-thinking-block check (400).`,
                  );
                }
                const kb = kbBlocks[0]!;
                const input = kb.input as Record<string, unknown>;
                const q = (input.q as string | undefined) ?? (input.query as string | undefined) ?? "";
                const id = typeof input.id === "string" ? input.id.trim() : "";
                const label = id.startsWith("json:") ? id : id ? `json:${id}` : "";
                let chunks: ReturnType<typeof store.search> = [];
                try {
                  chunks = label
                    ? store.searchByLabel(label, q, { sessionId, limit: 5 })
                    : store.search(q, { sessionId, limit: 5 });
                } catch { /* empty on DB error */ }
                const toolResultContent = formatKbResult(chunks);

                try {
                  const contRes = await continuationFn(assistantContent, kb.id, toolResultContent);
                  if (contRes.body) {
                    // Stream continuation, stripping its message_start (client already has ours)
                    // and rewriting content block indices when text blocks were pre-flushed.
                    const contDec = new TextDecoder("utf-8", { fatal: false });
                    for await (const { rawEvent: ce, parsed: cp } of parseSseStream(contRes.body, contDec)) {
                      if (!cp) { flush(ce); continue; }
                      const ctype = cp.type as string;

                      if (ctype === "message_start") {
                        // Always skip: client already has one message_start
                        continue;
                      }

                      if (flushedBlockCount > 0 && CONTENT_BLOCK_EVENTS.has(ctype)) {
                        // Renumber content blocks to follow the already-flushed text blocks
                        const origIdx = (cp.index as number | undefined) ?? 0;
                        const rewritten = { ...cp, index: origIdx + flushedBlockCount };
                        flush(encodeEvent(ctype, rewritten));
                      } else {
                        flush(ce);
                      }
                    }
                  }
                  try {
                    const inJson = JSON.stringify(kb.input ?? {});
                    getLedger().record("fewtok-core", {
                      bytesIn: Buffer.byteLength(inJson, "utf-8"),
                      bytesOut: Buffer.byteLength(toolResultContent, "utf-8"),
                      hits: 1,
                    });
                  } catch { /* swallow */ }
                } catch {
                  // Continuation failed: fall back to passthrough so client sees error context
                  for (const ev of bufferedEvents) flush(ev);
                  flush(rawEvent);
                }

              } else {
                // ── Passthrough path ─────────────────────────────────────
                // No tool_use, or non-kb tool only (events already flushed live):
                // flush any remaining buffered events then stop.
                for (const ev of bufferedEvents) flush(ev);
                flush(rawEvent);
              }

              break; // message_stop = end of stream

            } else {
              // message_delta, ping, etc.
              flushOrBuffer(rawEvent);
            }
          }
        } catch (err) {
          didError = true;
          controller.error(err);
        } finally {
          if (!didError) controller.close();
        }
      },
    });
  }
}