import { describe, expect, it, vi } from "vitest";
import { ModelCatalog } from "../src/model-catalog.js";
import type { ModelRegistry } from "@earendil-works/pi-coding-agent";
import type { Storage } from "../src/storage.js";
import type { Session } from "../src/types.js";

// ─── Helpers ───

type TestModel = {
  id: string;
  name: string;
  provider: string;
  contextWindow: number;
  baseUrl?: string;
};

function makeRegistry(
  available: TestModel[] = [],
  all?: TestModel[],
  oauthModelIds: string[] = [],
): ModelRegistry {
  const oauth = new Set(oauthModelIds);
  return {
    refresh: vi.fn(async () => {}),
    getAvailable: vi.fn(() => available),
    getAll: vi.fn(() => all ?? available),
    isUsingOAuth: vi.fn((model: TestModel) => oauth.has(`${model.provider}/${model.id}`)),
  } as unknown as ModelRegistry;
}

function makeStorage(sessions: Session[] = []): Storage & { saved: Session[] } {
  const saved: Session[] = [];
  return {
    saved,
    saveSession: vi.fn((s: Session) => saved.push(s)),
    listSessions: vi.fn(() => sessions),
  } as unknown as Storage & { saved: Session[] };
}

function makeSession(overrides: Partial<Session> = {}): Session {
  const now = Date.now();
  return {
    id: "s1",
    workspaceId: "w1",
    status: "ready",
    createdAt: now,
    lastActivity: now,
    messageCount: 0,
    tokens: { input: 0, output: 0 },
    cost: 0,
    ...overrides,
  };
}

const SONNET = {
  id: "claude-sonnet-4-20250514",
  name: "Claude Sonnet 4",
  provider: "anthropic",
  contextWindow: 200000,
};

const GPT = {
  id: "gpt-5.3-codex",
  name: "GPT-5.3 Codex",
  provider: "openai",
  contextWindow: 272000,
};

const GEMINI = {
  id: "gemini-3.0-pro",
  name: "Gemini 3.0 Pro",
  provider: "google",
  contextWindow: 2000000,
};

// ─── Tests ───

