import type { CapturedLog, JsonValue } from "../../schemas";
import { trustAsJsonValue } from "../jsonSchema";
import { SseEventSchema } from "./schemas";

type StreamTextBlock = { type: "text"; text: string };
type StreamThinkingBlock = { type: "thinking"; thinking: string; signature: string };
type StreamThinkBlock = { type: "think"; thinking: string; signature: string };
type StreamToolUseBlock = { type: "tool_use"; id: string; name: string; inputJson: string };
type StreamBlock = StreamTextBlock | StreamThinkingBlock | StreamThinkBlock | StreamToolUseBlock;

function parseInputJson(json: string): Record<string, unknown> {
  if (json === "") return {};
  try {
    const parsed: unknown = JSON.parse(json);
    if (parsed !== null && typeof parsed === "object" && !Array.isArray(parsed)) {
      return { ...parsed };
    }
    return {};
  } catch {
    return {};
  }
}

function finalizeBlock(block: StreamBlock): Record<string, unknown> {
  switch (block.type) {
    case "text":
      return { type: "text", text: block.text };
    case "thinking": {
      const out: Record<string, unknown> = { type: "thinking", thinking: block.thinking };
      if (block.signature !== "") out.signature = block.signature;
      return out;
    }
    case "think": {
      const out: Record<string, unknown> = { type: "think", thinking: block.thinking };
      if (block.signature !== "") out.signature = block.signature;
      return out;
    }
    case "tool_use":
      return {
        type: "tool_use",
        id: block.id,
        name: block.name,
        input: parseInputJson(block.inputJson),
      };
  }
}

/**
 * Extract Anthropic streaming response and return normalized JSON.
 * Optionally collects SSE chunks into log.streamingChunks for debugging/display.
 */
export function extractAnthropicStream(
  raw: string,
  log: CapturedLog,
  fallbackModel?: string,
  collectChunks?: boolean,
): string {
  const blocks = new Map<number, StreamBlock>();
  let id = "";
  let model = "";
  let stopReason: string | null = null;
  let stopSequence: string | null = null;
  let inputTokens = 0;
  let outputTokens = 0;

  const MAX_CHUNKS = 1000;
  let chunkIndex = 0;
  let streamStartMs = 0;
  const chunks: Array<{ index: number; timestamp: number; type: string; data: JsonValue }> = [];

  for (const line of raw.split("\n")) {
    const trimmedLine = line.trim();
    if (!trimmedLine.startsWith("data: ")) continue;
    try {
      const json: unknown = JSON.parse(trimmedLine.slice(6));
      const parsed = SseEventSchema.safeParse(json);
      if (!parsed.success) continue;
      const data = parsed.data;

      if (chunkIndex === 0) streamStartMs = Date.now();

      if (collectChunks === true && chunks.length < MAX_CHUNKS) {
        chunks.push({
          index: chunkIndex,
          timestamp: Date.now() - streamStartMs,
          type: data.type,
          // `data` has already been validated by SseEventSchema above; the
          // recursive JsonValue walk is a redundant full-tree zod parse on
          // every SSE event. See trustAsJsonValue for the rationale.
          data: trustAsJsonValue(data),
        });
      }
      chunkIndex++;

      switch (data.type) {
        case "message_start":
          id = data.message.id;
          model = data.message.model;
          inputTokens = data.message.usage.input_tokens;
          log.inputTokens = inputTokens;
          log.cacheCreationInputTokens = data.message.usage.cache_creation_input_tokens ?? null;
          log.cacheReadInputTokens = data.message.usage.cache_read_input_tokens ?? null;
          if (log.model === null) log.model = model;
          break;
        case "content_block_start": {
          const cb = data.content_block;
          switch (cb.type) {
            case "text":
              blocks.set(data.index, { type: "text", text: cb.text ?? "" });
              break;
            case "thinking":
              blocks.set(data.index, {
                type: "thinking",
                thinking: cb.thinking ?? "",
                signature: cb.signature ?? "",
              });
              break;
            case "think":
              blocks.set(data.index, {
                type: "think",
                thinking: cb.thinking ?? "",
                signature: cb.signature ?? "",
              });
              break;
            case "tool_use":
              blocks.set(data.index, {
                type: "tool_use",
                id: cb.id,
                name: cb.name,
                inputJson: "",
              });
              break;
          }
          break;
        }
        case "content_block_delta": {
          const block = blocks.get(data.index);
          if (block === undefined) break;
          const delta = data.delta;
          if (delta.type === "text_delta" && block.type === "text") {
            block.text += delta.text;
          } else if (
            delta.type === "thinking_delta" &&
            (block.type === "thinking" || block.type === "think")
          ) {
            block.thinking += delta.thinking;
          } else if (
            delta.type === "signature_delta" &&
            (block.type === "thinking" || block.type === "think")
          ) {
            block.signature += delta.signature;
          } else if (delta.type === "input_json_delta" && block.type === "tool_use") {
            block.inputJson += delta.partial_json;
          }
          break;
        }
        case "message_delta":
          stopReason = data.delta.stop_reason;
          stopSequence = data.delta.stop_sequence ?? null;
          outputTokens = data.usage.output_tokens;
          log.outputTokens = outputTokens;
          // Cache tokens may also be present in message_delta usage
          if (data.usage.cache_creation_input_tokens !== undefined) {
            log.cacheCreationInputTokens = data.usage.cache_creation_input_tokens;
          }
          if (data.usage.cache_read_input_tokens !== undefined) {
            log.cacheReadInputTokens = data.usage.cache_read_input_tokens;
          }
          break;
        case "error":
          // Store error info on the log for display
          log.error = data.error.message;
          break;
        case "content_block_stop":
        case "message_stop":
        case "ping":
          break;
      }
    } catch {
      // non-JSON SSE line, skip
    }
  }

  if (collectChunks === true) {
    log.streamingChunks = {
      chunks,
      truncated: chunkIndex > MAX_CHUNKS,
    };
  }

  const orderedContent = [...blocks.entries()]
    .sort(([a], [b]) => a - b)
    .map(([, block]) => finalizeBlock(block));

  return JSON.stringify({
    id,
    type: "message",
    model: model !== "" ? model : (fallbackModel ?? ""),
    role: "assistant",
    content: orderedContent,
    stop_reason: stopReason,
    stop_sequence: stopSequence,
    usage: { input_tokens: inputTokens, output_tokens: outputTokens },
  });
}
