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

/**
 * Extract OpenAI streaming response and return normalized JSON.
 * Optionally collects SSE chunks into log.streamingChunks for debugging/display.
 */
export function extractOpenAIStream(
  raw: string,
  log: CapturedLog,
  fallbackModel?: string,
  collectChunks: boolean = true,
): string {
  let id = "";
  let model = "";
  let completionContent = "";
  let reasoningContent = "";
  let finishReason: string | null = null;
  let promptTokens = 0;
  let completionTokens = 0;
  let usageCaptured = false;
  let started = false;
  type ToolCallEntry = {
    index: number;
    id?: string;
    type: "function";
    function: { name: string; arguments: string };
  };
  // Map keyed by `index` so delta accumulation is O(1) per delta instead of
  // O(N) (the previous Array#find). Map preserves insertion order, so the
  // final `[...toolCalls.values()]` keeps tool calls in the order they were
  // first seen.
  const toolCalls: Map<number, ToolCallEntry> = new Map();

  // Chunk collection state
  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;
    const dataStr = trimmedLine.slice(6);
    if (dataStr === "[DONE]") break;

    try {
      const parsed: unknown = JSON.parse(dataStr);
      // Check for error object in SSE stream before trying to parse as chunk
      if (typeof parsed === "object" && parsed !== null && !Array.isArray(parsed)) {
        const errorDesc = Object.getOwnPropertyDescriptor(parsed, "error");
        if (
          errorDesc !== undefined &&
          typeof errorDesc.value === "object" &&
          errorDesc.value !== null
        ) {
          const msgDesc = Object.getOwnPropertyDescriptor(errorDesc.value, "message");
          if (msgDesc !== undefined && typeof msgDesc.value === "string") {
            log.error = msgDesc.value;
          }
          continue;
        }
      }
      const chunkResult = OpenAISSERawChunkSchema.safeParse(parsed);
      if (!chunkResult.success) continue;
      const chunk = chunkResult.data;

      // Track stream start for timestamps
      if (chunkIndex === 0) streamStartMs = Date.now();

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

      if (!started) {
        id = chunk.id;
        model = chunk.model;
        started = true;
      }

      if (!usageCaptured && chunk.usage !== undefined && chunk.usage !== null) {
        promptTokens = chunk.usage.prompt_tokens ?? 0;
        completionTokens = chunk.usage.completion_tokens ?? 0;
        log.inputTokens = promptTokens;
        // Extract cached_tokens from raw parsed object (passthrough field in usage)
        if (typeof parsed === "object" && parsed !== null && !Array.isArray(parsed)) {
          const usageDesc = Object.getOwnPropertyDescriptor(parsed, "usage");
          if (
            usageDesc !== undefined &&
            typeof usageDesc.value === "object" &&
            usageDesc.value !== null
          ) {
            const detailsDesc = Object.getOwnPropertyDescriptor(
              usageDesc.value,
              "prompt_tokens_details",
            );
            if (
              detailsDesc !== undefined &&
              typeof detailsDesc.value === "object" &&
              detailsDesc.value !== null
            ) {
              const cacheDesc = Object.getOwnPropertyDescriptor(detailsDesc.value, "cached_tokens");
              if (cacheDesc !== undefined && typeof cacheDesc.value === "number") {
                log.cacheReadInputTokens = cacheDesc.value;
              }
            }
          }
        }
        usageCaptured = true;
      }

      for (const choice of chunk.choices) {
        const delta = choice.delta;
        if (delta.content !== undefined && delta.content !== null) {
          completionContent += delta.content;
        }
        if (delta.reasoning_content !== undefined && delta.reasoning_content !== null) {
          reasoningContent += delta.reasoning_content;
        }
        if (delta.thinking !== undefined && delta.thinking !== null) {
          reasoningContent += delta.thinking;
        }
        // MiniMax uses 'think' field for thinking content - check via bracket notation
        const thinkValue = delta.think;
        if (thinkValue !== undefined && thinkValue !== null) {
          reasoningContent += thinkValue;
        }
        if (choice.finish_reason !== undefined && choice.finish_reason !== null) {
          finishReason = choice.finish_reason;
        }
        // Accumulate tool_calls
        if (delta.tool_calls !== undefined && delta.tool_calls !== null) {
          for (const tc of delta.tool_calls) {
            // Find or create tool call entry by index
            let existing = toolCalls.get(tc.index);
            if (existing === undefined) {
              existing = {
                index: tc.index,
                type: "function",
                function: { name: "", arguments: "" },
              };
              toolCalls.set(tc.index, existing);
            }
            if (tc.id !== undefined) existing.id = tc.id;
            if (tc.function?.name !== undefined) existing.function.name += tc.function.name;
            if (tc.function?.arguments !== undefined)
              existing.function.arguments += tc.function.arguments;
          }
        }
      }
    } catch {
      // non-JSON SSE line, skip
    }
  }

  // Store collected chunks on the log
  if (collectChunks === true) {
    log.streamingChunks = {
      chunks,
      truncated: chunkIndex > MAX_CHUNKS,
    };
  }

  if (!usageCaptured) {
    completionTokens = Math.ceil(completionContent.length / 4);
  }

  log.outputTokens = completionTokens;

  const message: Record<string, unknown> = {
    role: "assistant",
    content: completionContent,
  };
  if (reasoningContent) message.reasoning_content = reasoningContent;
  if (toolCalls.size > 0) message.tool_calls = [...toolCalls.values()];

  return JSON.stringify({
    id,
    object: "chat.completion",
    created: Math.floor(Date.now() / 1000),
    model: model !== "" ? model : (fallbackModel ?? ""),
    choices: [
      {
        index: 0,
        message,
        finish_reason: finishReason,
      },
    ],
    usage: {
      prompt_tokens: promptTokens,
      completion_tokens: completionTokens,
      total_tokens: promptTokens + completionTokens,
    },
  });
}
