// src/layer/ExecResponseInterceptor.ts
import { parseSseStream } from "../codec/sseStream.ts";
import { encodeEvent } from "../codec/sseParser.ts";
import type { ContinuationFn, ContentBlock, ToolUseContentBlock } from "../kb/types.ts";
import { Executor } from "../sandbox/Executor.ts";
import { detectLanguage } from "../sandbox/detectLanguage.ts";
import { getLedger } from "../stats/ledger.ts";

const SANDBOX_OUTPUT_CAP_BYTES = 100_000;

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

export class ExecResponseInterceptor {
  private readonly executor = new Executor();

  constructor(private readonly available: string[]) {}

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

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

    return new ReadableStream<Uint8Array>({
      async start(ctrl) {
        let didError = false;
        const flush = (raw: string): void => { ctrl.enqueue(enc.encode(raw)); };
        const buf: string[] = [];
        const assistantBlocks: ContentBlock[] = [];
        let hasExecBlock = false;
        // Track the current tool_use block — exec or non-exec
        let activeToolUse: { id: string; name: string; inputJson: string; isExec: boolean } | null = null;
        // Raw block clones — mutated by deltas, pushed verbatim at content_block_stop.
        // Preserves every field Anthropic emits so the continuation message matches
        // the signed original byte-for-byte (signature check is fragile).
        let activeThinkingRaw: Record<string, unknown> | null = null;
        let activeRedactedThinkingRaw: Record<string, unknown> | null = null;
        let activeOtherRaw: Record<string, unknown> | null = null;
        let flushedCount = 0;
        // Invariant: assistantBlocks.length MUST equal contentBlockStartCount at message_stop.
        let contentBlockStartCount = 0;

        const flushOrBuf = (raw: string): void => {
          if (hasExecBlock) buf.push(raw);
          else flush(raw);
        };

        try {
          for await (const { rawEvent, parsed } of parseSseStream(upstream, dec)) {
            if (!parsed) { flushOrBuf(rawEvent); continue; }
            const t = parsed.type as string;

            if (t === "message_start") { flush(rawEvent); continue; }

            if (t === "content_block_start") {
              contentBlockStartCount++;
              const cb = parsed.content_block as Record<string, unknown>;
              if (cb.type === "tool_use") {
                const isExec = cb.name === "exec";
                if (isExec) hasExecBlock = true;
                activeToolUse = { id: cb.id as string, name: cb.name as string, inputJson: "", isExec };
                activeThinkingRaw = null;
                activeRedactedThinkingRaw = null;
                activeOtherRaw = null;
                // Exec blocks buffered; non-exec flushed live so client sees them
                if (isExec) buf.push(rawEvent);
                else flush(rawEvent);
              } else if (cb.type === "text") {
                activeThinkingRaw = null;
                activeRedactedThinkingRaw = null;
                activeOtherRaw = null;
                assistantBlocks.push({ type: "text", text: "" });
                flushOrBuf(rawEvent);
              } else if (cb.type === "thinking") {
                activeThinkingRaw = { ...cb };
                if (typeof activeThinkingRaw.thinking !== "string") activeThinkingRaw.thinking = "";
                if (typeof activeThinkingRaw.signature !== "string") activeThinkingRaw.signature = "";
                activeRedactedThinkingRaw = null;
                activeOtherRaw = null;
                flushOrBuf(rawEvent);
              } else if (cb.type === "redacted_thinking") {
                activeThinkingRaw = null;
                activeRedactedThinkingRaw = { ...cb };
                if (typeof activeRedactedThinkingRaw.data !== "string") activeRedactedThinkingRaw.data = "";
                activeOtherRaw = null;
                flushOrBuf(rawEvent);
              } else {
                // Unknown block type — capture verbatim so it round-trips into the
                // continuation message. Dropping it would index-shift later blocks
                // and break Anthropic's signed thinking-block check.
                activeThinkingRaw = null;
                activeRedactedThinkingRaw = null;
                activeOtherRaw = { ...cb };
                flushOrBuf(rawEvent);
              }
              continue;
            }

            if (t === "content_block_delta") {
              const d = parsed.delta as Record<string, unknown>;
              if (d.type === "input_json_delta" && activeToolUse) {
                activeToolUse.inputJson += (d.partial_json as string) ?? "";
                // Route same as block start: exec → buffer, non-exec → flush live
                if (activeToolUse.isExec) buf.push(rawEvent);
                else flush(rawEvent);
              } else if (d.type === "text_delta") {
                const last = assistantBlocks[assistantBlocks.length - 1];
                if (last?.type === "text") {
                  (last as { type: "text"; text: string }).text += (d.text as string) ?? "";
                }
                flushOrBuf(rawEvent);
              } else if (d.type === "thinking_delta" && activeThinkingRaw) {
                activeThinkingRaw.thinking = String(activeThinkingRaw.thinking ?? "") + ((d.thinking as string) ?? "");
                flushOrBuf(rawEvent);
              } else if (d.type === "signature_delta" && activeThinkingRaw) {
                activeThinkingRaw.signature = String(activeThinkingRaw.signature ?? "") + ((d.signature as string) ?? "");
                flushOrBuf(rawEvent);
              } else {
                flushOrBuf(rawEvent);
              }
              continue;
            }

            if (t === "content_block_stop") {
              if (activeToolUse) {
                let input: Record<string, unknown> = {};
                try { input = JSON.parse(activeToolUse.inputJson || "{}"); } catch { /* ignore */ }
                // ALL tool_use blocks tracked — enables mixed-exec detection
                assistantBlocks.push({ type: "tool_use", id: activeToolUse.id, name: activeToolUse.name, input });
                if (activeToolUse.isExec) {
                  buf.push(rawEvent);
                } else {
                  // Non-exec stop flushed live. Increment flushedCount so that
                  // if a LATER exec block triggers a continuation, the continuation's
                  // content block indices are offset past these already-visible blocks.
                  // In practice, exec-only responses (the continuation path) never
                  // reach this branch — non-exec tool_use + exec = mixed → drop exec.
                  flush(rawEvent);
                  flushedCount++;
                }
                activeToolUse = null;
              } else if (activeThinkingRaw) {
                assistantBlocks.push({ ...activeThinkingRaw } as unknown as ContentBlock);
                activeThinkingRaw = null;
                if (hasExecBlock) {
                  buf.push(rawEvent);
                } else {
                  flush(rawEvent);
                  flushedCount++;
                }
              } else if (activeRedactedThinkingRaw) {
                assistantBlocks.push({ ...activeRedactedThinkingRaw } as unknown as ContentBlock);
                activeRedactedThinkingRaw = null;
                if (hasExecBlock) {
                  buf.push(rawEvent);
                } else {
                  flush(rawEvent);
                  flushedCount++;
                }
              } else if (activeOtherRaw) {
                // Unknown block ends: push verbatim into assistantBlocks so the
                // continuation message preserves content[].length and block ordering.
                assistantBlocks.push({ ...activeOtherRaw } as unknown as ContentBlock);
                activeOtherRaw = null;
                if (hasExecBlock) {
                  buf.push(rawEvent);
                } else {
                  flush(rawEvent);
                  flushedCount++;
                }
              } else {
                // text block stop
                if (!hasExecBlock) flushedCount++;
                flushOrBuf(rawEvent);
              }
              continue;
            }

            if (t === "message_stop") {
              const execBlocks = assistantBlocks.filter(
                (b): b is ToolUseContentBlock => b.type === "tool_use" && b.name === "exec",
              );
              const nonExecToolBlocks = assistantBlocks.filter(
                b => b.type === "tool_use" && (b as ToolUseContentBlock).name !== "exec",
              );

              if (execBlocks.length > 0 && nonExecToolBlocks.length > 0) {
                // Mixed exec+non-exec: non-exec events already flushed live; drop buffered exec
                console.warn("[exec-interceptor] mixed exec+non-exec tool_use; dropping exec events");
                flush(rawEvent);

              } else if (execBlocks.length === 0) {
                // Passthrough: no exec tool_use
                for (const ev of buf) flush(ev);
                flush(rawEvent);

              } else {
                // Exec continuation path
                if (assistantBlocks.length !== contentBlockStartCount) {
                  console.warn(
                    `[exec-interceptor] reconstruction invariant violated: assistantBlocks.length=${assistantBlocks.length} != contentBlockStartCount=${contentBlockStartCount}. ` +
                    `Continuation may trip Anthropic's signed-thinking-block check (400).`,
                  );
                }
                const eb = execBlocks[0]!;
                const code     = (eb.input["code"]     as string) ?? "";
                const langHint = (eb.input["language"] as string) ?? "";
                const language = langHint || (detectLanguage(code, self.available) ?? "");

                let toolResult: string;
                let execTruncated = false;
                if (!language) {
                  toolResult = "error: language required — could not auto-detect";
                } else {
                  try {
                    const r = await self.executor.run(language, code);
                    const flags = [
                      r.timedOut ? " TIMEOUT" : "",
                      r.truncated ? " TRUNCATED" : "",
                    ].join("");
                    toolResult = `[${language}] exit=${r.exitCode}${flags}\n${r.output}`;
                    execTruncated = r.truncated;
                  } catch (err) {
                    toolResult = `error: ${err instanceof Error ? err.message : String(err)}`;
                  }
                }

                try {
                  const contRes = await continuationFn(assistantBlocks, eb.id, toolResult);
                  if (contRes.body) {
                    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") continue; // skip: client already has one
                      if (flushedCount > 0 && CONTENT_BLOCK_EVENTS.has(ctype)) {
                        const origIdx = (cp.index as number | undefined) ?? 0;
                        flush(encodeEvent(ctype, { ...cp, index: origIdx + flushedCount }));
                      } else {
                        flush(ce);
                      }
                    }
                  }
                  if (execTruncated) {
                    try {
                      getLedger().record("fewtok-core", {
                        bytesIn: SANDBOX_OUTPUT_CAP_BYTES,
                        bytesOut: Buffer.byteLength(toolResult, "utf-8"),
                        hits: 1,
                      });
                    } catch { /* swallow */ }
                  }
                } catch {
                  // Continuation failed: fall back to buffered events so client sees context
                  for (const ev of buf) flush(ev);
                  flush(rawEvent);
                }
              }
              break;
            }

            flushOrBuf(rawEvent);
          }
        } catch (err) {
          didError = true;
          ctrl.error(err);
        } finally {
          if (!didError) ctrl.close();
        }
      },
    });
  }
}