import { describe, it, expect, vi } from 'vitest';
import express from 'express';
import request from 'supertest';
import { buildBudgetCheck, buildBudgetCheckMiddleware } from './budget-check.js';
import type { GatewayKeyUsage } from '../db/repository.js';

function fakeRepo(opts: {
  usage?: Partial<GatewayKeyUsage> | null;
  tokensBudget?: number | null;
}): {
  getGatewayKeyUsage: ReturnType<typeof vi.fn>;
  findGatewayVirtualKeyById: ReturnType<typeof vi.fn>;
} {
  return {
    getGatewayKeyUsage: vi.fn(() => {
      if (opts.usage === undefined || opts.usage === null) return null;
      return {
        keyId: 'k',
        periodStart: '2026-05',
        tokensIn: 0,
        tokensOut: 0,
        requests: 0,
        lastUpdatedAt: '2026-05-19T00:00:00Z',
        ...opts.usage,
      };
    }),
    findGatewayVirtualKeyById: vi.fn(() => ({ tokensBudget: opts.tokensBudget ?? null })),
  };
}

describe('buildBudgetCheck (pure)', () => {
  it('null budget = always allowed', () => {
    const check = buildBudgetCheck({ repo: fakeRepo({}), now: () => new Date(Date.UTC(2026, 4, 19)) });
    const result = check('k', null);
    expect(result.allowed).toBe(true);
    expect(result.budget).toBeNull();
    expect(result.period).toBe('2026-05');
  });

  it('used < budget = allowed', () => {
    const check = buildBudgetCheck({
      repo: fakeRepo({ usage: { tokensIn: 100, tokensOut: 200 } }),
      now: () => new Date(Date.UTC(2026, 4, 19)),
    });
    const result = check('k', 1000);
    expect(result.allowed).toBe(true);
    expect(result.used).toBe(300);
    expect(result.budget).toBe(1000);
  });

  it('used >= budget = rejected', () => {
    const check = buildBudgetCheck({
      repo: fakeRepo({ usage: { tokensIn: 500, tokensOut: 500 } }),
      now: () => new Date(Date.UTC(2026, 4, 19)),
    });
    const result = check('k', 1000);
    expect(result.allowed).toBe(false);
    expect(result.used).toBe(1000);
  });

  it('treats no usage row as zero', () => {
    const check = buildBudgetCheck({
      repo: fakeRepo({ usage: null }),
      now: () => new Date(Date.UTC(2026, 4, 19)),
    });
    const result = check('k', 100);
    expect(result.allowed).toBe(true);
    expect(result.used).toBe(0);
  });

  it('non-positive budget = treated as no limit (defensive)', () => {
    const check = buildBudgetCheck({ repo: fakeRepo({}) });
    expect(check('k', 0).allowed).toBe(true);
    expect(check('k', -5).allowed).toBe(true);
  });
});

describe('buildBudgetCheckMiddleware (express)', () => {
  function buildApp(repo: ReturnType<typeof fakeRepo>, gatewayAuth: unknown): express.Express {
    const app = express();
    app.use((req, _res, next) => {
      (req as unknown as { gatewayAuth: unknown }).gatewayAuth = gatewayAuth;
      next();
    });
    app.use(buildBudgetCheckMiddleware({ repo, now: () => new Date(Date.UTC(2026, 4, 19)) }));
    app.get('/ok', (_req, res) => res.json({ ok: true }));
    return app;
  }

  it('passes through when budget allows', async () => {
    const repo = fakeRepo({ usage: { tokensIn: 100 }, tokensBudget: 1000 });
    const app = buildApp(repo, { team: 'alpha', source: 'db', keyId: 'k1' });
    const res = await request(app).get('/ok');
    expect(res.status).toBe(200);
    expect(res.body.ok).toBe(true);
  });

  it('rejects with 402 + structured detail when exhausted', async () => {
    const repo = fakeRepo({ usage: { tokensIn: 600, tokensOut: 600 }, tokensBudget: 1000 });
    const app = buildApp(repo, { team: 'alpha', source: 'db', keyId: 'k1' });
    const res = await request(app).get('/ok');
    expect(res.status).toBe(402);
    expect(res.body.error).toBe('budget_exhausted');
    expect(res.body.detail).toContain('team=alpha');
    expect(res.body.detail).toContain('used=1200');
    expect(res.body.detail).toContain('budget=1000');
    expect(res.body.detail).toContain('period=2026-05');
  });

  it('skips check when auth came from config (no keyId)', async () => {
    const repo = fakeRepo({ tokensBudget: 1000 });
    const app = buildApp(repo, { team: 'alpha', source: 'config' });
    const res = await request(app).get('/ok');
    expect(res.status).toBe(200);
    // Never read the DB
    expect(repo.findGatewayVirtualKeyById).not.toHaveBeenCalled();
  });

  it('skips check when no gatewayAuth on request (defensive)', async () => {
    const repo = fakeRepo({ tokensBudget: 1000 });
    const app = buildApp(repo, undefined);
    const res = await request(app).get('/ok');
    expect(res.status).toBe(200);
  });

  it('null tokensBudget on the row = always allowed', async () => {
    const repo = fakeRepo({ usage: { tokensIn: 10_000_000 }, tokensBudget: null });
    const app = buildApp(repo, { team: 'alpha', source: 'db', keyId: 'k1' });
    const res = await request(app).get('/ok');
    expect(res.status).toBe(200);
  });
});
