import { rec, str, num, parseProviderMetadata, extractTelemetryMetadata } from "./aiHelpers";
import type { AISpanData, DisplayItem, ToolDefinition, ToolUse } from "./types";

/**
 * Extracts structured AI span data from unflattened OTEL span properties.
 *
 * Works with the nested object produced by `unflattenAttributes()` — expects
 * keys like `gen_ai.response.model`, `ai.prompt.messages`, `trigger.llm.total_cost`, etc.
 *
 * @param properties  Unflattened span properties object
 * @param durationMs  Span duration in milliseconds
 * @returns Structured AI data, or undefined if this isn't an AI generation span
 */
export function extractAISpanData(
  properties: Record<string, unknown>,
  durationMs: number
): AISpanData | undefined {
  const genAi = properties.gen_ai;
  if (!genAi || typeof genAi !== "object") return undefined;

  const g = genAi as Record<string, unknown>;
  const ai = rec(properties.ai);
  const trigger = rec(properties.trigger);

  const gResponse = rec(g.response);
  const gRequest = rec(g.request);
  const gUsage = rec(g.usage);
  const gOperation = rec(g.operation);
  const gProvider = rec(g.provider);
  const gInput = rec(g.input);
  const gOutput = rec(g.output);
  const gTool = rec(g.tool);
  const aiModel = rec(ai.model);
  const aiResponse = rec(ai.response);
  const aiPrompt = rec(ai.prompt);
  const aiUsage = rec(ai.usage);
  const triggerLlm = rec(trigger.llm);

  const model = str(gResponse.model) ?? str(gRequest.model) ?? str(aiModel.id);
  if (!model) return undefined;

  // Prefer ai.usage (richer) over gen_ai.usage.
  // Gateway/some providers emit promptTokens/completionTokens instead of inputTokens/outputTokens.
  const inputTokens =
    num(aiUsage.inputTokens) ?? num(aiUsage.promptTokens) ?? num(gUsage.input_tokens) ?? 0;
  const outputTokens =
    num(aiUsage.outputTokens) ?? num(aiUsage.completionTokens) ?? num(gUsage.output_tokens) ?? 0;
  const totalTokens = num(aiUsage.totalTokens) ?? inputTokens + outputTokens;

  const tokensPerSecond =
    num(aiResponse.avgOutputTokensPerSecond) ??
    (outputTokens > 0 && durationMs > 0
      ? Math.round((outputTokens / (durationMs / 1000)) * 10) / 10
      : undefined);

  // AI SDK 7 moves span emission into `@ai-sdk/otel`, which emits OTel GenAI
  // semantic-convention attributes (gen_ai.input/output.messages, gen_ai.provider.name,
  // gen_ai.tool.definitions, gen_ai.response.finish_reasons). AI SDK 6 emits the older
  // `ai.*` keys (ai.prompt.messages, ai.response.text/toolCalls/finishReason). Customers
  // run either major, so detect the shape and read both.
  const isV7 =
    typeof gInput.messages === "string" ||
    typeof gOutput.messages === "string" ||
    typeof gProvider.name === "string";

  const toolDefs = parseToolDefinitions(isV7 ? gTool.definitions : aiPrompt.tools);
  const providerMeta = parseProviderMetadata(aiResponse.providerMetadata);
  const aiTelemetry = rec(ai.telemetry);
  const telemetryMetaRaw = rec(aiTelemetry.metadata);
  const promptMeta = rec(telemetryMetaRaw.prompt);
  const promptSlug = str(promptMeta.slug);
  const promptVersion = str(promptMeta.version);
  const promptModel = str(promptMeta.model);
  const promptLabels = str(promptMeta.labels);
  const promptInput = str(promptMeta.input);
  const telemetryMeta = extractTelemetryMetadata(aiTelemetry.metadata);

  return {
    model,
    provider: str(gProvider.name) ?? str(g.system) ?? str(aiModel.provider) ?? "unknown",
    operationName: str(gOperation.name) ?? str(ai.operationId) ?? "",
    responseId: str(gResponse.id) || undefined,
    finishReason: isV7 ? firstFinishReason(gResponse.finish_reasons) : str(aiResponse.finishReason),
    serviceTier: providerMeta?.serviceTier,
    resolvedProvider: providerMeta?.resolvedProvider,
    toolChoice: parseToolChoice(aiPrompt.toolChoice),
    toolCount: toolDefs?.length,
    messageCount: isV7 ? countGenAiMessages(gInput.messages) : countMessages(aiPrompt.messages),
    telemetryMetadata: telemetryMeta,
    promptSlug: promptSlug || undefined,
    promptVersion: promptVersion || undefined,
    promptModel: promptModel || undefined,
    promptLabels: promptLabels || undefined,
    promptInput: promptInput || undefined,
    inputTokens,
    outputTokens,
    totalTokens,
    cachedTokens:
      num(aiUsage.cachedInputTokens) ??
      num(gUsage.cache_read_input_tokens) ??
      num(rec(gUsage.cache_read).input_tokens),
    cacheCreationTokens:
      num(aiUsage.cacheCreationInputTokens) ??
      num(gUsage.cache_creation_input_tokens) ??
      num(rec(gUsage.cache_creation).input_tokens),
    reasoningTokens: num(aiUsage.reasoningTokens) ?? num(gUsage.reasoning_tokens),
    tokensPerSecond,
    msToFirstChunk: num(aiResponse.msToFirstChunk),
    durationMs,
    inputCost: num(triggerLlm.input_cost),
    outputCost: num(triggerLlm.output_cost),
    totalCost: num(triggerLlm.total_cost),
    cachedCost: num(triggerLlm.cached_cost),
    cacheCreationCost: num(triggerLlm.cache_creation_cost),
    responseText: isV7
      ? extractGenAiAssistantText(gOutput.messages) || undefined
      : str(aiResponse.text) || undefined,
    responseObject: str(aiResponse.object) || undefined,
    toolDefinitions: toolDefs,
    items: isV7
      ? buildGenAiDisplayItems(g.system_instructions, gInput.messages, gOutput.messages, toolDefs)
      : buildDisplayItems(aiPrompt.messages, aiResponse.toolCalls, toolDefs),
  };
}

