import type { GeneratorController } from "@lmstudio/sdk";
import type { ToolBridge } from "./toolBridge";

export async function* sseData(body: ReadableStream<Uint8Array>): AsyncGenerator<string> {
  const reader = body.getReader();
  const decoder = new TextDecoder();
  let buffer = "";
  let data: string[] = [];
  try {
    while (true) {
      const result = await reader.read();
      buffer += decoder.decode(result.value, { stream: !result.done });
      // CRLF pairs may be split across network chunks. Wait for a complete line.
      let newline: number;
      while ((newline = buffer.indexOf("\n")) !== -1) {
        const line = buffer.slice(0, newline).replace(/\r$/, "");
        buffer = buffer.slice(newline + 1);
        if (line === "") {
          if (data.length) { yield data.join("\n"); data = []; }
        } else if (line.startsWith("data:")) data.push(line.slice(5).replace(/^ /, ""));
      }
      if (buffer.length > 8 * 1024 * 1024 || data.join("\n").length > 8 * 1024 * 1024) throw new Error("Strata stream event exceeds 8 MiB.");
      if (result.done) {
        if (buffer.startsWith("data:")) data.push(buffer.slice(5).replace(/^ /, ""));
        if (data.length) yield data.join("\n");
        break;
      }
    }
  } finally {
    // A server may close/reset the transport immediately after its final event.
    // Cleanup failure must not turn an already completed answer into an error.
    try { await reader.cancel(); } catch { /* best-effort cleanup */ }
    finally { reader.releaseLock(); }
  }
}

type ToolState = { id: string; name: string; arguments: string };

export async function consumeStream(body: ReadableStream<Uint8Array>, ctl: GeneratorController, bridge?: ToolBridge): Promise<void> {
  const tools = new Map<number, ToolState>();
  let thinking = false;
  let finished = false;
  let finishReason: string | undefined;
  let pending = "";
  let mode: "probe" | "xml" | "text" = bridge?.tools.length ? "probe" : "text";
  const textFragment = (fragment: string) => {
    if (mode === "text") { ctl.fragmentGenerated(fragment, { reasoningType: "none" }); return; }
    pending += fragment;
    if (pending.length > 1024 * 1024) throw new Error("Tool-shaped answer exceeds 1 MiB.");
    const trimmed = pending.trimStart();
    if (mode === "probe" && trimmed) {
      if (trimmed.startsWith("<tool_call>")) mode = "xml";
      else if (!"<tool_call>".startsWith(trimmed)) {
        mode = "text";
        ctl.fragmentGenerated(pending, { reasoningType: "none" });
        pending = "";
      }
    }
  };
  const closeThinking = () => {
    if (thinking) { ctl.fragmentGenerated("</think>", { reasoningType: "reasoningEndTag" }); thinking = false; }
  };
  try {
    for await (const data of sseData(body)) {
      ctl.abortSignal.throwIfAborted();
      if (data === "[DONE]") { finished = true; break; }
      const chunk = JSON.parse(data);
      if (chunk.error) throw new Error(`Strata stream error: ${chunk.error.message ?? "unknown error"}`);
      const choice = chunk.choices?.[0];
      if (!choice) continue; // usage / prompt progress events
      const delta = choice.delta ?? {};
      const reasoning = delta.reasoning_content;
      if (typeof reasoning === "string" && reasoning) {
        if (!thinking) { ctl.fragmentGenerated("<think>", { reasoningType: "reasoningStartTag" }); thinking = true; }
        ctl.fragmentGenerated(reasoning, { reasoningType: "reasoning" });
      }
      if (typeof delta.content === "string" && delta.content) {
        closeThinking();
        textFragment(delta.content);
      }
      for (const call of delta.tool_calls ?? []) {
        closeThinking();
        if (!Number.isInteger(call.index) || call.index < 0) throw new Error("Invalid streamed tool index.");
        const state = tools.get(call.index) ?? { id: "", name: "", arguments: "" };
        if (call.id) state.id = call.id;
        if (call.function?.name) state.name += call.function.name;
        if (call.function?.arguments) state.arguments += call.function.arguments;
        if (state.arguments.length > 1024 * 1024) throw new Error("Tool arguments exceed 1 MiB.");
        tools.set(call.index, state);
      }
      if (choice.finish_reason) {
        finished = true;
        finishReason = choice.finish_reason;
        if (choice.finish_reason === "length" && tools.size) throw new Error("Strata ran out of output tokens during a tool call. Increase Max output tokens.");
        break; // the final choice is authoritative; usage/[DONE] is optional
      }
    }
    if (!finished) throw new Error("Strata stream disconnected before completion. Retry the request.");
    // The SDK controller has one active tool-generation slot. Buffer interleaved
    // calls by index, validate all calls, then report them sequentially.
    const complete = [...tools.entries()].sort(([a], [b]) => a - b).map(([, tool]) => {
      if (!tool.id || !tool.name) throw new Error("Incomplete tool call returned by Strata.");
      const args: unknown = JSON.parse(tool.arguments || "{}");
      if (!args || typeof args !== "object" || Array.isArray(args)) throw new Error("Tool arguments must be a JSON object.");
      const name = bridge ? bridge.original(tool.name) : tool.name;
      if (!name) throw new Error(`Strata returned an undeclared tool: ${tool.name}`);
      return { ...tool, name, parsed: args as Record<string, unknown> };
    });
    if (pending) {
      const recovered = finishReason !== "length" ? bridge?.recoverXml(pending) : undefined;
      if (recovered && !complete.length) complete.push(...recovered);
      else if (!recovered) ctl.fragmentGenerated(pending, { reasoningType: "none" });
    }
    for (const tool of complete) {
      ctl.abortSignal.throwIfAborted();
      ctl.toolCallGenerationStarted({ toolCallId: tool.id });
      // LM Studio 0.4.26 does not reset its namedReceived flag between calls
      // in one generation. The complete request below already carries the name;
      // omit this optional early-name notification to keep later calls working.
      ctl.toolCallGenerationArgumentFragmentGenerated(tool.arguments || "{}");
      ctl.toolCallGenerationEnded({ type: "function", id: tool.id, name: tool.name, arguments: tool.parsed });
    }
  } finally { closeThinking(); }
}

