import { describe, it, expect, beforeEach, afterEach } from 'vitest';
import Database from 'better-sqlite3';
import { runMigrations } from '../db/migrate.js';
import { createAbuseRepo, type AbuseThresholds } from './abuse-repo.js';

const validKey = 'a'.repeat(64);

function bootstrapDb(): Database.Database {
  process.env.MCP_ENCRYPTION_KEY = validKey;
  const db = new Database(':memory:');
  db.exec(`CREATE TABLE users (id TEXT PRIMARY KEY);`);
  db.exec(`CREATE TABLE jobs (id TEXT PRIMARY KEY, wait_reason TEXT);`);
  db.exec(`CREATE TABLE local_tasks (id INTEGER PRIMARY KEY AUTOINCREMENT);`);
  runMigrations(db);
  return db;
}

const TIGHT: AbuseThresholds = {
  windowMinutes: 10,
  failureThreshold: 3,
  lockMinutes: 30,
};

describe('ssh/abuse-repo', () => {
  let db: Database.Database;
  beforeEach(() => {
    db = bootstrapDb();
  });
  afterEach(() => {
    db.close();
    delete process.env.MCP_ENCRYPTION_KEY;
  });

  describe('constructor', () => {
    it('rejects zero or negative thresholds', () => {
      expect(() =>
        createAbuseRepo(db, { windowMinutes: 0, failureThreshold: 1, lockMinutes: 1 }),
      ).toThrow();
      expect(() =>
        createAbuseRepo(db, { windowMinutes: 1, failureThreshold: 0, lockMinutes: 1 }),
      ).toThrow();
      expect(() =>
        createAbuseRepo(db, { windowMinutes: 1, failureThreshold: 1, lockMinutes: 0 }),
      ).toThrow();
    });
  });

  describe('checkAndRecordFailure', () => {
    const baseArgs = {
      connectionId: 'conn-1',
      ownerId: null,
      userId: 'alice',
      host: 'srv.example.com',
      username: 'deploy',
    };

    it('does NOT lock below threshold', () => {
      const repo = createAbuseRepo(db, TIGHT);
      const now = new Date('2026-05-12T10:00:00Z');
      const r1 = repo.checkAndRecordFailure({ ...baseArgs, now });
      expect(r1.locked).toBe(false);
      const r2 = repo.checkAndRecordFailure({ ...baseArgs, now });
      expect(r2.locked).toBe(false);
    });

    it('locks at exactly threshold for global connection (ownerId=null)', () => {
      const repo = createAbuseRepo(db, TIGHT);
      const now = new Date('2026-05-12T10:00:00Z');
      repo.checkAndRecordFailure({ ...baseArgs, now });
      repo.checkAndRecordFailure({ ...baseArgs, now });
      const r3 = repo.checkAndRecordFailure({ ...baseArgs, now });
      expect(r3.locked).toBe(true);
      expect(r3.lockedScope).toBe('conn');
      expect(r3.lockUntil).toBeTruthy();
      // notifyAdmin true because globalhost scope also hit threshold (global conn).
      expect(r3.notifyAdmin).toBe(true);
    });

    it('user-owned conn does NOT lock via globalhost scope, but notifies admin', () => {
      const repo = createAbuseRepo(db, TIGHT);
      const now = new Date('2026-05-12T10:00:00Z');
      const args = { ...baseArgs, ownerId: 'alice' };
      // Use a different user than the owner so that userhost scope tracks DIFFERENT
      // user (to isolate globalhost effect).
      // Hit threshold via repeated failures from another user on same host
      const otherUser = { ...args, userId: 'bob', connectionId: 'conn-other' };
      // Threshold hits userhost AND globalhost; but conn scope is per-id.
      // We use conn-other and switch userIds to isolate globalhost:
      const calls = [
        { ...otherUser, userId: 'u1', connectionId: 'c1' },
        { ...otherUser, userId: 'u2', connectionId: 'c2' },
        { ...otherUser, userId: 'u3', connectionId: 'c3' },
      ];
      let last;
      for (const a of calls) {
        last = repo.checkAndRecordFailure({ ...a, ownerId: 'someone', now });
      }
      // userhost differs each call (different userIds), so userhost won't lock.
      // conn differs each call, so conn won't lock.
      // globalhost is shared (same host+username) — would lock IF enforce. But
      // ownerId is not null, so enforce_lock=0 for globalhost. Hence not locked.
      expect(last?.locked).toBe(false);
      expect(last?.notifyAdmin).toBe(true);
    });

    it('userhost scope locks when same user retries from different conns', () => {
      const repo = createAbuseRepo(db, TIGHT);
      const now = new Date('2026-05-12T10:00:00Z');
      const userOwnedArgs = { ...baseArgs, ownerId: 'alice' };
      // Same userId, host, username — but different connectionIds.
      let last;
      for (let i = 0; i < 3; i++) {
        last = repo.checkAndRecordFailure({ ...userOwnedArgs, connectionId: `c${i}`, now });
      }
      // Each conn-scope fresh, but userhost accumulates and enforces.
      expect(last?.locked).toBe(true);
      expect(last?.lockedScope).toBe('userhost');
    });

    it('resets count when window has expired (rolling window)', () => {
      const repo = createAbuseRepo(db, TIGHT);
      const t0 = new Date('2026-05-12T10:00:00Z');
      repo.checkAndRecordFailure({ ...baseArgs, now: t0 });
      repo.checkAndRecordFailure({ ...baseArgs, now: t0 });
      // 11 minutes later — outside the 10-min window.
      const tLate = new Date('2026-05-12T10:11:00Z');
      const r = repo.checkAndRecordFailure({ ...baseArgs, now: tLate });
      expect(r.locked).toBe(false); // count was reset to 1
    });
  });

  describe('isLocked', () => {
    it('returns locked=false when no record exists', () => {
      const repo = createAbuseRepo(db, TIGHT);
      expect(repo.isLocked('missing').locked).toBe(false);
    });

    it('returns locked=true after threshold, until lock expires', () => {
      const repo = createAbuseRepo(db, TIGHT);
      const now = new Date('2026-05-12T10:00:00Z');
      for (let i = 0; i < 3; i++) {
        repo.checkAndRecordFailure({
          connectionId: 'conn-1',
          ownerId: null,
          userId: 'a',
          host: 'h',
          username: 'u',
          now,
        });
      }
      expect(repo.isLocked('conn-1', now).locked).toBe(true);
      const future = new Date('2026-05-12T10:31:00Z');
      expect(repo.isLocked('conn-1', future).locked).toBe(false);
    });
  });

  describe('recordSuccess', () => {
    it('clears the conn-scope counter only', () => {
      const repo = createAbuseRepo(db, TIGHT);
      const now = new Date('2026-05-12T10:00:00Z');
      repo.checkAndRecordFailure({
        connectionId: 'conn-1',
        ownerId: null,
        userId: 'alice',
        host: 'h',
        username: 'u',
        now,
      });
      repo.recordSuccess('conn-1');
      // conn scope removed
      expect(repo.getByScopeKey('conn:conn-1')).toBeNull();
      // userhost still present
      expect(repo.getByScopeKey('userhost:alice|h|u')).not.toBeNull();
      // globalhost still present
      expect(repo.getByScopeKey('globalhost:h|u')).not.toBeNull();
    });
  });

  describe('reset', () => {
    it('removes a specific scope row (admin force unlock)', () => {
      const repo = createAbuseRepo(db, TIGHT);
      const now = new Date('2026-05-12T10:00:00Z');
      for (let i = 0; i < 3; i++) {
        repo.checkAndRecordFailure({
          connectionId: 'conn-1',
          ownerId: null,
          userId: 'alice',
          host: 'h',
          username: 'u',
          now,
        });
      }
      expect(repo.isLocked('conn-1', now).locked).toBe(true);
      expect(repo.reset('conn:conn-1')).toBe(true);
      expect(repo.isLocked('conn-1', now).locked).toBe(false);
      expect(repo.reset('conn:conn-1')).toBe(false);
    });
  });
});
