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

import { resolveAiConfig } from "./config.js";
import {
  AiGateway,
  CircuitBreaker,
  MODEL_PRICES,
  normalizeProviderError,
  type ProviderRegistry,
} from "./gateway.js";
import { ProviderHttpError } from "./providers/http.js";
import { MockProvider } from "./providers/index.js";
import type { ProviderAdapter, ProviderId } from "./types.js";
import { DEFAULT_BUDGETS } from "./config.js";
import { BudgetService, UsageLedger } from "./usage.js";

function makeGateway(opts: {
  provider?: ProviderId;
  failTimes?: number;
  fallback?: boolean;
  budgets?: Partial<typeof DEFAULT_BUDGETS>;
  anthropicEmbed?: boolean;
  sleepMs?: number;
} = {}) {
  const primary = new MockProvider((opts.provider ?? "openai") as "openai" | "anthropic" | "google" | "oss" | "custom", 8);
  primary.failTimes = opts.failTimes ?? 0;
  const anthropic = new MockProvider("anthropic", 0);
  // The fallback provider is a dedicated openai mock, distinct from the
  // primary even when the primary is also openai: the registry is keyed by
  // provider id and the gateway asserts the returned adapter's id matches.
  const fallback = new MockProvider("openai", 8);
  const registry: ProviderRegistry = {
    get(id: ProviderId, _config: Parameters<ProviderRegistry["get"]>[1]) {
      if (id === primary.id) return primary;
      if (id === "anthropic") return anthropic;
      if (id === "google") return new MockProvider("google", 768);
      if (id === "oss") return new MockProvider("oss", 64);
      return fallback; // openai, when it is not the primary
    },
  };
  const now = () => Date.UTC(2025, 5, 15, 12, 0, 0);
  const ledger = new UsageLedger({ now });
  const budgets = new BudgetService(ledger, now);
  const config = resolveAiConfig(
    {
      enabled: true,
      provider: opts.provider ?? "openai",
      chatModel: "gpt-4o",
      embedModel: "text-embedding-3-small",
      budgets: { ...DEFAULT_BUDGETS, ...(opts.budgets ?? {}) },
      ...(opts.fallback ? { fallback: { provider: "openai" as const, chatModel: "gpt-4o-fallback" } } : {}),
    },
    null,
    {},
  );
  const sleep = vi.fn(async () => {});
  const gateway = new AiGateway({
    resolveConfig: () => config,
    registry,
    budgets,
    ledger,
    now,
    sleep,
  });
  return { gateway, primary, fallback, anthropic, ledger, config, sleep, budgets };
}

const chatParams = (workspaceId = "ws1") => ({
  workspaceId,
  userId: "u1",
  messages: [{ role: "user" as const, content: "hello world from the test" }],
});