import type { GeneratorController } from "@lmstudio/sdk";
import type { ToolBridge } from "./toolBridge";

export async function* sseData(body: ReadableStream<Uint8Array>): AsyncGenerator<string> {
  const reader = body.getReader();
  const decoder = new TextDecoder();
  let buffer = "";
  let data: string[] = [];
  try {
    while (true) {
      const result = await reader.read();
      buffer += decoder.decode(result.value, { stream: !result.done });
      // CRLF pairs may be split across network chunks. Wait for a complete line.
      let newline: number;
      while ((newline = buffer.indexOf("\n")) !== -1) {
        const line = buffer.slice(0, newline).replace(/\r$/, "");
        buffer = buffer.slice(newline + 1);
        if (line === "") {
          if (data.length) { yield data.join("\n"); data = []; }
        } else if (line.startsWith("data:")) data.push(line.slice(5).replace(/^ /, ""));
      }
      if (buffer.length > 8 * 1024 * 1024 || data.join("\n").length > 8 * 1024 * 1024) throw new Error("Strata stream event exceeds 8 MiB.");
      if (result.done) {
        if (buffer.startsWith("data:")) data.push(buffer.slice(5).replace(/^ /, ""));
        if (data.length) yield data.join("\n");
        break;
      }
    }
  } finally {
    // A server may close/reset the transport immediately after its final event.
    // Cleanup failure must not turn an already completed answer into an error.
    try { await reader.cancel(); } catch { /* best-effort cleanup */ }
    finally { reader.releaseLock(); }
  }
}

type ToolState = { id: string; name: string; arguments: string };

