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

describe('rateLimitMiddleware — concurrent tracking', () => {
  let incrementConcurrent, decrementConcurrent, getConcurrentCount, resetConcurrent;

  beforeEach(async () => {
    // Reset concurrent tracking module state and mock rateLimitRepo
    vi.resetModules();

    vi.doMock('@/lib/db/repos/rateLimitRepo', () => ({
      getByUserId: vi.fn().mockResolvedValue({
        maxConcurrent: 5,
        rtpMinute: 100,
        rtpDay: 100000,
        tokenLimitMonthly: 1000000,
        creditLimit: 5.0,
      }),
    }));

    const mod = await import('./rateLimitMiddleware.js');
    incrementConcurrent = mod.incrementConcurrent;
    decrementConcurrent = mod.decrementConcurrent;
    getConcurrentCount = mod.getConcurrentCount;
    resetConcurrent = mod.resetConcurrent;
    resetConcurrent();
  });

  afterEach(() => {
    resetConcurrent();
  });

  it('incrementConcurrent returns true for first request', async () => {
    expect(await incrementConcurrent('user-1')).toBe(true);
    expect(getConcurrentCount('user-1')).toBe(1);
  });

  it('incrementConcurrent tracks up to max (default 5)', async () => {
    const userId = 'user-2';
    for (let i = 0; i < 5; i++) {
      expect(await incrementConcurrent(userId)).toBe(true);
    }
    expect(getConcurrentCount(userId)).toBe(5);
  });

  it('incrementConcurrent returns false when at max capacity', async () => {
    const userId = 'user-3';
    for (let i = 0; i < 5; i++) await incrementConcurrent(userId);
    expect(await incrementConcurrent(userId)).toBe(false);
    expect(getConcurrentCount(userId)).toBe(5); // count stays at 5
  });

  it('decrementConcurrent reduces count', async () => {
    const userId = 'user-4';
    await incrementConcurrent(userId);
    await incrementConcurrent(userId);
    expect(getConcurrentCount(userId)).toBe(2);
    decrementConcurrent(userId);
    expect(getConcurrentCount(userId)).toBe(1);
    decrementConcurrent(userId);
    expect(getConcurrentCount(userId)).toBe(0);
  });

  it('decrementConcurrent is safe for unknown userId', () => {
    expect(() => decrementConcurrent('no-such-user')).not.toThrow();
    expect(getConcurrentCount('no-such-user')).toBe(0);
  });

  it('resetConcurrent clears all counts', async () => {
    await incrementConcurrent('user-a');
    await incrementConcurrent('user-b');
    resetConcurrent();
    expect(getConcurrentCount('user-a')).toBe(0);
    expect(getConcurrentCount('user-b')).toBe(0);
  });

  it('defaults maxConcurrent to 5 when limits not found', async () => {
    vi.resetModules();
    vi.doMock('@/lib/db/repos/rateLimitRepo', () => ({
      getByUserId: vi.fn().mockResolvedValue(null),
    }));
    const mod = await import('./rateLimitMiddleware.js');
    mod.resetConcurrent();

    const userId = 'no-limits-user';
    for (let i = 0; i < 5; i++) {
      expect(await mod.incrementConcurrent(userId)).toBe(true);
    }
    expect(await mod.incrementConcurrent(userId)).toBe(false);
    mod.resetConcurrent();
  });
});

describe('rateLimitMiddleware — checkRateLimits', () => {
  let checkRateLimits;

  beforeEach(async () => {
    vi.resetModules();
  });

  afterEach(() => {
    vi.restoreAllMocks();
  });

  it('returns allowed:true when no apiKey provided', async () => {
    // Fresh import with no mocks needed — first condition is !apiKey
    const mod = await import('./rateLimitMiddleware.js');
    checkRateLimits = mod.checkRateLimits;
    const result = await checkRateLimits(null);
    expect(result).toEqual({ allowed: true });
  });

  it('returns allowed:true for empty string apiKey', async () => {
    const mod = await import('./rateLimitMiddleware.js');
    const result = await mod.checkRateLimits('');
    expect(result).toEqual({ allowed: true });
  });

  it('returns allowed:true when apiKey not found in DB', async () => {
    vi.doMock('@/lib/db/driver.js', () => ({
      getAdapter: vi.fn().mockResolvedValue({
        get: vi.fn().mockReturnValue(null),
      }),
    }));

    const mod = await import('./rateLimitMiddleware.js');
    const result = await mod.checkRateLimits('unknown-key-123');
    expect(result).toEqual({ allowed: true });
  });

  it('fails open on unexpected errors', async () => {
    vi.doMock('@/lib/db/driver.js', () => ({
      getAdapter: vi.fn().mockRejectedValue(new Error('DB connection lost')),
    }));

    const mod = await import('./rateLimitMiddleware.js');
    const result = await mod.checkRateLimits('some-key');
    expect(result).toEqual({ allowed: true });
  });
});

