import { describe, it, expect, beforeEach } from "vitest";
import { ModelPricingRegistry } from "./registry.js";
import { defaultModelPrices } from "./defaultPrices.js";
import type { LlmModelWithPricing } from "./types.js";

// Convert POSIX-style (?i) inline flag to JS RegExp 'i' flag
function compilePattern(pattern: string): RegExp {
  if (pattern.startsWith("(?i)")) {
    return new RegExp(pattern.slice(4), "i");
  }
  return new RegExp(pattern);
}

// Create a mock registry that we can load with test data without Prisma
class TestableRegistry extends ModelPricingRegistry {
  loadPatterns(models: LlmModelWithPricing[]) {
    // Access private fields via any cast for testing
    const self = this as any;
    self._patterns = models.map((model) => ({
      regex: compilePattern(model.matchPattern),
      model,
    }));
    self._exactMatchCache = new Map();
    self._loaded = true;
  }
}

const gpt4o: LlmModelWithPricing = {
  id: "model-gpt4o",
  friendlyId: "llm_model_gpt4o",
  modelName: "gpt-4o",
  matchPattern: "^gpt-4o(-\\d{4}-\\d{2}-\\d{2})?$",
  startDate: null,
  pricingTiers: [
    {
      id: "tier-gpt4o-standard",
      name: "Standard",
      isDefault: true,
      priority: 0,
      conditions: [],
      prices: [
        { usageType: "input", price: 0.0000025 },
        { usageType: "output", price: 0.00001 },
        { usageType: "input_cached_tokens", price: 0.00000125 },
      ],
    },
  ],
};

const claudeSonnet: LlmModelWithPricing = {
  id: "model-claude-sonnet",
  friendlyId: "llm_model_claude_sonnet",
  modelName: "claude-sonnet-4-0",
  matchPattern: "^claude-sonnet-4-0(-\\d{8})?$",
  startDate: null,
  pricingTiers: [
    {
      id: "tier-claude-sonnet-standard",
      name: "Standard",
      isDefault: true,
      priority: 0,
      conditions: [],
      prices: [
        { usageType: "input", price: 0.000003 },
        { usageType: "output", price: 0.000015 },
        { usageType: "input_cached_tokens", price: 0.0000015 },
      ],
    },
  ],
};

// Prices cache reads under the Anthropic-style alias `cache_read_input_tokens` (not
// `input_cached_tokens`) plus a cache-creation price, to exercise alias resolution.
const claudeWithCache: LlmModelWithPricing = {
  id: "model-claude-with-cache",
  friendlyId: "llm_model_claude_with_cache",
  modelName: "claude-with-cache",
  matchPattern: "^claude-with-cache$",
  startDate: null,
  pricingTiers: [
    {
      id: "tier-claude-with-cache",
      name: "Standard",
      isDefault: true,
      priority: 0,
      conditions: [],
      prices: [
        { usageType: "input", price: 0.000003 },
        { usageType: "output", price: 0.000015 },
        { usageType: "cache_read_input_tokens", price: 0.0000003 },
        { usageType: "cache_creation_input_tokens", price: 0.00000375 },
      ],
    },
  ],
};

// No cache prices at all — cached tokens should fall back to the input price.
const noCachePrice: LlmModelWithPricing = {
  id: "model-no-cache-price",
  friendlyId: "llm_model_no_cache_price",
  modelName: "no-cache-price",
  matchPattern: "^no-cache-price$",
  startDate: null,
  pricingTiers: [
    {
      id: "tier-no-cache-price",
      name: "Standard",
      isDefault: true,
      priority: 0,
      conditions: [],
      prices: [
        { usageType: "input", price: 0.000003 },
        { usageType: "output", price: 0.000015 },
      ],
    },
  ],
};

