import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";

import {
  BudgetExceededError,
  collectStream,
  complete,
  streamCompletion,
  UnknownProviderError,
} from "./index.js";
import type { ResolvedProvider } from "./index.js";

/** Build a `text/event-stream` Response from a list of SSE event payloads. */
function sseResponse(events: string[], status = 200): Response {
  const encoder = new TextEncoder();
  const body = new ReadableStream<Uint8Array<ArrayBuffer>>({
    start(controller) {
      for (const event of events) {
        controller.enqueue(encoder.encode(`data: ${event}\n\n`));
      }
      controller.enqueue(encoder.encode("data: [DONE]\n\n"));
      controller.close();
    },
  });
  return new Response(body, {
    status,
    headers: { "content-type": "text/event-stream" },
  });
}

/** Build a JSON Response (for non-streaming completions). */
function jsonResponse(body: unknown, status = 200): Response {
  return new Response(JSON.stringify(body), {
    status,
    headers: { "content-type": "application/json" },
  });
}

const config: ResolvedProvider = {
  id: "test-openai",
  kind: "openai-compatible",
  baseUrl: "https://api.openai.com/v1",
  apiKey: "sk-test",
  model: "gpt-4o",
  contextWindow: 128_000,
};

let savedFetch: typeof fetch;

beforeEach(() => {
  savedFetch = globalThis.fetch;
});

afterEach(() => {
  globalThis.fetch = savedFetch;
  vi.restoreAllMocks();
});

describe("streamCompletion", () => {
  it("yields text deltas in order", async () => {
    globalThis.fetch = (() =>
      sseResponse([
        JSON.stringify({ choices: [{ delta: { content: "Hello" } }] }),
        JSON.stringify({ choices: [{ delta: { content: " world" } }] }),
      ])) as unknown as typeof fetch;

    const chunks: Array<{ type: string; text?: string }> = [];
    for await (const chunk of streamCompletion(
      config,
      [{ role: "user", content: "hi" }],
    )) {
      chunks.push(chunk);
    }

    const textChunks = chunks.filter((c) => c.type === "text");
    expect(textChunks.map((c) => c.text)).toEqual(["Hello", " world"]);
  });

  it("parses reasoning deltas separately from text", async () => {
    globalThis.fetch = (() =>
      sseResponse([
        JSON.stringify({ choices: [{ delta: { reasoning_content: "let me think" } }] }),
        JSON.stringify({ choices: [{ delta: { content: "the answer" } }] }),
      ])) as unknown as typeof fetch;

    const reasoning: string[] = [];
    const text: string[] = [];
    for await (const chunk of streamCompletion(
      config,
      [{ role: "user", content: "hi" }],
    )) {
      if (chunk.type === "reasoning" && chunk.reasoning) {
        reasoning.push(chunk.reasoning);
      } else if (chunk.type === "text" && chunk.text) {
        text.push(chunk.text);
      }
    }

    expect(reasoning).toEqual(["let me think"]);
    expect(text).toEqual(["the answer"]);
  });

  it("tracks usage on the final done chunk", async () => {
    globalThis.fetch = (() =>
      sseResponse([
        JSON.stringify({ choices: [{ delta: { content: "ok" } }] }),
        JSON.stringify({
          choices: [{ delta: {} }],
          usage: { prompt_tokens: 12, completion_tokens: 3, total_tokens: 15 },
        }),
      ])) as unknown as typeof fetch;

    const result = await collectStream(
      config,
      [{ role: "user", content: "hi" }],
    );
    expect(result.text).toBe("ok");
    expect(result.usage).toEqual({ promptTokens: 12, completionTokens: 3, total: 15 });
  });
});

describe("collectStream", () => {
  it("assembles full text, reasoning, and usage", async () => {
    globalThis.fetch = (() =>
      sseResponse([
        JSON.stringify({ choices: [{ delta: { reasoning_content: "hmm" } }] }),
        JSON.stringify({ choices: [{ delta: { content: "Final" } }] }),
        JSON.stringify({ choices: [{ delta: { content: " answer" } }] }),
        JSON.stringify({
          choices: [{ delta: {} }],
          usage: { prompt_tokens: 5, completion_tokens: 2, total_tokens: 7 },
        }),
      ])) as unknown as typeof fetch;

    const result = await collectStream(
      config,
      [{ role: "user", content: "q" }],
    );
    expect(result.text).toBe("Final answer");
    expect(result.reasoning).toBe("hmm");
    expect(result.usage).toEqual({ promptTokens: 5, completionTokens: 2, total: 7 });
  });
});

describe("complete (non-streaming)", () => {
  it("returns the full text and usage", async () => {
    globalThis.fetch = (() =>
      jsonResponse({
        choices: [{ message: { content: "full text" } }],
        usage: { prompt_tokens: 1, completion_tokens: 2, total_tokens: 3 },
      })) as unknown as typeof fetch;

    const result = await complete(
      config,
      [{ role: "user", content: "hi" }],
    );
    expect(result.text).toBe("full text");
    expect(result.usage).toEqual({ promptTokens: 1, completionTokens: 2, total: 3 });
  });
});

describe("retry on 429", () => {
  it("retries a rate-limited request and succeeds", async () => {
    let calls = 0;
    globalThis.fetch = ((() => {
      calls += 1;
      if (calls === 1) {
        return jsonResponse({ error: { message: "rate limited" } }, 429);
      }
      return sseResponse([JSON.stringify({ choices: [{ delta: { content: "recovered" } }] })]);
    })) as unknown as typeof fetch;

    const result = await collectStream(
      config,
      [{ role: "user", content: "hi" }],
    );
    expect(result.text).toBe("recovered");
    expect(calls).toBe(2);
  });
});

describe("budget enforcement", () => {
  it("throws BudgetExceededError before any network call", async () => {
    const tinyConfig: ResolvedProvider = { ...config, contextWindow: 8 };
    const fetchSpy = vi.fn();
    globalThis.fetch = fetchSpy as unknown as typeof fetch;

    await expect(
      collectStream(tinyConfig, [{ role: "system", content: "x".repeat(200) }]),
    ).rejects.toBeInstanceOf(BudgetExceededError);
    expect(fetchSpy).not.toHaveBeenCalled();
  });
});

describe("provider resolution", () => {
  it("throws UnknownProviderError for an unrecognized kind", async () => {
    const badConfig = { ...config, kind: "nope" } as unknown as ResolvedProvider;
    await expect(
      collectStream(badConfig, [{ role: "user", content: "hi" }]),
    ).rejects.toBeInstanceOf(UnknownProviderError);
  });
});