// ─── Shared mock helpers ──────────────────────────────────────────────────

function mockDriver(getImplementation) {
  vi.doMock('@/lib/db/driver.js', () => ({
    getAdapter: vi.fn().mockResolvedValue({
      get: vi.fn().mockImplementation(getImplementation),
    }),
  }));
}

function mockUser(role = 'user') {
  vi.doMock('@/lib/db/repos/usersRepo', () => ({
    findById: vi.fn().mockResolvedValue({ id: 'user-1', role }),
  }));
}

function mockRateLimits(overrides = {}) {
  vi.doMock('@/lib/db/repos/rateLimitRepo', () => ({
    getByUserId: vi.fn().mockResolvedValue({
      rtpMinute: 100,
      rtpDay: 100000,
      tokenLimitMonthly: 1000000,
      creditLimit: 5.0,
      maxConcurrent: 5,
      ...overrides,
    }),
    getMonthlyTokens: vi.fn().mockResolvedValue(overrides._monthlyTokens ?? 0),
    getMonthlyCost: vi.fn().mockResolvedValue(overrides._monthlyCost ?? 0),
  }));
}

function mockEmptyRates() {
  vi.doMock('@/lib/db/repos/rateLimitRepo', () => ({
    getByUserId: vi.fn().mockResolvedValue(null),
    getMonthlyTokens: vi.fn().mockResolvedValue(0),
    getMonthlyCost: vi.fn().mockResolvedValue(0),
  }));
}

describe('rateLimitMiddleware — enforceRateLimits', () => {
  beforeEach(async () => {
    vi.resetModules();
  });

  afterEach(() => {
    vi.restoreAllMocks();
  });

  it('returns null when rate limit check passes', async () => {
    const mod = await import('./rateLimitMiddleware.js');
    const result = await mod.enforceRateLimits('');
    expect(result).toBeNull();
  });

  it('returns 429 JSON response when rate limited', async () => {
    // First get() = apiKey lookup, second get() = usageHistory COUNT
    let calls = 0;
    vi.doMock('@/lib/db/driver.js', () => ({
      getAdapter: vi.fn().mockResolvedValue({
        get: vi.fn().mockImplementation(() => {
          calls++;
          return calls === 1 ? { userId: 'user-1' } : { cnt: 1 };
        }),
      }),
    }));

    vi.doMock('@/lib/db/repos/usersRepo', () => ({
      findById: vi.fn().mockResolvedValue({ id: 'user-1', role: 'user' }),
    }));

    vi.doMock('@/lib/db/repos/rateLimitRepo', () => ({
      getByUserId: vi.fn().mockResolvedValue({
        rtpMinute: 1,
        rtpDay: 100000,
        tokenLimitMonthly: 1000000,
        creditLimit: 5.0,
        maxConcurrent: 5,
      }),
      getMonthlyTokens: vi.fn().mockResolvedValue(0),
      getMonthlyCost: vi.fn().mockResolvedValue(0),
    }));

    const mod = await import('./rateLimitMiddleware.js');
    const result = await mod.enforceRateLimits('some-key');

    expect(result).not.toBeNull();
    // NextResponse.json() returns a Response-like object
    expect(result.status).toBe(429);
  });
});

// ═══════════════════════════════════════════════════════════════════════════
// Integration tests — all 4 rate limit strategies
// ═══════════════════════════════════════════════════════════════════════════