export async function consumeStream(body: ReadableStream<Uint8Array>, ctl: GeneratorController, bridge?: ToolBridge): Promise<void> {
  const tools = new Map<number, ToolState>();
  let thinking = false;
  let finished = false;
  let finishReason: string | undefined;
  let pending = "";
  let mode: "probe" | "xml" | "text" = bridge?.tools.length ? "probe" : "text";
  const textFragment = (fragment: string) => {
    if (mode === "text") { ctl.fragmentGenerated(fragment, { reasoningType: "none" }); return; }
    pending += fragment;
    if (pending.length > 1024 * 1024) throw new Error("Tool-shaped answer exceeds 1 MiB.");
    const trimmed = pending.trimStart();
    if (mode === "probe" && trimmed) {
      if (trimmed.startsWith("<tool_call>")) mode = "xml";
      else if (!"<tool_call>".startsWith(trimmed)) {
        mode = "text";
        ctl.fragmentGenerated(pending, { reasoningType: "none" });
        pending = "";
      }
    }
  };
  const closeThinking = () => {
    if (thinking) { ctl.fragmentGenerated("</think>", { reasoningType: "reasoningEndTag" }); thinking = false; }
  };
  try {
    for await (const data of sseData(body)) {
      ctl.abortSignal.throwIfAborted();
      if (data === "[DONE]") { finished = true; break; }
      const chunk = JSON.parse(data);
      if (chunk.error) throw new Error(`Strata stream error: ${chunk.error.message ?? "unknown error"}`);
      const choice = chunk.choices?.[0];
      if (!choice) continue; // usage / prompt progress events
      const delta = choice.delta ?? {};
      const reasoning = delta.reasoning_content;
      if (typeof reasoning === "string" && reasoning) {
        if (!thinking) { ctl.fragmentGenerated("<think>", { reasoningType: "reasoningStartTag" }); thinking = true; }
        ctl.fragmentGenerated(reasoning, { reasoningType: "reasoning" });
      }
      if (typeof delta.content === "string" && delta.content) {
        closeThinking();
        textFragment(delta.content);
      }
      for (const call of delta.tool_calls ?? []) {
        closeThinking();
        if (!Number.isInteger(call.index) || call.index < 0) throw new Error("Invalid streamed tool index.");
        const state = tools.get(call.index) ?? { id: "", name: "", arguments: "" };
        if (call.id) state.id = call.id;
        if (call.function?.name) state.name += call.function.name;
        if (call.function?.arguments) state.arguments += call.function.arguments;
        if (state.arguments.length > 1024 * 1024) throw new Error("Tool arguments exceed 1 MiB.");
        tools.set(call.index, state);
      }
      if (choice.finish_reason) {
        finished = true;
        finishReason = choice.finish_reason;
        if (choice.finish_reason === "length" && tools.size) throw new Error("Strata ran out of output tokens during a tool call. Increase Max output tokens.");
        break; // the final choice is authoritative; usage/[DONE] is optional
      }
    }
    if (!finished) throw new Error("Strata stream disconnected before completion. Retry the request.");
    // The SDK controller has one active tool-generation slot. Buffer interleaved
    // calls by index, validate all calls, then report them sequentially.
    const complete = [...tools.entries()].sort(([a], [b]) => a - b).map(([, tool]) => {
      if (!tool.id || !tool.name) throw new Error("Incomplete tool call returned by Strata.");
      const args: unknown = JSON.parse(tool.arguments || "{}");
      if (!args || typeof args !== "object" || Array.isArray(args)) throw new Error("Tool arguments must be a JSON object.");
      const name = bridge ? bridge.original(tool.name) : tool.name;
      if (!name) throw new Error(`Strata returned an undeclared tool: ${tool.name}`);
      return { ...tool, name, parsed: args as Record<string, unknown> };
    });
    if (pending) {
      const recovered = finishReason !== "length" ? bridge?.recoverXml(pending) : undefined;
      if (recovered && !complete.length) complete.push(...recovered);
      else if (!recovered) ctl.fragmentGenerated(pending, { reasoningType: "none" });
    }
    for (const tool of complete) {
      ctl.abortSignal.throwIfAborted();
      ctl.toolCallGenerationStarted({ toolCallId: tool.id });
      // LM Studio 0.4.26 does not reset its namedReceived flag between calls
      // in one generation. The complete request below already carries the name;
      // omit this optional early-name notification to keep later calls working.
      ctl.toolCallGenerationArgumentFragmentGenerated(tool.arguments || "{}");
      ctl.toolCallGenerationEnded({ type: "function", id: tool.id, name: tool.name, arguments: tool.parsed });
    }
  } finally { closeThinking(); }
}