// ---------------------------------------------------------------------------
// Message → DisplayItem transformation
// ---------------------------------------------------------------------------

type RawMessage = {
  role: string;
  content: unknown;
  toolCallId?: string;
  name?: string;
};

/**
 * Build display items from prompt messages and optionally response tool calls.
 * - Parses ai.prompt.messages and merges consecutive tool-call + tool-result pairs
 * - If ai.response.toolCalls is present (finishReason=tool-calls), appends those too
 */
function buildDisplayItems(
  messagesRaw: unknown,
  responseToolCallsRaw: unknown,
  toolDefs?: ToolDefinition[]
): DisplayItem[] | undefined {
  const items = parseMessagesToDisplayItems(messagesRaw);
  const responseToolCalls = parseResponseToolCalls(responseToolCallsRaw);

  if (!items && !responseToolCalls) return undefined;

  const result = items ?? [];

  if (responseToolCalls && responseToolCalls.length > 0) {
    result.push({ type: "tool-use", tools: responseToolCalls });
  }

  if (toolDefs && toolDefs.length > 0) {
    const defsByName = new Map(toolDefs.map((d) => [d.name, d]));
    for (const item of result) {
      if (item.type === "tool-use") {
        for (const tool of item.tools) {
          const def = defsByName.get(tool.toolName);
          if (def) {
            tool.description = def.description;
            tool.parametersJson = def.parametersJson;
          }
        }
      }
    }
  }

  return result.length > 0 ? result : undefined;
}

