import type { ContentPart, ToolCallPart, UnifiedChatResponse, UnifiedStreamEvent, Usage, ProviderId, FinishReason, Citation } from "./types";

/**
 * Accumulates a unified stream into a full response. Used by adapters to implement
 * `chat()` on top of `streamChat()` and by the server to persist assistant messages.
 */
export class StreamAccumulator {
  text = "";
  reasoning = "";
  reasoningSignature?: string;
  reasoningProviderData?: unknown;
  toolCalls = new Map<string, { name: string; args: string; final?: Record<string, unknown>; providerData?: unknown }>();
  citations: Citation[] = [];
  usage?: Usage;
  finishReason: FinishReason = "other";
  providerData: Record<string, unknown> = {};
  id?: string;
  model?: string;
  error?: UnifiedStreamEvent & { type: "error" };
  firstTokenAt?: number;

  constructor(private readonly startedAt = Date.now()) {}

  push(ev: UnifiedStreamEvent) {
    switch (ev.type) {
      case "start":
        this.id = ev.id ?? this.id;
        this.model = ev.model ?? this.model;
        break;
      case "text-delta":
        if (!this.firstTokenAt && ev.text) this.firstTokenAt = Date.now();
        this.text += ev.text;
        break;
      case "reasoning-delta":
        if (!this.firstTokenAt && ev.text) this.firstTokenAt = Date.now();
        this.reasoning += ev.text;
        break;
      case "reasoning-signature":
        this.reasoningSignature = ev.signature;
        this.reasoningProviderData = ev.providerData;
        break;
      case "tool-start":
        this.toolCalls.set(ev.id, { name: ev.name, args: "" });
        break;
      case "tool-delta": {
        const tc = this.toolCalls.get(ev.id);
        if (tc) tc.args += ev.argumentsDelta;
        break;
      }
      case "tool-end": {
        const prev = this.toolCalls.get(ev.id);
        // Adapters that stream deltas emit a partial tool-end (empty argumentsText) — keep the accumulated text.
        const args = ev.argumentsText || prev?.args || "";
        this.toolCalls.set(ev.id, { name: ev.name || prev?.name || "", args, final: Object.keys(ev.arguments ?? {}).length ? ev.arguments : undefined, providerData: ev.providerData ?? prev?.providerData });
        break;
      }
      case "citation":
        this.citations.push(ev.citation);
        break;
      case "usage":
        this.usage = { ...(this.usage ?? {}), ...ev.usage } as Usage;
        break;
      case "provider-data":
        Object.assign(this.providerData, ev.data);
        break;
      case "finish":
        this.finishReason = ev.reason;
        break;
      case "error":
        this.error = ev;
        this.finishReason = "error";
        break;
    }
  }

  get ttftMs(): number | undefined {
    return this.firstTokenAt ? this.firstTokenAt - this.startedAt : undefined;
  }

  toolCallParts(): ToolCallPart[] {
    return [...this.toolCalls.entries()].map(([id, tc]) => ({
      type: "tool-call",
      id,
      name: tc.name,
      arguments: tc.final ?? safeJson(tc.args),
      argumentsText: tc.args,
      providerData: tc.providerData,
    }));
  }

  contentParts(): ContentPart[] {
    const parts: ContentPart[] = [];
    if (this.reasoning) parts.push({ type: "reasoning", text: this.reasoning, signature: this.reasoningSignature, providerData: this.reasoningProviderData });
    if (this.text) parts.push({ type: "text", text: this.text });
    parts.push(...this.toolCallParts());
    return parts;
  }

  toResponse(provider: ProviderId, model: string): UnifiedChatResponse {
    return {
      id: this.id,
      provider,
      model: this.model ?? model,
      content: this.contentParts(),
      text: this.text,
      reasoning: this.reasoning || undefined,
      toolCalls: this.toolCallParts(),
      citations: this.citations.length ? this.citations : undefined,
      usage: this.usage,
      finishReason: this.finishReason,
      latencyMs: Date.now() - this.startedAt,
      providerData: Object.keys(this.providerData).length ? this.providerData : undefined,
    };
  }
}

export function safeJson(text: string): Record<string, unknown> {
  if (!text || !text.trim()) return {};
  try {
    const v = JSON.parse(text);
    return v && typeof v === "object" ? (v as Record<string, unknown>) : { value: v };
  } catch {
    return { _raw: text };
  }
}

/** Implements `chat()` from `streamChat()` for adapters (no duplicated request-building code). */
export async function collectStream(provider: ProviderId, model: string, stream: AsyncIterable<UnifiedStreamEvent>): Promise<UnifiedChatResponse> {
  const acc = new StreamAccumulator();
  for await (const ev of stream) {
    acc.push(ev);
    if (ev.type === "error") {
      const { PolyProviderError } = await import("./types");
      throw new PolyProviderError(ev.error);
    }
  }
  return acc.toResponse(provider, model);
}