describe('rateLimitMiddleware — strategy: RTPM (requests per minute)', () => {
  beforeEach(async () => {
    vi.resetModules();
  });
  afterEach(() => { vi.restoreAllMocks(); });

  it('allows when usage count < rtpMinute limit', async () => {
    // get() call 1 = apiKey lookup → userId, call 2 = COUNT query
    let calls = 0;
    mockDriver(function () {
      calls++;
      return calls === 1 ? { userId: 'user-1' } : { cnt: 49 };
    });
    mockUser('user');
    mockRateLimits({ rtpMinute: 50 }); // 49 < 50 → pass

    const mod = await import('./rateLimitMiddleware.js');
    const result = await mod.checkRateLimits('test-key');
    expect(result.allowed).toBe(true);
  });

  it('blocks when usage count >= rtpMinute limit', async () => {
    let calls = 0;
    mockDriver(function () {
      calls++;
      return calls === 1 ? { userId: 'user-1' } : { cnt: 50 };
    });
    mockUser('user');
    mockRateLimits({ rtpMinute: 50 }); // 50 >= 50 → block

    const mod = await import('./rateLimitMiddleware.js');
    const result = await mod.checkRateLimits('test-key');
    expect(result.allowed).toBe(false);
    expect(result.reason).toContain('requests per minute');
  });

  it('bypasses RTPM when limit is 0 (unlimited)', async () => {
    let calls = 0;
    mockDriver(function () {
      calls++;
      return calls === 1 ? { userId: 'user-1' } : { cnt: 999 };
    });
    mockUser('user');
    mockRateLimits({ rtpMinute: 0 }); // 0 = unlimited

    const mod = await import('./rateLimitMiddleware.js');
    const result = await mod.checkRateLimits('test-key');
    expect(result.allowed).toBe(true);
  });
});

describe('rateLimitMiddleware — strategy: RTPD (requests per day)', () => {
  beforeEach(async () => {
    vi.resetModules();
  });
  afterEach(() => { vi.restoreAllMocks(); });

  it('allows when daily usage < rtpDay limit', async () => {
    let calls = 0;
    mockDriver(function () {
      calls++;
      return calls === 1 ? { userId: 'user-1' } : { cnt: 49 };
    });
    mockUser('user');
    mockRateLimits({ rtpMinute: 0, rtpDay: 50 }); // 0 RTPM = unlimited, 49 < 50 RTPD

    const mod = await import('./rateLimitMiddleware.js');
    const result = await mod.checkRateLimits('test-key');
    expect(result.allowed).toBe(true);
  });

  it('blocks when daily usage >= rtpDay limit', async () => {
    let calls = 0;
    mockDriver(function () {
      calls++;
      return calls === 1 ? { userId: 'user-1' } : { cnt: 50 };
    });
    mockUser('user');
    mockRateLimits({ rtpMinute: 0, rtpDay: 50 }); // 0 RTPM = unlimited, 50 >= 50 RTPD

    const mod = await import('./rateLimitMiddleware.js');
    const result = await mod.checkRateLimits('test-key');
    expect(result.allowed).toBe(false);
    expect(result.reason).toContain('requests per day');
  });

  it('bypasses RTPD when limit is 0 (unlimited)', async () => {
    let calls = 0;
    mockDriver(function () {
      calls++;
      return calls === 1 ? { userId: 'user-1' } : { cnt: 9999 };
    });
    mockUser('user');
    mockRateLimits({ rtpMinute: 0, rtpDay: 0 }); // both unlimited

    const mod = await import('./rateLimitMiddleware.js');
    const result = await mod.checkRateLimits('test-key');
    expect(result.allowed).toBe(true);
  });
});

describe('rateLimitMiddleware — strategy: token count (monthly)', () => {
  beforeEach(async () => {
    vi.resetModules();
  });
  afterEach(() => { vi.restoreAllMocks(); });

  it('allows when monthly tokens < limit', async () => {
    mockDriver(() => ({ userId: 'user-1' }));
    mockUser('user');
    mockRateLimits({
      rtpMinute: 0,
      rtpDay: 0,
      tokenLimitMonthly: 100000,
      _monthlyTokens: 50000,
    });

    const mod = await import('./rateLimitMiddleware.js');
    const result = await mod.checkRateLimits('test-key');
    expect(result.allowed).toBe(true);
  });

  it('blocks when monthly tokens >= limit', async () => {
    mockDriver(() => ({ userId: 'user-1' }));
    mockUser('user');
    mockRateLimits({
      rtpMinute: 0,
      rtpDay: 0,
      tokenLimitMonthly: 100000,
      _monthlyTokens: 100000,
    });

    const mod = await import('./rateLimitMiddleware.js');
    const result = await mod.checkRateLimits('test-key');
    expect(result.allowed).toBe(false);
    expect(result.reason).toContain('Token limit exceeded');
  });

  it('bypasses token check when limit is 0 (unlimited)', async () => {
    mockDriver(() => ({ userId: 'user-1' }));
    mockUser('user');
    mockRateLimits({
      rtpMinute: 0,
      rtpDay: 0,
      tokenLimitMonthly: 0,
      _monthlyTokens: 99999999,
    });

    const mod = await import('./rateLimitMiddleware.js');
    const result = await mod.checkRateLimits('test-key');
    expect(result.allowed).toBe(true);
  });
});

