import type { GeneratorController } from "@lmstudio/sdk";
import type { GenerationStats, ToolCallState, UpstreamChunk } from "../types";
import { REASONING_CONTENT_KEY, REASONING_END_TAG, REASONING_START_TAG } from "../constants";
import { ShartiLogger } from "../utils/logger";
import { normalizeReasoningDelta, parseToolCallArguments } from "./messages";

/* -------------------------------------------------------------------------- */
/*                              Stream consumption                            */
/*                                                                            */
/* Renders upstream deltas to the controller (text / reasoning / tool calls)  */
/* and aggregates per-request statistics in a single pass.                    */
/* -------------------------------------------------------------------------- */

export async function consumeStream(
  stream: AsyncIterable<UpstreamChunk>,
  ctl: GeneratorController,
  logger: ShartiLogger,
  startTime: number,
): Promise<GenerationStats> {
  let current: ToolCallState | null = null;
  let reasoningOpen = false;

  // --- Timing tracking ---
  let firstTokenTimestamp: number | null = null;

  // --- Accumulated counters from stream chunks ---
  let inputTokens: number | null = null;
  let outputTokens: number | null = null;
  let totalTokens: number | null = null;
  const timingsAccumulator: Record<string, unknown> = {};

  function markFirstToken(): void {
    if (firstTokenTimestamp === null) firstTokenTimestamp = Date.now();
  }

  function openReasoningBlock(): void {
    if (reasoningOpen) return;
    markFirstToken();
    ctl.fragmentGenerated(REASONING_START_TAG, { reasoningType: "reasoningStartTag" });
    reasoningOpen = true;
  }

  function closeReasoningBlock(): void {
    if (!reasoningOpen) return;
    ctl.fragmentGenerated(REASONING_END_TAG, { reasoningType: "reasoningEndTag" });
    reasoningOpen = false;
  }

  function flushCurrentToolCall(): void {
    if (current === null || current.name === null) return;
    const args = parseToolCallArguments(current.arguments);
    if (args === null) {
      logger.warn(`Failed to parse tool call arguments for "${current.name}" — invalid JSON received from model.`);
      ctl.toolCallGenerationFailed(new Error(`Failed to parse tool call arguments as JSON for "${current.name}"`));
    } else {
      ctl.toolCallGenerationEnded({
        type: "function",
        name: current.name,
        arguments: args,
        id: current.id,
      });
    }
    current = null;
  }

  let lastFinishReason: string | null = null;

  for await (const chunk of stream) {
    // --- Usage / timing telemetry (carried on meta or final chunks) ---
    if (chunk.usage) {
      inputTokens = chunk.usage.prompt_tokens ?? inputTokens;
      outputTokens = chunk.usage.completion_tokens ?? outputTokens;
      totalTokens = chunk.usage.total_tokens ?? totalTokens;
    }
    if (chunk.timings) {
      Object.assign(timingsAccumulator, chunk.timings);
      // Splash reports token counts under alternate keys.
      if (timingsAccumulator.prompt_n != null) inputTokens = timingsAccumulator.prompt_n as number;
      if (timingsAccumulator.predicted_n != null) outputTokens = timingsAccumulator.predicted_n as number;
    }

    const choice = chunk.choices?.[0];
    const delta = choice?.delta;

    if (choice?.finish_reason) lastFinishReason = choice.finish_reason;

    // No delta => meta-chunk (telemetry only), nothing to render.
    if (!delta) continue;

    // --- Native reasoning content from the stream ---
    const reasoningDelta = normalizeReasoningDelta(delta[REASONING_CONTENT_KEY], REASONING_CONTENT_KEY);
    if (reasoningDelta !== null && reasoningDelta.length > 0) {
      openReasoningBlock();
      ctl.fragmentGenerated(reasoningDelta, { reasoningType: "reasoning" });
    }

    // --- Visible text ---
    if (delta.content) {
      closeReasoningBlock();
      markFirstToken();
      ctl.fragmentGenerated(delta.content);
    }

    // --- Tool-call streaming ---
    if ((delta.tool_calls?.length ?? 0) > 0) {
      closeReasoningBlock();
      markFirstToken();
    }

    for (const toolCall of delta.tool_calls ?? []) {
      if (toolCall.id !== undefined) {
        flushCurrentToolCall();
        current = { id: toolCall.id, name: null, index: toolCall.index, arguments: "" };
        ctl.toolCallGenerationStarted({ toolCallId: toolCall.id });
      }
      if (toolCall.function?.name && current) {
        current.name = toolCall.function.name;
        ctl.toolCallGenerationNameReceived(toolCall.function.name);
      }
      if (toolCall.function?.arguments && current) {
        current.arguments += toolCall.function.arguments;
        ctl.toolCallGenerationArgumentFragmentGenerated(toolCall.function.arguments);
      }
    }

    if (choice.finish_reason === "tool_calls" && current?.name) {
      flushCurrentToolCall();
    }
  }

  flushCurrentToolCall();
  closeReasoningBlock();

  // --- Surface abnormal termination reasons ---
  if (lastFinishReason === "length") {
    logger.warn("Response truncated: maximum token limit reached.");
  } else if (lastFinishReason === "content_filter") {
    logger.warn("Response blocked by content filter.");
  }

  // --- Derived metrics ---
  const endTime = Date.now();
  const totalTimeMs = endTime - startTime;
  const ttftMs = firstTokenTimestamp !== null ? firstTokenTimestamp - startTime : null;

  const totalTimeSeconds = totalTimeMs / 1000;
  const timeToFirstTokenSeconds = ttftMs !== null ? ttftMs / 1000 : null;

  const timings = Object.keys(timingsAccumulator).length > 0 ? timingsAccumulator : null;

  // Decode throughput: prefer server-reported rate, fall back to a computed estimate.
  const tokensPerSecond =
    timings?.predicted_per_second != null
      ? (timings.predicted_per_second as number)
      : outputTokens != null && totalTimeSeconds > 0
        ? outputTokens / totalTimeSeconds
        : null;

  const promptPerSecond = timings?.prompt_per_second != null ? (timings.prompt_per_second as number) : null;
  const cachedTokens = timings?.cache_n != null ? (timings.cache_n as number) : null;

  return {
    inputTokens,
    outputTokens,
    totalTokens: totalTokens ?? (inputTokens != null && outputTokens != null ? inputTokens + outputTokens : null),
    tokensPerSecond,
    timeToFirstTokenSeconds,
    totalTimeSeconds,
    promptPerSecond,
    cachedTokens,
    rawTimings: timings,
    finishReason: lastFinishReason,
  };
}

