import { describe, it, expect, beforeEach, afterEach } from 'vitest';
import { generateKeyPairSync } from 'node:crypto';
import Database from 'better-sqlite3';
import { runMigrations } from '../db/migrate.js';
import {
  bootstrapSystemDek,
  verifySystemDek,
  getSystemDek,
  getOrCreateUserDek,
  encryptPrivateKey,
  decryptPrivateKey,
  computeKeyFingerprint,
  formatPublicKey,
  generateKeypair,
  sanitizeError,
  clearBuffer,
} from './crypto.js';

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

// Generate a real RSA PKCS#1 PEM at module load. ssh2.utils.parseKey accepts
// PKCS#1 "RSA PRIVATE KEY" PEM. (ssh2 doesn't accept PKCS#8.)
// We use RSA-2048 (fast enough for tests) rather than Ed25519 because Node's
// Ed25519 export defaults to PKCS#8 which ssh2 rejects.
function makePem(): Buffer {
  const { privateKey } = generateKeyPairSync('rsa', {
    modulusLength: 2048,
    publicKeyEncoding: { type: 'spki', format: 'pem' },
    privateKeyEncoding: { type: 'pkcs1', format: 'pem' },
  });
  return Buffer.from(privateKey as string, 'utf-8');
}

const TEST_PEM = makePem();

const ENC_PASSPHRASE = 'test-pass-1';
function makeEncryptedOpenSshKey(passphrase: string): Buffer {
  // ssh2.utils.generateKeyPairSync emits OpenSSH-format private keys; passing
  // a passphrase encrypts the private portion (cipher defaults to aes256-ctr).
  // eslint-disable-next-line @typescript-eslint/no-require-imports
  const ssh2 = require('ssh2') as {
    utils: {
      generateKeyPairSync: (
        keyType: string,
        opts: { bits?: number; passphrase?: string; cipher?: string },
      ) => { private: string };
    };
  };
  const { private: priv } = ssh2.utils.generateKeyPairSync('rsa', {
    bits: 2048,
    passphrase,
    cipher: 'aes256-cbc',
  });
  return Buffer.from(priv, 'utf-8');
}

const ENC_PEM = makeEncryptedOpenSshKey(ENC_PASSPHRASE);