describe("ModelPricingRegistry", () => {
  let registry: TestableRegistry;

  beforeEach(() => {
    registry = new TestableRegistry(null as any);
    registry.loadPatterns([gpt4o, claudeSonnet, claudeWithCache, noCachePrice]);
  });

  describe("match", () => {
    it("should match exact model name", () => {
      const result = registry.match("gpt-4o");
      expect(result).not.toBeNull();
      expect(result!.modelName).toBe("gpt-4o");
    });

    it("should match model with date suffix", () => {
      const result = registry.match("gpt-4o-2024-08-06");
      expect(result).not.toBeNull();
      expect(result!.modelName).toBe("gpt-4o");
    });

    it("should match claude model", () => {
      const result = registry.match("claude-sonnet-4-0-20250514");
      expect(result).not.toBeNull();
      expect(result!.modelName).toBe("claude-sonnet-4-0");
    });

    it("should return null for unknown model", () => {
      const result = registry.match("unknown-model-xyz");
      expect(result).toBeNull();
    });

    it("should cache exact matches", () => {
      registry.match("gpt-4o");
      registry.match("gpt-4o");
      // Second call should use cache - no way to verify without mocking, but it shouldn't error
      expect(registry.match("gpt-4o")!.modelName).toBe("gpt-4o");
    });

    it("should cache misses", () => {
      expect(registry.match("unknown")).toBeNull();
      expect(registry.match("unknown")).toBeNull();
    });
  });

  describe("calculateCost", () => {
    it("should calculate cost for input and output tokens", () => {
      const result = registry.calculateCost("gpt-4o", {
        input: 1000,
        output: 100,
      });

      expect(result).not.toBeNull();
      expect(result!.matchedModelName).toBe("gpt-4o");
      expect(result!.pricingTierName).toBe("Standard");
      expect(result!.inputCost).toBeCloseTo(0.0025); // 1000 * 0.0000025
      expect(result!.outputCost).toBeCloseTo(0.001); // 100 * 0.00001
      expect(result!.totalCost).toBeCloseTo(0.0035);
    });

    it("should include cached token costs and charge input only on the fresh portion", () => {
      // input_tokens (500) is inclusive of the 200 cached read tokens, so the input price
      // applies to the 300 fresh tokens and the cache price to the 200 cached tokens — the
      // cached tokens must not be billed twice.
      const result = registry.calculateCost("gpt-4o", {
        input: 500,
        output: 50,
        input_cached_tokens: 200,
      });

      expect(result).not.toBeNull();
      expect(result!.costDetails["input"]).toBeCloseTo(0.00075); // (500 - 200) * 0.0000025
      expect(result!.costDetails["output"]).toBeCloseTo(0.0005); // 50 * 0.00001
      expect(result!.costDetails["input_cached_tokens"]).toBeCloseTo(0.00025); // 200 * 0.00000125
      expect(result!.totalCost).toBeCloseTo(0.0015);
    });

    it("should not double-charge cache creation tokens (subset of input)", () => {
      // input (1000) is inclusive of both the 400 cache-read and 300 cache-creation tokens.
      const result = registry.calculateCost("claude-with-cache", {
        input: 1000,
        output: 100,
        input_cached_tokens: 400,
        cache_creation_input_tokens: 300,
      });

      expect(result).not.toBeNull();
      // fresh input = 1000 - 400 - 300 = 300
      expect(result!.costDetails["input"]).toBeCloseTo(0.0009); // 300 * 0.000003
      expect(result!.costDetails["input_cached_tokens"]).toBeCloseTo(0.00012); // 400 * 0.0000003
      expect(result!.costDetails["cache_creation_input_tokens"]).toBeCloseTo(0.001125); // 300 * 0.00000375
      expect(result!.costDetails["output"]).toBeCloseTo(0.0015); // 100 * 0.000015
      // 0.0009 + 0.00012 + 0.001125 + 0.0015
      expect(result!.totalCost).toBeCloseTo(0.003645);
    });

    it("should apply the cache-read discount when priced under a provider alias key", () => {
      // The usage is normalized to `input_cached_tokens` but this model prices cache reads
      // under `cache_read_input_tokens` — the discount must still apply.
      const result = registry.calculateCost("claude-with-cache", {
        input: 1000,
        input_cached_tokens: 400,
      });

      expect(result).not.toBeNull();
      expect(result!.costDetails["input"]).toBeCloseTo(0.0018); // (1000 - 400) * 0.000003
      expect(result!.costDetails["input_cached_tokens"]).toBeCloseTo(0.00012); // 400 * 0.0000003
      expect(result!.totalCost).toBeCloseTo(0.00192);
    });

    it("should fall back to the input price for cache tokens when no cache price exists", () => {
      // no-cache-price model has only input/output prices; cached tokens must still be billed
      // (at the input price) — never free, never double-charged. Total equals input * price.
      const result = registry.calculateCost("no-cache-price", {
        input: 1000,
        input_cached_tokens: 400,
      });

      expect(result).not.toBeNull();
      expect(result!.costDetails["input"]).toBeCloseTo(0.0018); // (1000 - 400) * 0.000003
      expect(result!.costDetails["input_cached_tokens"]).toBeCloseTo(0.0012); // 400 * 0.000003
      expect(result!.totalCost).toBeCloseTo(0.003); // 1000 * 0.000003 — unchanged from no-cache behavior
    });

    it("should return null for unknown model", () => {
      const result = registry.calculateCost("unknown-model", { input: 100, output: 50 });
      expect(result).toBeNull();
    });

    it("should handle zero tokens", () => {
      const result = registry.calculateCost("gpt-4o", { input: 0, output: 0 });
      expect(result).not.toBeNull();
      expect(result!.totalCost).toBe(0);
    });

    it("should handle missing usage types gracefully", () => {
      const result = registry.calculateCost("gpt-4o", { input: 100 });
      expect(result).not.toBeNull();
      expect(result!.inputCost).toBeCloseTo(0.00025);
      expect(result!.outputCost).toBe(0); // No output tokens
      expect(result!.totalCost).toBeCloseTo(0.00025);
    });
  });

  describe("isLoaded", () => {
    it("should return false before loading", () => {
      const freshRegistry = new TestableRegistry(null as any);
      expect(freshRegistry.isLoaded).toBe(false);
    });

    it("should return true after loading", () => {
      expect(registry.isLoaded).toBe(true);
    });
  });

  describe("prefix stripping", () => {
    it("should match gateway-prefixed model names", () => {
      const result = registry.match("openai/gpt-4o");
      expect(result).not.toBeNull();
      expect(result!.modelName).toBe("gpt-4o");
    });

    it("should match openrouter-prefixed model names with date suffix", () => {
      const result = registry.match("openai/gpt-4o-2024-08-06");
      expect(result).not.toBeNull();
      expect(result!.modelName).toBe("gpt-4o");
    });

    it("should return null for prefixed unknown model", () => {
      const result = registry.match("xai/unknown-model");
      expect(result).toBeNull();
    });
  });

  describe("tier matching", () => {
    const multiTierModel: LlmModelWithPricing = {
      id: "model-gemini-pro",
      friendlyId: "llm_model_gemini_pro",
      modelName: "gemini-2.5-pro",
      matchPattern: "^gemini-2\\.5-pro$",
      startDate: null,
      pricingTiers: [
        {
          id: "tier-large-context",
          name: "Large Context",
          isDefault: false,
          priority: 0,
          conditions: [{ usageDetailPattern: "input", operator: "gt" as const, value: 200000 }],
          prices: [
            { usageType: "input", price: 0.0000025 },
            { usageType: "output", price: 0.00001 },
          ],
        },
        {
          id: "tier-standard",
          name: "Standard",
          isDefault: true,
          priority: 1,
          conditions: [],
          prices: [
            { usageType: "input", price: 0.00000125 },
            { usageType: "output", price: 0.000005 },
          ],
        },
      ],
    };

    it("should use conditional tier when conditions match", () => {
      const tieredRegistry = new TestableRegistry(null as any);
      tieredRegistry.loadPatterns([multiTierModel]);

      const result = tieredRegistry.calculateCost("gemini-2.5-pro", {
        input: 250000,
        output: 1000,
      });

      expect(result).not.toBeNull();
      expect(result!.pricingTierName).toBe("Large Context");
      expect(result!.inputCost).toBeCloseTo(0.625); // 250000 * 0.0000025
    });

    it("should fall back to default tier when conditions do not match", () => {
      const tieredRegistry = new TestableRegistry(null as any);
      tieredRegistry.loadPatterns([multiTierModel]);

      const result = tieredRegistry.calculateCost("gemini-2.5-pro", {
        input: 1000,
        output: 100,
      });

      expect(result).not.toBeNull();
      expect(result!.pricingTierName).toBe("Standard");
      expect(result!.inputCost).toBeCloseTo(0.00125); // 1000 * 0.00000125
    });

    it("should not let unconditional tier win over conditional match", () => {
      // Model where unconditional tier has lower priority than conditional
      const model: LlmModelWithPricing = {
        ...multiTierModel,
        pricingTiers: [
          {
            id: "tier-unconditional",
            name: "Unconditional",
            isDefault: false,
            priority: 0,
            conditions: [],
            prices: [{ usageType: "input", price: 0.001 }],
          },
          {
            id: "tier-conditional",
            name: "Conditional",
            isDefault: false,
            priority: 1,
            conditions: [{ usageDetailPattern: "input", operator: "gt" as const, value: 100 }],
            prices: [{ usageType: "input", price: 0.0001 }],
          },
          {
            id: "tier-default",
            name: "Default",
            isDefault: true,
            priority: 2,
            conditions: [],
            prices: [{ usageType: "input", price: 0.01 }],
          },
        ],
      };

      const tieredRegistry = new TestableRegistry(null as any);
      tieredRegistry.loadPatterns([model]);

      // Condition matches — conditional tier should win, not the unconditional one
      const result = tieredRegistry.calculateCost("gemini-2.5-pro", { input: 500 });
      expect(result).not.toBeNull();
      expect(result!.pricingTierName).toBe("Conditional");
    });

    it("should fall back to isDefault tier when no conditions match", () => {
      const model: LlmModelWithPricing = {
        ...multiTierModel,
        pricingTiers: [
          {
            id: "tier-conditional",
            name: "Conditional",
            isDefault: false,
            priority: 0,
            conditions: [{ usageDetailPattern: "input", operator: "gt" as const, value: 999999 }],
            prices: [{ usageType: "input", price: 0.001 }],
          },
          {
            id: "tier-default",
            name: "Default",
            isDefault: true,
            priority: 1,
            conditions: [],
            prices: [{ usageType: "input", price: 0.0001 }],
          },
        ],
      };

      const tieredRegistry = new TestableRegistry(null as any);
      tieredRegistry.loadPatterns([model]);

      const result = tieredRegistry.calculateCost("gemini-2.5-pro", { input: 100 });
      expect(result).not.toBeNull();
      expect(result!.pricingTierName).toBe("Default");
    });
  });

  describe("defaultModelPrices (Langfuse JSON)", () => {
    it("should load all models from the JSON file", () => {
      expect(defaultModelPrices.length).toBeGreaterThan(100);
    });

    it("should compile all match patterns without errors", () => {
      const langfuseRegistry = new TestableRegistry(null as any);
      const models: LlmModelWithPricing[] = defaultModelPrices.map((def, i) => ({
        id: `test-${i}`,
        friendlyId: `llm_model_test${i}`,
        modelName: def.modelName,
        matchPattern: def.matchPattern,
        startDate: def.startDate ? new Date(def.startDate) : null,
        pricingTiers: def.pricingTiers.map((tier, j) => ({
          id: `tier-${i}-${j}`,
          name: tier.name,
          isDefault: tier.isDefault,
          priority: tier.priority,
          conditions: tier.conditions,
          prices: Object.entries(tier.prices).map(([usageType, price]) => ({
            usageType,
            price,
          })),
        })),
      }));

      // This should not throw — all 141 patterns should compile
      expect(() => langfuseRegistry.loadPatterns(models)).not.toThrow();
      expect(langfuseRegistry.isLoaded).toBe(true);
    });

    it("should match real-world model names from Langfuse patterns", () => {
      const langfuseRegistry = new TestableRegistry(null as any);
      const models: LlmModelWithPricing[] = defaultModelPrices.map((def, i) => ({
        id: `test-${i}`,
        friendlyId: `llm_model_test${i}`,
        modelName: def.modelName,
        matchPattern: def.matchPattern,
        startDate: null,
        pricingTiers: def.pricingTiers.map((tier, j) => ({
          id: `tier-${i}-${j}`,
          name: tier.name,
          isDefault: tier.isDefault,
          priority: tier.priority,
          conditions: tier.conditions,
          prices: Object.entries(tier.prices).map(([usageType, price]) => ({
            usageType,
            price,
          })),
        })),
      }));
      langfuseRegistry.loadPatterns(models);

      // Test real model strings that SDKs send
      expect(langfuseRegistry.match("gpt-4o")).not.toBeNull();
      expect(langfuseRegistry.match("gpt-4o-mini")).not.toBeNull();
      expect(langfuseRegistry.match("claude-sonnet-4-5-20250929")).not.toBeNull();
      expect(langfuseRegistry.match("claude-sonnet-4-20250514")).not.toBeNull();
    });
  });
});