import type { GeneratorController } from "@lmstudio/sdk";
import type { GenerationStats, ToolCallState, UpstreamChunk } from "../types";
import { REASONING_CONTENT_KEY, REASONING_END_TAG, REASONING_START_TAG } from "../constants";
import { ShartiLogger } from "../utils/logger";
import { normalizeReasoningDelta, parseToolCallArguments } from "./messages";

/* -------------------------------------------------------------------------- */
/*                              Stream consumption                            */
/*                                                                            */
/* Renders upstream deltas to the controller (text / reasoning / tool calls)  */
/* and aggregates per-request statistics in a single pass.                    */
/* -------------------------------------------------------------------------- */

export async function consumeStream(
  stream: AsyncIterable<UpstreamChunk>,
  ctl: GeneratorController,
  logger: ShartiLogger,
  startTime: number,
): Promise<GenerationStats> {
  let current: ToolCallState | null = null;
  let reasoningOpen = false;

  // --- Timing tracking ---
  let firstTokenTimestamp: number | null = null;

  // --- Accumulated counters from stream chunks ---
  let inputTokens: number | null = null;
  let outputTokens: number | null = null;
  let totalTokens: number | null = null;
  const timingsAccumulator: Record<string, unknown> = {};

  function markFirstToken(): void {
    if (firstTokenTimestamp === null) firstTokenTimestamp = Date.now();
  }

  function openReasoningBlock(): void {
    if (reasoningOpen) return;
    markFirstToken();
    ctl.fragmentGenerated(REASONING_START_TAG, { reasoningType: "reasoningStartTag" });
    reasoningOpen = true;
  }

  function closeReasoningBlock(): void {
    if (!reasoningOpen) return;
    ctl.fragmentGenerated(REASONING_END_TAG, { reasoningType: "reasoningEndTag" });
    reasoningOpen = false;
  }

  function flushCurrentToolCall(): void {
    if (current === null || current.name === null) return;
    const args = parseToolCallArguments(current.arguments);
    if (args === null) {
      logger.warn(`Failed to parse tool call arguments for "${current.name}" — invalid JSON received from model.`);
      ctl.toolCallGenerationFailed(new Error(`Failed to parse tool call arguments as JSON for "${current.name}"`));
    } else {
      ctl.toolCallGenerationEnded({
        type: "function",
        name: current.name,
        arguments: args,
        id: current.id,
      });
    }
    current = null;
  }

  let lastFinishReason: string | null = null;

  for await (const chunk of stream) {
    // --- Usage / timing telemetry (carried on meta or final chunks) ---
    if (chunk.usage) {
      inputTokens = chunk.usage.prompt_tokens ?? inputTokens;
      outputTokens = chunk.usage.completion_tokens ?? outputTokens;
      totalTokens = chunk.usage.total_tokens ?? totalTokens;
    }
    if (chunk.timings) {
      Object.assign(timingsAccumulator, chunk.timings);
      // Splash reports token counts under alternate keys.
      if (timingsAccumulator.prompt_n != null) inputTokens = timingsAccumulator.prompt_n as number;
      if (timingsAccumulator.predicted_n != null) outputTokens = timingsAccumulator.predicted_n as number;
    }

    const choice = chunk.choices?.[0];
    const delta = choice?.delta;

    if (choice?.finish_reason) lastFinishReason = choice.finish_reason;

    // No delta => meta-chunk (telemetry only), nothing to render.
    if (!delta) continue;

    // --- Native reasoning content from the stream ---
    const reasoningDelta = normalizeReasoningDelta(delta[REASONING_CONTENT_KEY], REASONING_CONTENT_KEY);
    if (reasoningDelta !== null && reasoningDelta.length > 0) {
      openReasoningBlock();
      ctl.fragmentGenerated(reasoningDelta, { reasoningType: "reasoning" });
    }

    // --- Visible text ---
    if (delta.content) {
      closeReasoningBlock();
      markFirstToken();
      ctl.fragmentGenerated(delta.content);
    }

    // --- Tool-call streaming ---
    if ((delta.tool_calls?.length ?? 0) > 0) {
      closeReasoningBlock();
      markFirstToken();
    }

    for (const toolCall of delta.tool_calls ?? []) {
      if (toolCall.id !== undefined) {
        flushCurrentToolCall();
        current = { id: toolCall.id, name: null, index: toolCall.index, arguments: "" };
        ctl.toolCallGenerationStarted({ toolCallId: toolCall.id });
      }
      if (toolCall.function?.name && current) {
        current.name = toolCall.function.name;
        ctl.toolCallGenerationNameReceived(toolCall.function.name);
      }
      if (toolCall.function?.arguments && current) {
        current.arguments += toolCall.function.arguments;
        ctl.toolCallGenerationArgumentFragmentGenerated(toolCall.function.arguments);
      }
    }

    if (choice.finish_reason === "tool_calls" && current?.name) {
      flushCurrentToolCall();
    }
  }

  flushCurrentToolCall();
  closeReasoningBlock();

  // --- Surface abnormal termination reasons ---
  if (lastFinishReason === "length") {
    logger.warn("Response truncated: maximum token limit reached.");
  } else if (lastFinishReason === "content_filter") {
    logger.warn("Response blocked by content filter.");
  }

  // --- Derived metrics ---
  const endTime = Date.now();
  const totalTimeMs = endTime - startTime;
  const ttftMs = firstTokenTimestamp !== null ? firstTokenTimestamp - startTime : null;

  const totalTimeSeconds = totalTimeMs / 1000;
  const timeToFirstTokenSeconds = ttftMs !== null ? ttftMs / 1000 : null;

  const timings = Object.keys(timingsAccumulator).length > 0 ? timingsAccumulator : null;

  // Decode throughput: prefer server-reported rate, fall back to a computed estimate.
  const tokensPerSecond =
    timings?.predicted_per_second != null
      ? (timings.predicted_per_second as number)
      : outputTokens != null && totalTimeSeconds > 0
        ? outputTokens / totalTimeSeconds
        : null;

  const promptPerSecond = timings?.prompt_per_second != null ? (timings.prompt_per_second as number) : null;
  const cachedTokens = timings?.cache_n != null ? (timings.cache_n as number) : null;

  return {
    inputTokens,
    outputTokens,
    totalTokens: totalTokens ?? (inputTokens != null && outputTokens != null ? inputTokens + outputTokens : null),
    tokensPerSecond,
    timeToFirstTokenSeconds,
    totalTimeSeconds,
    promptPerSecond,
    cachedTokens,
    rawTimings: timings,
    finishReason: lastFinishReason,
  };
}