describe("ModelCatalog", () => {
  describe("refresh", () => {
    it("populates catalog from available models", () => {
      const registry = makeRegistry([SONNET, GPT]);
      const catalog = new ModelCatalog(registry, makeStorage());

      catalog.refresh();

      const models = catalog.getAll();
      expect(models).toHaveLength(2);
      expect(models[0].id).toBe("anthropic/claude-sonnet-4-20250514");
      expect(models[1].id).toBe("openai/gpt-5.3-codex");
    });

    it("returns an empty catalog when no providers are authenticated/configured", () => {
      const registry = makeRegistry([], [SONNET]);
      const catalog = new ModelCatalog(registry, makeStorage());

      catalog.refresh();

      expect(catalog.getAll()).toEqual([]);
    });

    it("keeps configured models while hiding unauthenticated built-ins", () => {
      const local = {
        id: "Qwen3.5-27B-8bit",
        name: "Qwen3.5 27B (Local VLM)",
        provider: "omlx",
        contextWindow: 262144,
      };
      const registry = makeRegistry([local], [SONNET, GPT, local]);
      const catalog = new ModelCatalog(registry, makeStorage());

      catalog.refresh();

      const models = catalog.getAll();
      expect(models).toHaveLength(1);
      expect(models[0].id).toBe("omlx/Qwen3.5-27B-8bit");
    });

    it("includes localhost models even when the registry marks auth as unavailable", () => {
      const local = {
        id: "Qwen3.6-27B-8bit",
        name: "Qwen3.6 27B (Local VLM 8-bit)",
        provider: "omlx",
        contextWindow: 262144,
        baseUrl: "http://localhost:8400/v1",
      };
      const registry = makeRegistry([], [SONNET, local]);
      const catalog = new ModelCatalog(registry, makeStorage());

      catalog.refresh();

      const models = catalog.getAll();
      expect(models).toHaveLength(1);
      expect(models[0].id).toBe("omlx/Qwen3.6-27B-8bit");
    });

    it("deduplicates by provider/id", () => {
      const dup = { ...SONNET };
      const registry = makeRegistry([SONNET, dup]);
      const catalog = new ModelCatalog(registry, makeStorage());

      catalog.refresh();

      expect(catalog.getAll()).toHaveLength(1);
    });

    it("sets updatedAt timestamp on success", () => {
      const catalog = new ModelCatalog(makeRegistry([SONNET]), makeStorage());
      expect(catalog.getUpdatedAt()).toBe(0);

      catalog.refresh();

      expect(catalog.getUpdatedAt()).toBeGreaterThan(0);
    });

    it("does not wait for a stalled registry refresh", async () => {
      const registry = makeRegistry([SONNET]);
      (registry.refresh as ReturnType<typeof vi.fn>).mockReturnValue(new Promise<void>(() => {}));
      const catalog = new ModelCatalog(registry, makeStorage());
      let completed = false;

      void catalog.refresh().then(() => {
        completed = true;
      });
      await Promise.resolve();

      expect(completed).toBe(true);
      expect(catalog.getAll()).toEqual([
        expect.objectContaining({ id: "anthropic/claude-sonnet-4-20250514" }),
      ]);
    });

    it("updates the cached catalog after a background refresh completes", async () => {
      let finishRefresh: (() => void) | undefined;
      const reload = new Promise<void>((resolve) => {
        finishRefresh = resolve;
      });
      const registry = makeRegistry([SONNET]);
      (registry.refresh as ReturnType<typeof vi.fn>).mockReturnValue(reload);
      const catalog = new ModelCatalog(registry, makeStorage());

      await catalog.refresh();
      expect(catalog.getAll().map((model) => model.id)).toEqual([
        "anthropic/claude-sonnet-4-20250514",
      ]);

      (registry.getAvailable as ReturnType<typeof vi.fn>).mockReturnValue([GPT]);
      finishRefresh?.();
      await reload;

      expect(catalog.getAll().map((model) => model.id)).toEqual(["openai/gpt-5.3-codex"]);
    });

    it("coalesces calls while a background refresh is in flight", async () => {
      const registry = makeRegistry([SONNET]);
      (registry.refresh as ReturnType<typeof vi.fn>).mockReturnValue(new Promise<void>(() => {}));
      const catalog = new ModelCatalog(registry, makeStorage());

      await catalog.refresh();
      await catalog.refresh();

      expect(registry.refresh).toHaveBeenCalledTimes(1);
    });

    it("survives registry throwing", () => {
      const registry = makeRegistry();
      (registry.refresh as ReturnType<typeof vi.fn>).mockImplementation(() => {
        throw new Error("disk on fire");
      });

      const catalog = new ModelCatalog(registry, makeStorage());
      expect(() => catalog.refresh()).not.toThrow();
      expect(catalog.getAll()).toEqual([]);
    });

    it("defaults contextWindow to 200000 when SDK returns 0", () => {
      const noCtx = { ...SONNET, contextWindow: 0 };
      const catalog = new ModelCatalog(makeRegistry([noCtx]), makeStorage());

      catalog.refresh();

      expect(catalog.getAll()[0].contextWindow).toBe(200000);
    });

    it("uses model id as name when SDK name is empty", () => {
      const noName = { ...SONNET, name: "" };
      const catalog = new ModelCatalog(makeRegistry([noName]), makeStorage());

      catalog.refresh();

      expect(catalog.getAll()[0].name).toBe("claude-sonnet-4-20250514");
    });

    it("marks OAuth-backed models as subscription auth", () => {
      const catalog = new ModelCatalog(
        makeRegistry([SONNET, GPT], undefined, ["anthropic/claude-sonnet-4-20250514"]),
        makeStorage(),
      );

      catalog.refresh();

      expect(catalog.getAll()).toEqual(
        expect.arrayContaining([
          expect.objectContaining({
            id: "anthropic/claude-sonnet-4-20250514",
            authKind: "subscription",
          }),
          expect.objectContaining({ id: "openai/gpt-5.3-codex", authKind: "apiKey" }),
        ]),
      );
    });
  });

  describe("getContextWindow", () => {
    function catalogWith(...models: (typeof SONNET)[]) {
      const catalog = new ModelCatalog(makeRegistry(models), makeStorage());
      catalog.refresh();
      return catalog;
    }

    it("matches by exact canonical id", () => {
      const c = catalogWith(SONNET);
      expect(c.getContextWindow("anthropic/claude-sonnet-4-20250514")).toBe(200000);
    });

    it("matches by tail segment (no provider prefix)", () => {
      const c = catalogWith(GPT);
      expect(c.getContextWindow("gpt-5.3-codex")).toBe(272000);
    });

    it("matches by model name", () => {
      const c = catalogWith(GPT);
      expect(c.getContextWindow("GPT-5.3 Codex")).toBe(272000);
    });

    it("matches by normalized token (case-insensitive, punctuation-stripped)", () => {
      const c = catalogWith(GPT);
      expect(c.getContextWindow("GPT 5.3 Codex")).toBe(272000);
    });

    it("matches by tail suffix (provider/tail → tail)", () => {
      const c = catalogWith(SONNET);
      expect(c.getContextWindow("claude-sonnet-4-20250514")).toBe(200000);
    });

    it("falls back to NNNk suffix parsing", () => {
      const c = catalogWith(SONNET);
      expect(c.getContextWindow("unknown-model-128k")).toBe(128000);
      expect(c.getContextWindow("some-272k-model")).toBe(272000);
    });

    it("returns 200000 as ultimate fallback", () => {
      const c = catalogWith(SONNET);
      expect(c.getContextWindow("totally-unknown")).toBe(200000);
    });

    it("handles empty string", () => {
      const c = catalogWith(SONNET);
      expect(c.getContextWindow("")).toBe(200000);
    });

    it("trims whitespace", () => {
      const c = catalogWith(GPT);
      expect(c.getContextWindow("  gpt-5.3-codex  ")).toBe(272000);
    });

    it("matches large context windows correctly", () => {
      const c = catalogWith(GEMINI);
      expect(c.getContextWindow("gemini-3.0-pro")).toBe(2000000);
    });
  });

  describe("ensureSessionContextWindow", () => {
    it("sets context window when missing", () => {
      const storage = makeStorage();
      const catalog = new ModelCatalog(makeRegistry([GPT]), storage);
      catalog.refresh();

      const session = makeSession({ model: "gpt-5.3-codex", contextWindow: undefined });
      catalog.ensureSessionContextWindow(session);

      expect(session.contextWindow).toBe(272000);
      expect(storage.saveSession).toHaveBeenCalledWith(session);
    });

    it("sets context window when zero", () => {
      const storage = makeStorage();
      const catalog = new ModelCatalog(makeRegistry([GPT]), storage);
      catalog.refresh();

      const session = makeSession({ model: "gpt-5.3-codex", contextWindow: 0 });
      catalog.ensureSessionContextWindow(session);

      expect(session.contextWindow).toBe(272000);
    });

    it("heals stale 200k fallback to correct value", () => {
      const storage = makeStorage();
      const catalog = new ModelCatalog(makeRegistry([GPT]), storage);
      catalog.refresh();

      const session = makeSession({ model: "gpt-5.3-codex", contextWindow: 200000 });
      catalog.ensureSessionContextWindow(session);

      expect(session.contextWindow).toBe(272000);
      expect(storage.saveSession).toHaveBeenCalled();
    });

    it("does not overwrite a correct non-fallback value", () => {
      const storage = makeStorage();
      const catalog = new ModelCatalog(makeRegistry([GPT]), storage);
      catalog.refresh();

      const session = makeSession({ model: "gpt-5.3-codex", contextWindow: 272000 });
      catalog.ensureSessionContextWindow(session);

      expect(session.contextWindow).toBe(272000);
      expect(storage.saveSession).not.toHaveBeenCalled();
    });

    it("does not overwrite a custom value even if different from resolved", () => {
      const storage = makeStorage();
      const catalog = new ModelCatalog(makeRegistry([GPT]), storage);
      catalog.refresh();

      // User set 100000 manually — should be preserved (not 200k fallback)
      const session = makeSession({ model: "gpt-5.3-codex", contextWindow: 100000 });
      catalog.ensureSessionContextWindow(session);

      expect(session.contextWindow).toBe(100000);
      expect(storage.saveSession).not.toHaveBeenCalled();
    });

    it("returns the session for chaining", () => {
      const catalog = new ModelCatalog(makeRegistry([GPT]), makeStorage());
      catalog.refresh();

      const session = makeSession({ model: "gpt-5.3-codex" });
      const result = catalog.ensureSessionContextWindow(session);

      expect(result).toBe(session);
    });
  });

  describe("allowlist filtering", () => {
    it("only includes models whose canonical ID is in the allowlist", () => {
      const registry = makeRegistry([SONNET, GPT, GEMINI]);
      const catalog = new ModelCatalog(registry, makeStorage(), [
        "anthropic/claude-sonnet-4-20250514",
        "google/gemini-3.0-pro",
      ]);

      catalog.refresh();

      const models = catalog.getAll();
      expect(models).toHaveLength(2);
      expect(models.map((m) => m.id)).toEqual([
        "anthropic/claude-sonnet-4-20250514",
        "google/gemini-3.0-pro",
      ]);
    });

    it("excludes models not in the allowlist", () => {
      const registry = makeRegistry([SONNET, GPT, GEMINI]);
      const catalog = new ModelCatalog(registry, makeStorage(), ["openai/gpt-5.3-codex"]);

      catalog.refresh();

      const models = catalog.getAll();
      expect(models).toHaveLength(1);
      expect(models[0].id).toBe("openai/gpt-5.3-codex");
    });

    it("does not filter when allowlist is undefined", () => {
      const registry = makeRegistry([SONNET, GPT, GEMINI]);
      const catalog = new ModelCatalog(registry, makeStorage(), undefined);

      catalog.refresh();

      expect(catalog.getAll()).toHaveLength(3);
    });

    it("does not filter when allowlist is empty array", () => {
      const registry = makeRegistry([SONNET, GPT, GEMINI]);
      const catalog = new ModelCatalog(registry, makeStorage(), []);

      catalog.refresh();

      expect(catalog.getAll()).toHaveLength(3);
    });

    it("supports Pi enabledModels glob patterns against canonical IDs", () => {
      const registry = makeRegistry([SONNET, GPT, GEMINI]);
      const catalog = new ModelCatalog(registry, makeStorage(), ["openai/*", "google/gemini-*"]);

      catalog.refresh();

      expect(catalog.getAll().map((m) => m.id)).toEqual([
        "openai/gpt-5.3-codex",
        "google/gemini-3.0-pro",
      ]);
    });

    it("supports Pi enabledModels glob patterns against model IDs", () => {
      const registry = makeRegistry([SONNET, GPT, GEMINI]);
      const catalog = new ModelCatalog(registry, makeStorage(), ["*sonnet*", "gpt-*"]);

      catalog.refresh();

      expect(catalog.getAll().map((m) => m.id)).toEqual([
        "anthropic/claude-sonnet-4-20250514",
        "openai/gpt-5.3-codex",
      ]);
    });

    it("ignores optional Pi thinking suffixes in allowlist patterns", () => {
      const registry = makeRegistry([SONNET, GPT, GEMINI]);
      const catalog = new ModelCatalog(registry, makeStorage(), ["openai/*:high"]);

      catalog.refresh();

      expect(catalog.getAll().map((m) => m.id)).toEqual(["openai/gpt-5.3-codex"]);
    });

    it("re-reads dynamic allowlist providers on refresh", () => {
      const registry = makeRegistry([SONNET, GPT, GEMINI]);
      let allowlist = ["openai/*"];
      const catalog = new ModelCatalog(registry, makeStorage(), () => allowlist);

      catalog.refresh();
      expect(catalog.getAll().map((m) => m.id)).toEqual(["openai/gpt-5.3-codex"]);

      allowlist = ["google/*"];
      catalog.refresh();
      expect(catalog.getAll().map((m) => m.id)).toEqual(["google/gemini-3.0-pro"]);
    });

    it("returns empty when allowlist is set but no model is authenticated/configured", () => {
      const registry = makeRegistry([], [SONNET, GPT, GEMINI]);
      const catalog = new ModelCatalog(registry, makeStorage(), ["openai/gpt-5.3-codex"]);

      catalog.refresh();

      expect(catalog.getAll()).toEqual([]);
    });

    it("returns empty catalog when no models match the allowlist", () => {
      const registry = makeRegistry([SONNET, GPT]);
      const catalog = new ModelCatalog(registry, makeStorage(), ["google/gemini-3.0-pro"]);

      catalog.refresh();

      expect(catalog.getAll()).toEqual([]);
    });

    it("getContextWindow falls back to 200k for model excluded by allowlist", () => {
      const registry = makeRegistry([SONNET, GPT]);
      const catalog = new ModelCatalog(registry, makeStorage(), [
        "anthropic/claude-sonnet-4-20250514",
      ]);
      catalog.refresh();

      // GPT is excluded from catalog — can't resolve its context window
      expect(catalog.getContextWindow("gpt-5.3-codex")).toBe(200000);
      // SONNET is in the catalog — resolves normally
      expect(catalog.getContextWindow("claude-sonnet-4-20250514")).toBe(200000);
    });

    it("getContextWindow still uses NNNk fallback for excluded models", () => {
      const registry = makeRegistry([GPT]);
      const catalog = new ModelCatalog(registry, makeStorage(), [
        "anthropic/claude-sonnet-4-20250514",
      ]);
      catalog.refresh();

      // GPT excluded, but "gpt-5.3-codex" has no NNNk suffix → 200k
      expect(catalog.getContextWindow("gpt-5.3-codex")).toBe(200000);
      // An unknown model with NNNk suffix still parses
      expect(catalog.getContextWindow("some-model-128k")).toBe(128000);
    });

    it("ensureSessionContextWindow uses 200k default for model excluded by allowlist", () => {
      const storage = makeStorage();
      const registry = makeRegistry([GPT]);
      const catalog = new ModelCatalog(registry, storage, ["anthropic/claude-sonnet-4-20250514"]);
      catalog.refresh();

      const session = makeSession({ model: "gpt-5.3-codex", contextWindow: undefined });
      catalog.ensureSessionContextWindow(session);

      // GPT excluded from catalog → getContextWindow returns 200k fallback
      expect(session.contextWindow).toBe(200000);
      expect(storage.saveSession).toHaveBeenCalledWith(session);
    });

    it("healPersistedSessionContextWindows with allowlisted vs excluded models", () => {
      const s1 = makeSession({ id: "s1", model: "gpt-5.3-codex", contextWindow: undefined });
      const s2 = makeSession({
        id: "s2",
        model: "claude-sonnet-4-20250514",
        contextWindow: undefined,
      });

      const storage = makeStorage([s1, s2]);
      const registry = makeRegistry([SONNET, GPT]);
      const catalog = new ModelCatalog(registry, storage, ["anthropic/claude-sonnet-4-20250514"]);
      catalog.refresh();

      catalog.healPersistedSessionContextWindows();

      // SONNET is allowlisted → resolves to 200k (its actual value)
      expect(s2.contextWindow).toBe(200000);
      // GPT is excluded → falls back to 200k default (not 272k)
      expect(s1.contextWindow).toBe(200000);
      expect(storage.saveSession).toHaveBeenCalledTimes(2);
    });
  });

  describe("healPersistedSessionContextWindows", () => {
    it("heals sessions with stale fallback context windows", () => {
      const s1 = makeSession({ id: "s1", model: "gpt-5.3-codex", contextWindow: 200000 });
      const s2 = makeSession({ id: "s2", model: "gpt-5.3-codex", contextWindow: 272000 });
      const s3 = makeSession({ id: "s3", model: "gemini-3.0-pro", contextWindow: undefined });

      const storage = makeStorage([s1, s2, s3]);
      const catalog = new ModelCatalog(makeRegistry([GPT, GEMINI]), storage);
      catalog.refresh();

      catalog.healPersistedSessionContextWindows();

      // s1 healed from 200k → 272k, s3 healed from undefined → 2M
      expect(s1.contextWindow).toBe(272000);
      expect(s3.contextWindow).toBe(2000000);
      // s2 was already correct — not saved
      expect(storage.saveSession).toHaveBeenCalledTimes(2);
    });

    it("does nothing when all sessions are correct", () => {
      const s1 = makeSession({ id: "s1", model: "gpt-5.3-codex", contextWindow: 272000 });
      const storage = makeStorage([s1]);
      const catalog = new ModelCatalog(makeRegistry([GPT]), storage);
      catalog.refresh();

      catalog.healPersistedSessionContextWindows();

      expect(storage.saveSession).not.toHaveBeenCalled();
    });
  });
});