describe('ssh/crypto', () => {
  let db: Database.Database;

  beforeEach(() => {
    process.env.MCP_ENCRYPTION_KEY = validKey;
    db = new Database(':memory:');
    db.pragma('foreign_keys = ON');
    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);
    for (const id of ['alice', 'bob']) {
      db.prepare('INSERT INTO users(id) VALUES(?)').run(id);
    }
  });

  afterEach(() => {
    db.close();
    delete process.env.MCP_ENCRYPTION_KEY;
  });

  describe('system DEK', () => {
    it('bootstrap is idempotent', () => {
      bootstrapSystemDek(db);
      const row1 = db.prepare('SELECT encrypted_dek FROM system_deks WHERE id=1').get() as
        | { encrypted_dek: Buffer }
        | undefined;
      expect(row1).toBeTruthy();
      bootstrapSystemDek(db);
      const row2 = db.prepare('SELECT encrypted_dek FROM system_deks WHERE id=1').get() as
        | { encrypted_dek: Buffer }
        | undefined;
      // The second bootstrap is a no-op: blob should match.
      expect(row2!.encrypted_dek.equals(row1!.encrypted_dek)).toBe(true);
    });

    it('verifySystemDek succeeds with matching master key', () => {
      bootstrapSystemDek(db);
      expect(() => verifySystemDek(db)).not.toThrow();
    });

    it('verifySystemDek throws if master key changed', () => {
      bootstrapSystemDek(db);
      process.env.MCP_ENCRYPTION_KEY = otherKey;
      expect(() => verifySystemDek(db)).toThrow();
    });

    it('verifySystemDek throws if bootstrap missing', () => {
      expect(() => verifySystemDek(db)).toThrow(/bootstrap row missing/);
    });

    it('getSystemDek returns 32 bytes', () => {
      bootstrapSystemDek(db);
      const dek = getSystemDek(db);
      try {
        expect(dek.length).toBe(32);
      } finally {
        clearBuffer(dek);
      }
    });
  });

  describe('per-user DEK', () => {
    it('lazy-creates on first call', () => {
      const row1 = db.prepare('SELECT * FROM ssh_user_deks WHERE user_id = ?').get('alice');
      expect(row1).toBeUndefined();
      const dek = getOrCreateUserDek(db, 'alice');
      try {
        expect(dek.length).toBe(32);
        const row2 = db.prepare('SELECT * FROM ssh_user_deks WHERE user_id = ?').get('alice');
        expect(row2).toBeTruthy();
      } finally {
        clearBuffer(dek);
      }
    });

    it('returns same DEK across calls', () => {
      const a1 = getOrCreateUserDek(db, 'alice');
      const a2 = getOrCreateUserDek(db, 'alice');
      try {
        expect(a1.equals(a2)).toBe(true);
      } finally {
        clearBuffer(a1);
        clearBuffer(a2);
      }
    });

    it('returns different DEKs across users', () => {
      const a = getOrCreateUserDek(db, 'alice');
      const b = getOrCreateUserDek(db, 'bob');
      try {
        expect(a.equals(b)).toBe(false);
      } finally {
        clearBuffer(a);
        clearBuffer(b);
      }
    });
  });

  describe('encrypt/decrypt private key (user-owned)', () => {
    it('round-trips a PEM buffer under per-user DEK', () => {
      const { blob } = encryptPrivateKey(db, 'alice', TEST_PEM);
      const dec = decryptPrivateKey(db, 'alice', blob);
      try {
        expect(dec.equals(TEST_PEM)).toBe(true);
      } finally {
        clearBuffer(dec);
      }
    });

    it('decryption fails if attempted under wrong user', () => {
      const { blob } = encryptPrivateKey(db, 'alice', TEST_PEM);
      expect(() => decryptPrivateKey(db, 'bob', blob)).toThrow();
    });
  });

  describe('encrypt/decrypt private key (global / system DEK)', () => {
    it('round-trips a PEM buffer under system DEK (ownerId=null)', () => {
      bootstrapSystemDek(db);
      const { blob } = encryptPrivateKey(db, null, TEST_PEM);
      const dec = decryptPrivateKey(db, null, blob);
      try {
        expect(dec.equals(TEST_PEM)).toBe(true);
      } finally {
        clearBuffer(dec);
      }
    });

    it('throws if system DEK not bootstrapped', () => {
      expect(() => encryptPrivateKey(db, null, TEST_PEM)).toThrow(/not bootstrapped/);
    });

    it('user-encrypted key cannot be decrypted as global', () => {
      bootstrapSystemDek(db);
      const { blob } = encryptPrivateKey(db, 'alice', TEST_PEM);
      expect(() => decryptPrivateKey(db, null, blob)).toThrow();
    });
  });

  describe('computeKeyFingerprint', () => {
    it('produces a stable SHA256:<b64> fingerprint', () => {
      const fp = computeKeyFingerprint(TEST_PEM);
      expect(fp).toMatch(/^SHA256:[A-Za-z0-9+/]+$/);
      // Same input → same fingerprint
      expect(computeKeyFingerprint(TEST_PEM)).toBe(fp);
    });

    it('throws on invalid PEM', () => {
      expect(() => computeKeyFingerprint(Buffer.from('not a pem'))).toThrow();
    });

    it('parses encrypted OpenSSH PEM when passphrase is supplied', () => {
      const fp = computeKeyFingerprint(ENC_PEM, Buffer.from(ENC_PASSPHRASE, 'utf-8'));
      expect(fp).toMatch(/^SHA256:[A-Za-z0-9+/]+$/);
    });

    it('throws on encrypted OpenSSH PEM when passphrase is missing', () => {
      expect(() => computeKeyFingerprint(ENC_PEM)).toThrow(/passphrase/i);
    });

    it('throws on encrypted OpenSSH PEM when passphrase is wrong', () => {
      expect(() => computeKeyFingerprint(ENC_PEM, Buffer.from('wrong', 'utf-8'))).toThrow();
    });
  });

  describe('formatPublicKey', () => {
    it('returns "<algo> <base64>" for plain PKCS#1 PEM', () => {
      const pub = formatPublicKey(TEST_PEM);
      // ssh2 reports PKCS#1 RSA as `ssh-rsa`.
      expect(pub).toMatch(/^ssh-rsa [A-Za-z0-9+/]+={0,2}$/);
    });

    it('returns "<algo> <base64>" for encrypted OpenSSH PEM with passphrase', () => {
      const pub = formatPublicKey(ENC_PEM, Buffer.from(ENC_PASSPHRASE, 'utf-8'));
      expect(pub).toMatch(/^ssh-rsa [A-Za-z0-9+/]+={0,2}$/);
    });

    it('public key is consistent with computeKeyFingerprint', () => {
      // Same private key → same public bytes → same SHA256 fingerprint.
      const fp = computeKeyFingerprint(TEST_PEM);
      const pub = formatPublicKey(TEST_PEM);
      const pubB64 = pub.split(' ')[1];
      const { createHash } = require('node:crypto') as typeof import('node:crypto');
      const fp2 = `SHA256:${createHash('sha256').update(Buffer.from(pubB64, 'base64')).digest('base64').replace(/=+$/, '')}`;
      expect(fp2).toBe(fp);
    });
  });

  describe('generateKeypair', () => {
    it('generates a usable ed25519 keypair', () => {
      const { privateKeyPem, publicKey } = generateKeypair('ed25519');
      expect(privateKeyPem.toString('utf-8')).toContain('OPENSSH PRIVATE KEY');
      expect(publicKey).toMatch(/^ssh-ed25519 [A-Za-z0-9+/]+={0,2}$/);
      // The generated private key parses back to the same public portion.
      const derived = formatPublicKey(privateKeyPem);
      expect(derived).toBe(publicKey);
    });

    it('generates a usable rsa-4096 keypair', () => {
      const { privateKeyPem, publicKey } = generateKeypair('rsa-4096');
      expect(privateKeyPem.toString('utf-8')).toContain('OPENSSH PRIVATE KEY');
      expect(publicKey).toMatch(/^ssh-rsa [A-Za-z0-9+/]+={0,2}$/);
      const derived = formatPublicKey(privateKeyPem);
      expect(derived).toBe(publicKey);
    }, 30_000);  // RSA-4096 generation can be slow
  });

  describe('sanitizeError', () => {
    it('strips PEM blocks from message', () => {
      const dirty = new Error(
        `decrypt failed for key: -----BEGIN OPENSSH PRIVATE KEY-----\nAAAA\nBBBB\n-----END OPENSSH PRIVATE KEY-----`,
      );
      const cleaned = sanitizeError(dirty);
      expect(cleaned.message).toContain('[REDACTED PEM]');
      expect(cleaned.message).not.toContain('BEGIN OPENSSH');
      expect(cleaned.message).not.toContain('AAAA');
    });

    it('leaves messages without PEM untouched', () => {
      const e = new Error('something else failed');
      expect(sanitizeError(e).message).toBe('something else failed');
    });

    it('handles non-Error input', () => {
      expect(sanitizeError('plain string').message).toContain('unknown');
      expect(sanitizeError(null).message).toContain('unknown');
    });
  });

  describe('clearBuffer', () => {
    it('zeros out a buffer', () => {
      const b = Buffer.from([1, 2, 3, 4]);
      clearBuffer(b);
      expect(b.equals(Buffer.from([0, 0, 0, 0]))).toBe(true);
    });

    it('no-ops on empty / null', () => {
      expect(() => clearBuffer(null)).not.toThrow();
      expect(() => clearBuffer(undefined)).not.toThrow();
      expect(() => clearBuffer(Buffer.alloc(0))).not.toThrow();
    });
  });
});