function parseMessagesToDisplayItems(raw: unknown): DisplayItem[] | undefined {
  if (typeof raw !== "string") return undefined;

  let messages: RawMessage[];
  try {
    const parsed = JSON.parse(raw);
    if (!Array.isArray(parsed)) return undefined;
    messages = parsed.map((item: unknown) => {
      const m = rec(item);
      return {
        role: str(m.role) ?? "user",
        content: m.content,
        toolCallId: str(m.toolCallId),
        name: str(m.name),
      };
    });
  } catch {
    return undefined;
  }

  const items: DisplayItem[] = [];
  let i = 0;

  while (i < messages.length) {
    const msg = messages[i];

    if (msg.role === "system") {
      items.push({ type: "system", text: extractTextContent(msg.content) });
      i++;
      continue;
    }

    if (msg.role === "user") {
      items.push({ type: "user", text: extractTextContent(msg.content) });
      i++;
      continue;
    }

    // Assistant message — check if it contains tool calls
    if (msg.role === "assistant") {
      const toolCalls = extractToolCalls(msg.content);

      if (toolCalls.length > 0) {
        // Collect subsequent tool result messages that match these tool calls
        const _toolCallIds = new Set(toolCalls.map((tc) => tc.toolCallId));
        let j = i + 1;
        while (j < messages.length && messages[j].role === "tool") {
          j++;
        }
        // Gather tool result messages between i+1 and j
        const toolResultMsgs = messages.slice(i + 1, j);

        // Build ToolUse entries by pairing calls with results
        const tools: ToolUse[] = toolCalls.map((tc) => {
          const resultMsg = toolResultMsgs.find((m) => {
            // Match by toolCallId in the message's content parts
            const results = extractToolResults(m.content);
            return results.some((r) => r.toolCallId === tc.toolCallId);
          });

          const result = resultMsg
            ? extractToolResults(resultMsg.content).find((r) => r.toolCallId === tc.toolCallId)
            : undefined;

          return {
            toolCallId: tc.toolCallId,
            toolName: tc.toolName,
            inputJson: JSON.stringify(tc.input, null, 2),
            resultSummary: result?.summary,
            resultOutput: result?.formattedOutput,
          };
        });

        items.push({ type: "tool-use", tools });
        i = j; // skip past the tool result messages
        continue;
      }

      // Assistant message with just text
      const text = extractTextContent(msg.content);
      if (text) {
        items.push({ type: "assistant", text });
      }
      i++;
      continue;
    }

    // Skip any other message types (tool messages that weren't consumed above)
    i++;
  }

  return items.length > 0 ? items : undefined;
}

// ---------------------------------------------------------------------------
// Response tool calls (from ai.response.toolCalls, used when finishReason=tool-calls)
// ---------------------------------------------------------------------------

/**
 * Parse ai.response.toolCalls JSON string into ToolUse entries.
 * These are tool calls the model requested but haven't been executed yet in this span.
 */
function parseResponseToolCalls(raw: unknown): ToolUse[] | undefined {
  if (typeof raw !== "string") return undefined;
  try {
    const parsed = JSON.parse(raw);
    if (!Array.isArray(parsed)) return undefined;
    const tools: ToolUse[] = [];
    for (const item of parsed) {
      const tc = rec(item);
      if (tc.type === "tool-call" || tc.toolName || tc.toolCallId) {
        tools.push({
          toolCallId: str(tc.toolCallId) ?? "",
          toolName: str(tc.toolName) ?? "",
          inputJson: JSON.stringify(
            tc.input && typeof tc.input === "object" ? tc.input : {},
            null,
            2
          ),
        });
      }
    }
    return tools.length > 0 ? tools : undefined;
  } catch {
    return undefined;
  }
}

// ---------------------------------------------------------------------------
// Content part extraction
// ---------------------------------------------------------------------------

function extractTextContent(content: unknown): string {
  if (typeof content === "string") return content;
  if (!Array.isArray(content)) return "";

  const texts: string[] = [];
  for (const raw of content) {
    const p = rec(raw);
    if (p.type === "text" && typeof p.text === "string") {
      texts.push(p.text);
    } else if (typeof p.text === "string") {
      texts.push(p.text);
    }
  }
  return texts.join("\n");
}