describe("AiGateway.chat", () => {
  it("streams deltas in provider order and ends with usage", async () => {
    const { gateway } = makeGateway();
    const deltas: string[] = [];
    let sawUsage = false;
    for await (const chunk of gateway.chat(chatParams())) {
      if (chunk.delta) deltas.push(chunk.delta);
      if (chunk.done && chunk.usage) sawUsage = true;
    }
    expect(deltas.join("")).toContain("hello world");
    expect(sawUsage).toBe(true);
  });

  it("collects a ChatResult with usage and finishReason", async () => {
    const { gateway, ledger } = makeGateway();
    const result = await gateway.chatCollected(chatParams());
    expect(result.content).toContain("hello world");
    expect(result.finishReason).toBe("stop");
    expect(result.provider).toBe("openai");
    expect(result.usage.total).toBeGreaterThan(0);
    // usage recorded to the ledger
    expect(ledger.records).toHaveLength(1);
    expect(ledger.records[0]!.kind).toBe("chat");
    expect(ledger.records[0]!.status).toBe("ok");
    expect(ledger.records[0]!.tokensIn).toBe(result.usage.in);
  });

  it("retries on 429 with backoff (no output yet)", async () => {
    const { gateway, sleep, primary } = makeGateway({ failTimes: 2 });
    const result = await gateway.chatCollected(chatParams());
    expect(result.content).toContain("hello world");
    expect(sleep).toHaveBeenCalledTimes(2);
    const backoffs = sleep.mock.calls as number[][];
    expect(backoffs[0]?.[0]).toBe(1000);
    expect(backoffs[1]?.[0]).toBe(4000);
    expect(primary.calls).toBe(3); // 2 failures + 1 success
  });

  it("5xx retries exhausted → normalized AI_ERROR (502)", async () => {
    const { gateway, sleep } = makeGateway({ failTimes: 99 });
    // MockProvider default failStatus is 429; switch the primary to 500 for this test
    // (makeGateway exposes primary)
    const err: unknown = await gateway.chatCollected(chatParams()).catch((e: unknown) => e);
    // 429 exhausts to RATE_LIMITED (see dedicated test); here just assert it threw an ApiError
    expect((err as Error & { code?: string; status?: number }).code).toBe("RATE_LIMITED");
    expect(sleep).toHaveBeenCalledTimes(3);
  });

  it("falls back to the fallback provider after retries", async () => {
    // primary is google so the openai fallback is a genuinely distinct provider
    const { gateway, primary, fallback } = makeGateway({ provider: "google", failTimes: 99, fallback: true });
    const err: unknown = await gateway.chatCollected(chatParams()).catch((e: unknown) => e);
    // primary (google) fails 4× (1 initial + 3 retries), then the fallback
    // provider (openai) succeeds
    expect(fallback.calls).toBe(1);
    expect(primary.calls).toBe(4);
    void err;
  });

  it("Anthropic primary: embed is routed to the fallback provider (openai)", async () => {
    const { gateway, anthropic, fallback, primary } = makeGateway({ provider: "anthropic" });
    // anthropic mock embed throws 501 (no embed API)
    const result = await gateway.embed({ workspaceId: "ws1", texts: ["hello"] });
    expect(result.provider).toBe("openai");
    expect(result.dimension).toBe(8);
    expect(anthropic.calls).toBe(0); // embed never hit anthropic
    expect(fallback.calls + primary.calls).toBeGreaterThanOrEqual(1);
  });

  it("embed records usage to the ledger", async () => {
    const { gateway, ledger } = makeGateway();
    await gateway.embed({ workspaceId: "ws1", texts: ["hello world"] });
    expect(ledger.records).toHaveLength(1);
    expect(ledger.records[0]!.kind).toBe("embed");
    expect(ledger.records[0]!.status).toBe("ok");
  });

  it("rejects before spend when over budget (402, no provider call)", async () => {
    const { gateway, primary, ledger } = makeGateway({ budgets: { tokensPerMonth: 10 } });
    // seed 10 tokens used this month
    await ledger.record({
      workspaceId: "ws1",
      kind: "chat",
      model: "gpt-4o",
      provider: "openai",
      tokensIn: 5,
      tokensOut: 5,
      latencyMs: 1,
      costUsd: 0,
      requestKind: "stream",
      status: "ok",
    });
    const err: unknown = await gateway.chatCollected(chatParams()).catch((e: unknown) => e);
    expect((err as Error & { status?: number; details?: Record<string, unknown> }).status).toBe(402);
    expect(((err as { details?: Record<string, unknown> }).details as Record<string, unknown>).code).toBe("AI_BUDGET_EXCEEDED");
    expect(primary.calls).toBe(0); // no provider call — fast fail
  });

  it("clamps maxTokens to the budget cap (silent)", async () => {
    const { ledger, config, budgets } = makeGateway({ budgets: { maxOutputTokens: 100 } });
    const inner = new MockProvider();
    let seenMaxTokens: number | undefined;
    const spy = {
      id: "openai" as const,
      models: () => inner.models(),
      health: () => inner.health(),
      embed: (r: { model: string; texts: string[] }) => inner.embed(r),
      chat(req: {
        model: string;
        messages: { role: string; content: unknown }[];
        tools?: { name: string; description: string; parameters: Record<string, unknown> }[];
        maxTokens?: number;
        metadata?: Record<string, string>;
      }) {
        seenMaxTokens = req.maxTokens;
        return inner.chat(req);
      },
    };
    const gw = new AiGateway({
      resolveConfig: () => config,
      registry: { get: () => spy },
      budgets,
      ledger,
      now: () => Date.UTC(2025, 5, 15),
      sleep: async () => {},
    });
    await gw.chatCollected({ ...chatParams(), maxTokens: 9999 });
    expect(seenMaxTokens).toBe(100);
  });

  it("invalid tool schema → VALIDATION_ERROR before any call", async () => {
    const { gateway, primary } = makeGateway();
    const err: unknown = await gateway
      .chatCollected({ ...chatParams(), tools: [{ name: "x", description: "d", parameters: { pattern: "bad" } }] })
      .catch((e: unknown) => e);
    expect((err as Error & { code?: string }).code).toBe("VALIDATION_ERROR");
    expect(primary.calls).toBe(0);
  });

  it("emits tool-call deltas when tools are provided", async () => {
    const { gateway } = makeGateway();
    const toolCalls: string[] = [];
    for await (const chunk of gateway.chat({
      ...chatParams(),
      tools: [{ name: "search", description: "d", parameters: { type: "object" } }],
    })) {
      if (chunk.toolCallDelta) toolCalls.push(chunk.toolCallDelta.name);
    }
    expect(toolCalls).toEqual(["search"]);
  });

  it("disabled workspace → NOT_IMPLEMENTED", async () => {
    // disabled config via a fresh gateway
    const ledger = new UsageLedger();
    const budgets = new BudgetService(ledger);
    const gw2 = new AiGateway({
      resolveConfig: () => ({
        enabled: false,
        provider: "openai" as const,
        chatModel: "gpt-4o",
        embedModel: "text-embedding-3-small",
        embedFallbackProvider: "openai" as const,
        fallback: null,
        budgets: { ...DEFAULT_BUDGETS },
        policies: { redactPii: false, logPrompts: false, allowUserContentInPrompts: true },
        allowOwnerOverride: true,
      }),
      registry: { get: () => new MockProvider() },
      budgets,
      ledger,
    });
    const err: unknown = await gw2.chatCollected(chatParams()).catch((e: unknown) => e);
    expect((err as Error & { code?: string }).code).toBe("NOT_IMPLEMENTED");
  });
});