describe("loadFromModels / toSerializable (worker in-memory load)", () => {
  it("matches and prices from in-memory models without a DB", () => {
    const reg = new ModelPricingRegistry();
    reg.loadFromModels([gpt4o, claudeSonnet]);

    expect(reg.isLoaded).toBe(true);
    expect(reg.match("gpt-4o")).not.toBeNull();
    expect(reg.match("gpt-4o-2024-08-06")).not.toBeNull();

    const cost = reg.calculateCost("gpt-4o", { input: 1000, output: 500 });
    expect(cost).not.toBeNull();
    expect(cost!.totalCost).toBeGreaterThan(0);
  });

  it("round-trips through toSerializable (the main->worker broadcast shape)", () => {
    const source = new ModelPricingRegistry();
    source.loadFromModels([gpt4o, claudeSonnet]);

    const worker = new ModelPricingRegistry();
    worker.loadFromModels(source.toSerializable());

    for (const m of ["gpt-4o", "gpt-4o-2024-08-06", "claude-sonnet-4-0", "unknown-model"]) {
      expect(worker.match(m)).toEqual(source.match(m));
    }
    const usage = { input: 1234, output: 567 };
    expect(worker.calculateCost("gpt-4o", usage)).toEqual(source.calculateCost("gpt-4o", usage));
  });

  it("throws if loadFromDatabase is called without a prisma client", async () => {
    const reg = new ModelPricingRegistry();
    await expect(reg.loadFromDatabase()).rejects.toThrow(/requires a prisma client/);
  });
});