type ParsedToolCall = {
  toolCallId: string;
  toolName: string;
  input: Record<string, unknown>;
};

function extractToolCalls(content: unknown): ParsedToolCall[] {
  if (!Array.isArray(content)) return [];
  const calls: ParsedToolCall[] = [];
  for (const raw of content) {
    const p = rec(raw);
    if (p.type === "tool-call") {
      calls.push({
        toolCallId: str(p.toolCallId) ?? "",
        toolName: str(p.toolName) ?? "",
        input: p.input && typeof p.input === "object" ? (p.input as Record<string, unknown>) : {},
      });
    }
  }
  return calls;
}

type ParsedToolResult = {
  toolCallId: string;
  toolName: string;
  summary: string;
  formattedOutput: string;
};

function extractToolResults(content: unknown): ParsedToolResult[] {
  if (!Array.isArray(content)) return [];
  const results: ParsedToolResult[] = [];
  for (const raw of content) {
    const p = rec(raw);
    if (p.type === "tool-result") {
      const { summary, formattedOutput } = summarizeToolOutput(p.output);
      results.push({
        toolCallId: str(p.toolCallId) ?? "",
        toolName: str(p.toolName) ?? "",
        summary,
        formattedOutput,
      });
    }
  }
  return results;
}

/**
 * Summarize a tool output into a short label and a formatted string for display.
 * Handles the AI SDK's `{ type: "json", value: { status, contentType, body, truncated } }` shape.
 */
function summarizeToolOutput(output: unknown): { summary: string; formattedOutput: string } {
  if (typeof output === "string") {
    return {
      summary: output.length > 80 ? output.slice(0, 80) + "..." : output,
      formattedOutput: output,
    };
  }

  if (!output || typeof output !== "object") {
    return { summary: "result", formattedOutput: JSON.stringify(output, null, 2) };
  }

  const o = output as Record<string, unknown>;

  // AI SDK wraps tool results as { type: "json", value: { status, contentType, body, ... } }
  if (o.type === "json" && o.value && typeof o.value === "object") {
    const v = o.value as Record<string, unknown>;
    const parts: string[] = [];
    if (typeof v.status === "number") parts.push(`${v.status}`);
    if (typeof v.contentType === "string") parts.push(v.contentType);
    if (v.truncated === true) parts.push("truncated");
    return {
      summary: parts.length > 0 ? parts.join(" · ") : "json result",
      formattedOutput: JSON.stringify(v, null, 2),
    };
  }

  return { summary: "result", formattedOutput: JSON.stringify(output, null, 2) };
}

// ---------------------------------------------------------------------------
// Tool definitions (from ai.prompt.tools)
// ---------------------------------------------------------------------------

/**
 * Parse ai.prompt.tools — after the array fix, this arrives as a JSON array string
 * where each element is itself a JSON string of a tool definition.
 */
function parseToolDefinitions(raw: unknown): ToolDefinition[] | undefined {
  if (typeof raw !== "string") return undefined;
  try {
    const parsed = JSON.parse(raw);
    if (!Array.isArray(parsed)) return undefined;
    const defs: ToolDefinition[] = [];
    for (const item of parsed) {
      // Each item is either a JSON string or already an object
      const obj = typeof item === "string" ? JSON.parse(item) : item;
      if (!obj || typeof obj !== "object") continue;
      const o = obj as Record<string, unknown>;
      const name = str(o.name);
      if (!name) continue;
      const schema = o.parameters ?? o.inputSchema;
      defs.push({
        name,
        description: str(o.description),
        parametersJson:
          schema && typeof schema === "object" ? JSON.stringify(schema, null, 2) : undefined,
      });
    }
    return defs.length > 0 ? defs : undefined;
  } catch {
    return undefined;
  }
}

// ---------------------------------------------------------------------------
// Tool choice parsing
// ---------------------------------------------------------------------------