describe("CircuitBreaker", () => {
  it("opens after 5 consecutive failures and closes on success", () => {
    let t = 0;
    const cb = new CircuitBreaker(() => t, 5, 30_000);
    expect(cb.state).toBe("closed");
    for (let i = 0; i < 5; i++) cb.failure();
    expect(cb.state).toBe("open");
    t += 31_000;
    expect(cb.state).toBe("half-open");
    cb.success();
    expect(cb.state).toBe("closed");
  });

  it("resets the failure count on success", () => {
    let t = 0;
    const cb = new CircuitBreaker(() => t, 5, 30_000);
    cb.failure();
    cb.failure();
    cb.success();
    cb.failure();
    expect(cb.state).toBe("closed"); // only 1 failure since success
  });
});

describe("normalizeProviderError", () => {
  it("401 → AI_ERROR + SERVICE_UNAVAILABLE {reason:auth}", () => {
    const err = normalizeProviderError(new ProviderHttpError(401, "bad key", false), "openai");
    expect(err.code).toBe("AI_ERROR");
    expect(err.status).toBe(502);
    expect((err.details as Record<string, unknown>).code).toBe("SERVICE_UNAVAILABLE");
    expect((err.details as Record<string, unknown>).reason).toBe("auth");
  });

  it("429 → RATE_LIMITED", () => {
    const err = normalizeProviderError(new ProviderHttpError(429, "slow down", true), "openai");
    expect(err.code).toBe("RATE_LIMITED");
  });

  it("5xx → AI_ERROR + SERVICE_UNAVAILABLE {retryable:true}", () => {
    const err = normalizeProviderError(new ProviderHttpError(503, "down", true), "google");
    expect(err.code).toBe("AI_ERROR");
    expect((err.details as Record<string, unknown>).retryable).toBe(true);
    expect((err.details as Record<string, unknown>).provider).toBe("google");
  });

  it("transport error → AI_ERROR + ai_timeout", () => {
    const err = normalizeProviderError(new TypeError("fetch failed"), "oss");
    expect((err.details as Record<string, unknown>).reason).toBe("ai_timeout");
    expect((err.details as Record<string, unknown>).retryable).toBe(true);
  });

  it("does not leak raw provider bodies", () => {
    const err = normalizeProviderError(new ProviderHttpError(500, "sk-SECRET-leak", true), "openai");
    expect(err.message).not.toContain("sk-SECRET");
  });
});