describe('rateLimitMiddleware — strategy: credit/quota (monthly cost)', () => {
  beforeEach(async () => {
    vi.resetModules();
  });
  afterEach(() => { vi.restoreAllMocks(); });

  it('allows when monthly cost < creditLimit', async () => {
    mockDriver(() => ({ userId: 'user-1' }));
    mockUser('user');
    mockRateLimits({
      rtpMinute: 0,
      rtpDay: 0,
      tokenLimitMonthly: 0,
      creditLimit: 5.0,
      _monthlyCost: 3.5,
    });

    const mod = await import('./rateLimitMiddleware.js');
    const result = await mod.checkRateLimits('test-key');
    expect(result.allowed).toBe(true);
  });

  it('blocks when monthly cost >= creditLimit', async () => {
    mockDriver(() => ({ userId: 'user-1' }));
    mockUser('user');
    mockRateLimits({
      rtpMinute: 0,
      rtpDay: 0,
      tokenLimitMonthly: 0,
      creditLimit: 5.0,
      _monthlyCost: 5.0,
    });

    const mod = await import('./rateLimitMiddleware.js');
    const result = await mod.checkRateLimits('test-key');
    expect(result.allowed).toBe(false);
    expect(result.reason).toContain('Credit/quota exceeded');
  });

  it('bypasses credit check when limit is 0 (unlimited)', async () => {
    mockDriver(() => ({ userId: 'user-1' }));
    mockUser('user');
    mockRateLimits({
      rtpMinute: 0,
      rtpDay: 0,
      tokenLimitMonthly: 0,
      creditLimit: 0,
      _monthlyCost: 999.99,
    });

    const mod = await import('./rateLimitMiddleware.js');
    const result = await mod.checkRateLimits('test-key');
    expect(result.allowed).toBe(true);
  });
});

describe('rateLimitMiddleware — admin bypass', () => {
  beforeEach(async () => {
    vi.resetModules();
  });
  afterEach(() => { vi.restoreAllMocks(); });

  it('admin user bypasses all rate limits even when exceeded', async () => {
    let calls = 0;
    mockDriver(function () {
      calls++;
      return calls === 1 ? { userId: 'admin-1' } : { cnt: 999 };
    });
    mockUser('admin'); // admin role
    // Set all limits to be exceeded
    mockRateLimits({
      rtpMinute: 1,
      rtpDay: 1,
      tokenLimitMonthly: 1,
      creditLimit: 0.01,
      _monthlyTokens: 999999,
      _monthlyCost: 999.99,
    });

    const mod = await import('./rateLimitMiddleware.js');
    const result = await mod.checkRateLimits('admin-key');
    expect(result.allowed).toBe(true);
  });
});

describe('rateLimitMiddleware — edge cases', () => {
  beforeEach(async () => {
    vi.resetModules();
  });
  afterEach(() => { vi.restoreAllMocks(); });

  it('passes when user has no rateLimits row', async () => {
    mockDriver(() => ({ userId: 'user-1' }));
    mockUser('user');
    mockEmptyRates(); // getByUserId returns null

    const mod = await import('./rateLimitMiddleware.js');
    const result = await mod.checkRateLimits('test-key');
    expect(result.allowed).toBe(true);
  });

  it('first strategy hit short-circuits remaining checks', async () => {
    // RTPM blocks first, RTPD and others never checked
    let calls = 0;
    mockDriver(function () {
      calls++;
      // call 1: apiKey lookup, call 2: RTPM COUNT
      return calls === 1 ? { userId: 'user-1' } : { cnt: 100 };
    });
    mockUser('user');
    mockRateLimits({ rtpMinute: 1, rtpDay: 1 }); // RTPM hits first

    const mod = await import('./rateLimitMiddleware.js');
    const result = await mod.checkRateLimits('test-key');
    expect(result.allowed).toBe(false);
    expect(result.reason).toContain('requests per minute');
    // Only 2 DB calls: apiKey lookup + RTPM count; RTPD never reached
    expect(calls).toBe(2);
  });
});