function parseToolChoice(raw: unknown): string | undefined {
  if (typeof raw !== "string") return undefined;
  try {
    const parsed = JSON.parse(raw);
    if (typeof parsed === "string") return parsed;
    if (parsed && typeof parsed === "object") {
      const obj = parsed as Record<string, unknown>;
      if (typeof obj.type === "string") return obj.type;
    }
    return undefined;
  } catch {
    return undefined;
  }
}

// ---------------------------------------------------------------------------
// Message count
// ---------------------------------------------------------------------------

function countMessages(raw: unknown): number | undefined {
  if (typeof raw !== "string") return undefined;
  try {
    const parsed = JSON.parse(raw);
    if (!Array.isArray(parsed)) return undefined;
    return parsed.length > 0 ? parsed.length : undefined;
  } catch {
    return undefined;
  }
}

// ---------------------------------------------------------------------------
// AI SDK 7 — @ai-sdk/otel GenAI semantic-convention shape
// ---------------------------------------------------------------------------
//
// v7 moved span emission out of `ai` core into `@ai-sdk/otel`, which emits OTel
// GenAI semantic-convention attributes instead of the v6 `ai.*` keys:
//   gen_ai.input.messages        JSON string: [{ role, parts: [...] }]
//   gen_ai.output.messages       JSON string: [{ role:"assistant", parts, finish_reason }]
//   gen_ai.system_instructions   system prompt (plain string, or [{ type:"text", content }])
// Message parts (per @ai-sdk/otel's convertMessagePartToSemConv):
//   { type:"text" | "reasoning", content }
//   { type:"tool_call", id, name, arguments }
//   { type:"tool_call_response", id, response }   // response already unwrapped from the AI SDK envelope
// Media / approval / custom parts are not surfaced in the display yet.

type GenAiMessage = { role: string; parts: Record<string, unknown>[]; finishReason?: string };

function parseGenAiMessages(raw: unknown): GenAiMessage[] | undefined {
  if (typeof raw !== "string") return undefined;
  try {
    const parsed = JSON.parse(raw);
    if (!Array.isArray(parsed)) return undefined;
    return parsed.map((m) => {
      const o = rec(m);
      return {
        role: str(o.role) ?? "user",
        parts: Array.isArray(o.parts) ? o.parts.map(rec) : [],
        finishReason: str(o.finish_reason),
      };
    });
  } catch {
    return undefined;
  }
}

/** `gen_ai.response.finish_reasons` arrives as a JSON array string (e.g. `["stop"]`); take the first. */
function firstFinishReason(raw: unknown): string | undefined {
  if (typeof raw !== "string") return undefined;
  try {
    const parsed = JSON.parse(raw);
    if (Array.isArray(parsed)) {
      const first = parsed.find((r) => typeof r === "string");
      return typeof first === "string" ? first : undefined;
    }
    if (typeof parsed === "string") return parsed || undefined;
  } catch {
    return raw || undefined;
  }
  return undefined;
}

function countGenAiMessages(raw: unknown): number | undefined {
  const msgs = parseGenAiMessages(raw);
  return msgs && msgs.length > 0 ? msgs.length : undefined;
}

/** Concatenated text of all `text` parts across the assistant output messages. */
function extractGenAiAssistantText(outputRaw: unknown): string {
  const msgs = parseGenAiMessages(outputRaw);
  if (!msgs) return "";
  const texts: string[] = [];
  for (const m of msgs) {
    if (m.role !== "assistant") continue;
    for (const p of m.parts) {
      if (p.type === "text" && typeof p.content === "string") texts.push(p.content);
    }
  }
  return texts.join("\n");
}

/** Plain text of a message's `text` parts (reasoning parts aren't surfaced yet). */
function genAiMessageText(parts: Record<string, unknown>[]): string {
  const texts: string[] = [];
  for (const p of parts) {
    if (p.type === "text" && typeof p.content === "string") texts.push(p.content);
  }
  return texts.join("\n");
}