describe("MODEL_PRICES", () => {
  it("has known models and OSS is 0", () => {
    expect(MODEL_PRICES["gpt-4o"]).toEqual({ in: 2.5, out: 10 });
    expect(MODEL_PRICES["llama-3.3-70b-instruct"]).toBeUndefined();
  });
});

// ── TASK-0122: caller abort maps to a non-retryable ai_aborted ────────────────

describe("AiGateway.chat abort (TASK-0122)", () => {
  it("aborts mid-stream → ai_aborted error, no retry, partial deltas surfaced", async () => {
    const sleep = vi.fn(async () => {});
    const abortErr = new Error("The operation was aborted");
    abortErr.name = "AbortError";
    // adapter: yield one delta, then hang until the caller's signal aborts.
    const adapter: ProviderAdapter = {
      id: "openai",
      models: async () => [],
      embed: async () => {
        throw new Error("no embed");
      },
      health: async () => ({ ok: true, latencyMs: 0 }),
      async *chat(req) {
        yield { delta: "par", done: false };
        await new Promise<void>((resolve) => {
          if (!req.signal || req.signal.aborted) return resolve();
          req.signal.addEventListener("abort", () => resolve(), { once: true });
        });
        throw abortErr;
      },
    };
    const registry: ProviderRegistry = { get: () => adapter };
    const ledger = new UsageLedger({ now: () => 0 });
    const budgets = new BudgetService(ledger, () => 0);
    const config = resolveAiConfig(
      { enabled: true, provider: "openai", chatModel: "gpt-4o", embedModel: "text-embedding-3-small", budgets: DEFAULT_BUDGETS },
      null,
      {},
    );
    const gateway = new AiGateway({ resolveConfig: () => config, registry, budgets, ledger, now: () => 0, sleep });

    const ac = new AbortController();
    const deltas: string[] = [];
    const run = (async () => {
      for await (const chunk of gateway.chat({ ...chatParams(), signal: ac.signal })) {
        if (chunk.delta) deltas.push(chunk.delta);
      }
    })();
    await new Promise((r) => setTimeout(r, 10));
    ac.abort();
    let err: unknown;
    try {
      await run;
    } catch (e) {
      err = e;
    }
    expect(err).toBeDefined();
    const details = (err as { details?: Record<string, unknown> }).details;
    expect(details?.reason).toBe("ai_aborted");
    // the abort is never retried (no backoff sleep) — and no fallback re-route.
    expect(sleep).not.toHaveBeenCalled();
    expect(deltas).toEqual(["par"]);
    // usage was recorded as an error row (partial tokens before the abort).
    expect(ledger.records[ledger.records.length - 1]!.status).toBe("error");
  });
});