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

// Unmock usage middleware — we want to test the real implementation
vi.unmock("../middleware/usage.js");

// DB and keys are still mocked by setup.ts
import { queryWithRetry } from "../services/db.js";
import { isProKey } from "../services/keys.js";

describe("usage middleware", () => {
  let usage: typeof import("../middleware/usage.js");

  const mockJson = vi.fn();
  const mockStatus = vi.fn(() => ({ json: mockJson }));
  const mockNext = vi.fn();

  function makeReq(key = "df_free_testkey123"): any {
    return { apiKeyInfo: { key } };
  }

  function makeRes(): any {
    mockJson.mockClear();
    mockStatus.mockClear();
    return { status: mockStatus, json: mockJson };
  }

  beforeEach(async () => {
    vi.clearAllMocks();
    vi.resetModules();
    // Re-import to get fresh internal state (new Map)
    usage = await import("../middleware/usage.js");
  });

  // --- loadUsageData ---

  describe("loadUsageData", () => {
    it("populates in-memory map from DB rows", async () => {
      vi.mocked(queryWithRetry).mockResolvedValueOnce({
        rows: [
          { key: "df_free_abc12345", count: 42, month_key: "2026-02" },
        ],
        rowCount: 1,
      } as any);

      await usage.loadUsageData();

      const stats = usage.getUsageStats("df_free_abc12345");
      expect(stats["df_free_..."]).toEqual({ count: 42, month: "2026-02" });
    });

    it("handles empty DB result gracefully", async () => {
      vi.mocked(queryWithRetry).mockResolvedValueOnce({ rows: [], rowCount: 0 } as any);
      await usage.loadUsageData();
      const stats = usage.getUsageStats("df_free_nonexist");
      expect(stats).toEqual({});
    });

    it("handles DB error gracefully (starts fresh)", async () => {
      vi.mocked(queryWithRetry).mockRejectedValueOnce(new Error("DB down"));
      await usage.loadUsageData();
      const stats = usage.getUsageStats("anything");
      expect(stats).toEqual({});
    });
  });

  // --- getUsageStats ---

  describe("getUsageStats", () => {
    it("returns empty object when key not found", () => {
      expect(usage.getUsageStats("nonexistent")).toEqual({});
    });

    it("returns empty object when no key provided", () => {
      expect(usage.getUsageStats()).toEqual({});
    });

    it("masks key to first 8 chars + '...'", async () => {
      vi.mocked(queryWithRetry).mockResolvedValueOnce({
        rows: [{ key: "df_free_longkeyvalue", count: 5, month_key: "2026-02" }],
        rowCount: 1,
      } as any);
      await usage.loadUsageData();

      const stats = usage.getUsageStats("df_free_longkeyvalue");
      const keys = Object.keys(stats);
      expect(keys).toHaveLength(1);
      expect(keys[0]).toBe("df_free_...");
    });
  });

  // --- usageMiddleware ---

  describe("usageMiddleware", () => {
    it("calls next() for free key under limit", () => {
      vi.mocked(isProKey).mockReturnValue(false);
      const req = makeReq();
      const res = makeRes();

      usage.usageMiddleware(req, res, mockNext);

      expect(mockNext).toHaveBeenCalled();
      expect(mockStatus).not.toHaveBeenCalled();
    });

    it("calls next() for pro key under limit", () => {
      vi.mocked(isProKey).mockReturnValue(true);
      const req = makeReq("df_pro_testkey123");
      const res = makeRes();

      usage.usageMiddleware(req, res, mockNext);

      expect(mockNext).toHaveBeenCalled();
      expect(mockStatus).not.toHaveBeenCalled();
    });

    it("returns 429 when free key exceeds 100/month", async () => {
      vi.mocked(isProKey).mockReturnValue(false);
      const key = "df_free_limited1234";

      // Seed usage at the limit
      const monthKey = `${new Date().getFullYear()}-${String(new Date().getMonth() + 1).padStart(2, "0")}`;
      vi.mocked(queryWithRetry).mockResolvedValueOnce({
        rows: [{ key, count: 100, month_key: monthKey }],
        rowCount: 1,
      } as any);
      await usage.loadUsageData();

      const req = makeReq(key);
      const res = makeRes();

      usage.usageMiddleware(req, res, mockNext);

      expect(mockStatus).toHaveBeenCalledWith(429);
      expect(mockJson).toHaveBeenCalledWith(
        expect.objectContaining({ error: expect.stringContaining("Free tier limit") })
      );
      expect(mockNext).not.toHaveBeenCalled();
    });

    it("returns 429 when pro key exceeds 5000/month", async () => {
      vi.mocked(isProKey).mockReturnValue(true);
      const key = "df_pro_limited12345";

      const monthKey = `${new Date().getFullYear()}-${String(new Date().getMonth() + 1).padStart(2, "0")}`;
      vi.mocked(queryWithRetry).mockResolvedValueOnce({
        rows: [{ key, count: 5000, month_key: monthKey }],
        rowCount: 1,
      } as any);
      await usage.loadUsageData();

      const req = makeReq(key);
      const res = makeRes();

      usage.usageMiddleware(req, res, mockNext);

      expect(mockStatus).toHaveBeenCalledWith(429);
      expect(mockJson).toHaveBeenCalledWith(
        expect.objectContaining({ error: expect.stringContaining("Pro tier limit") })
      );
      expect(mockNext).not.toHaveBeenCalled();
    });

    it("increments usage count on each call", () => {
      vi.mocked(isProKey).mockReturnValue(false);
      const key = "df_free_counting123";

      // Call middleware 3 times
      for (let i = 0; i < 3; i++) {
        usage.usageMiddleware(makeReq(key), makeRes(), mockNext);
      }

      const stats = usage.getUsageStats(key);
      const masked = Object.keys(stats)[0];
      expect(stats[masked].count).toBe(3);
    });

    it("resets count when month changes", async () => {
      vi.mocked(isProKey).mockReturnValue(false);
      const key = "df_free_monthreset";

      // Seed with old month data at limit
      vi.mocked(queryWithRetry).mockResolvedValueOnce({
        rows: [{ key, count: 100, month_key: "2025-01" }],
        rowCount: 1,
      } as any);
      await usage.loadUsageData();

      // Current month is different, so it should reset and allow
      const req = makeReq(key);
      const res = makeRes();
      usage.usageMiddleware(req, res, mockNext);

      expect(mockNext).toHaveBeenCalled();
      expect(mockStatus).not.toHaveBeenCalled();

      // Count should be 1 (reset + this request)
      const stats = usage.getUsageStats(key);
      const masked = Object.keys(stats)[0];
      expect(stats[masked].count).toBe(1);
    });

    it("allows free key at count 99 (just under limit)", async () => {
      vi.mocked(isProKey).mockReturnValue(false);
      const key = "df_free_almost1234";
      const monthKey = `${new Date().getFullYear()}-${String(new Date().getMonth() + 1).padStart(2, "0")}`;

      vi.mocked(queryWithRetry).mockResolvedValueOnce({
        rows: [{ key, count: 99, month_key: monthKey }],
        rowCount: 1,
      } as any);
      await usage.loadUsageData();

      const req = makeReq(key);
      const res = makeRes();
      usage.usageMiddleware(req, res, mockNext);

      expect(mockNext).toHaveBeenCalled();
      expect(mockStatus).not.toHaveBeenCalled();
    });

    it("handles missing apiKeyInfo gracefully (uses 'unknown')", () => {
      vi.mocked(isProKey).mockReturnValue(false);
      const req = { apiKeyInfo: undefined } as any;
      const res = makeRes();

      usage.usageMiddleware(req, res, mockNext);

      expect(mockNext).toHaveBeenCalled();
    });
  });
});