/** Parse `gen_ai.system_instructions` (plain string, or a JSON array of `{ type:"text", content }`). */
function parseSystemInstructions(raw: unknown): string | undefined {
  if (typeof raw !== "string") return undefined;
  const trimmed = raw.trim();
  if (trimmed.startsWith("[")) {
    try {
      const parsed = JSON.parse(trimmed);
      if (Array.isArray(parsed)) {
        const text = parsed
          .map((p) => str(rec(p).content))
          .filter((t): t is string => Boolean(t))
          .join("\n");
        return text || undefined;
      }
    } catch {
      // fall through to raw
    }
  }
  return raw || undefined;
}

/**
 * Build display items from the v7 GenAI message attributes: the system prompt,
 * the input message history, and the assistant's output for this span. Assistant
 * tool_call parts are paired with tool_call_response parts from following tool messages.
 */
function buildGenAiDisplayItems(
  systemInstructionsRaw: unknown,
  inputMessagesRaw: unknown,
  outputMessagesRaw: unknown,
  toolDefs?: ToolDefinition[]
): DisplayItem[] | undefined {
  const items: DisplayItem[] = [];

  const systemText = parseSystemInstructions(systemInstructionsRaw);
  if (systemText) items.push({ type: "system", text: systemText });

  const messages = [
    ...(parseGenAiMessages(inputMessagesRaw) ?? []),
    ...(parseGenAiMessages(outputMessagesRaw) ?? []),
  ];
  appendGenAiMessages(items, messages);

  if (toolDefs && toolDefs.length > 0) {
    const defsByName = new Map(toolDefs.map((d) => [d.name, d]));
    for (const item of items) {
      if (item.type === "tool-use") {
        for (const tool of item.tools) {
          const def = defsByName.get(tool.toolName);
          if (def) {
            tool.description = def.description;
            tool.parametersJson = def.parametersJson;
          }
        }
      }
    }
  }

  return items.length > 0 ? items : undefined;
}

function appendGenAiMessages(items: DisplayItem[], messages: GenAiMessage[]): void {
  let i = 0;
  while (i < messages.length) {
    const msg = messages[i];

    if (msg.role === "system") {
      const text = genAiMessageText(msg.parts);
      if (text) items.push({ type: "system", text });
      i++;
      continue;
    }

    if (msg.role === "user") {
      const text = genAiMessageText(msg.parts);
      if (text) items.push({ type: "user", text });
      i++;
      continue;
    }

    if (msg.role === "assistant") {
      const text = genAiMessageText(msg.parts);
      if (text) items.push({ type: "assistant", text });

      const toolCalls = msg.parts.filter((p) => p.type === "tool_call");
      if (toolCalls.length > 0) {
        // Collect tool_call_response parts from the tool messages that follow.
        const responsesById = new Map<string, unknown>();
        let j = i + 1;
        while (j < messages.length && messages[j].role === "tool") {
          for (const p of messages[j].parts) {
            if (p.type === "tool_call_response") {
              const id = str(p.id);
              if (id) responsesById.set(id, p.response);
            }
          }
          j++;
        }

        const tools: ToolUse[] = toolCalls.map((tc) => {
          const id = str(tc.id) ?? "";
          let resultSummary: string | undefined;
          let resultOutput: string | undefined;
          if (id && responsesById.has(id)) {
            const summarized = summarizeToolOutput(responsesById.get(id));
            resultSummary = summarized.summary;
            resultOutput = summarized.formattedOutput;
          }
          return {
            toolCallId: id,
            toolName: str(tc.name) ?? "",
            inputJson: JSON.stringify(tc.arguments ?? {}, null, 2),
            resultSummary,
            resultOutput,
          };
        });

        items.push({ type: "tool-use", tools });
        i = j;
        continue;
      }

      i++;
      continue;
    }

    // tool-role messages are consumed via the assistant pairing above; skip stragglers.
    i++;
  }
}