import { describe, it, expect, beforeEach, afterEach } from 'vitest';
import Database from 'better-sqlite3';
import { runMigrations } from '../db/migrate.js';
import {
  createConnectionRepo,
  parseHostKeyType,
  type CreateConnectionInput,
} from './connection-repo.js';

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

function bootstrapDb(): Database.Database {
  process.env.MCP_ENCRYPTION_KEY = validKey;
  const 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);
  db.prepare(`INSERT INTO users(id) VALUES (?), (?)`).run('alice', 'bob');
  return db;
}

function baseInput(overrides: Partial<CreateConnectionInput> = {}): CreateConnectionInput {
  return {
    ownerId: 'alice',
    label: 'prod-srv',
    host: 'srv.example.com',
    port: 22,
    username: 'deploy',
    privateKeyEnc: Buffer.from([1, 2, 3]),
    keyFingerprint: 'SHA256:abc',
    remotePathPrefix: '/home/deploy',
    ...overrides,
  };
}

// Build a wire-format SSH host key blob with the given algorithm name + dummy payload,
// then base64-encode it. This matches what an actual hostVerifier raw buffer contains.
function wireFormatHostKeyB64(algorithm: string): string {
  const nameBuf = Buffer.from(algorithm, 'utf8');
  const lenBuf = Buffer.alloc(4);
  lenBuf.writeUInt32BE(nameBuf.length, 0);
  const payload = Buffer.from([0xaa, 0xbb, 0xcc, 0xdd]);
  return Buffer.concat([lenBuf, nameBuf, payload]).toString('base64');
}

describe('parseHostKeyType', () => {
  it('extracts the algorithm name from valid wire format', () => {
    expect(parseHostKeyType(wireFormatHostKeyB64('ssh-ed25519'))).toBe('ssh-ed25519');
    expect(parseHostKeyType(wireFormatHostKeyB64('ssh-rsa'))).toBe('ssh-rsa');
    expect(parseHostKeyType(wireFormatHostKeyB64('ecdsa-sha2-nistp256'))).toBe(
      'ecdsa-sha2-nistp256',
    );
  });

  it('returns null for malformed/too-short input', () => {
    expect(parseHostKeyType('')).toBeNull();
    expect(parseHostKeyType('AAAA')).toBeNull(); // valid base64 but only 3 bytes decoded
    expect(parseHostKeyType('this is not base64 padding!!')).toBeNull();
  });

  it('returns null when claimed length exceeds buffer', () => {
    const lenBuf = Buffer.alloc(4);
    lenBuf.writeUInt32BE(1000, 0); // way too big
    const malformed = Buffer.concat([lenBuf, Buffer.from('short')]).toString('base64');
    expect(parseHostKeyType(malformed)).toBeNull();
  });

  it('rejects non-printable name bytes', () => {
    const nameBuf = Buffer.from([0xff, 0xfe, 0xfd]);
    const lenBuf = Buffer.alloc(4);
    lenBuf.writeUInt32BE(nameBuf.length, 0);
    const b64 = Buffer.concat([lenBuf, nameBuf, Buffer.from([0xaa])]).toString('base64');
    expect(parseHostKeyType(b64)).toBeNull();
  });
});

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

  it('creates and round-trips through resolveConnection', () => {
    const repo = createConnectionRepo(db);
    const created = repo.create(baseInput({ now: '2026-05-12T10:00:00.000Z' }));
    expect(created.id).toBeTruthy();
    expect(created.label).toBe('prod-srv');
    expect(created.enabled).toBe(true);
    expect(created.hostKeyPending).toBe(false);
    expect(created.allowPrivateAddresses).toBe(false);
    const got = repo.resolveConnection(created.id);
    expect(got?.host).toBe('srv.example.com');
    expect(got?.privateKeyEnc.equals(Buffer.from([1, 2, 3]))).toBe(true);
  });

  it('rejects empty remotePathPrefix at create', () => {
    const repo = createConnectionRepo(db);
    expect(() => repo.create(baseInput({ remotePathPrefix: '' }))).toThrow(/non-empty/);
  });

  it('rejects empty remotePathPrefix at update', () => {
    const repo = createConnectionRepo(db);
    const c = repo.create(baseInput());
    expect(() => repo.update(c.id, { remotePathPrefix: '' })).toThrow(/non-empty/);
  });

  it('updates patch-by-patch and bumps updated_at', () => {
    const repo = createConnectionRepo(db);
    const c = repo.create(baseInput({ now: '2026-05-12T10:00:00.000Z' }));
    const ok = repo.update(c.id, { label: 'renamed' }, '2026-05-12T11:00:00.000Z');
    expect(ok).toBe(true);
    const after = repo.resolveConnection(c.id);
    expect(after?.label).toBe('renamed');
    expect(after?.updatedAt).toBe('2026-05-12T11:00:00.000Z');
  });

  it('update returns false when no patch fields supplied', () => {
    const repo = createConnectionRepo(db);
    const c = repo.create(baseInput());
    expect(repo.update(c.id, {})).toBe(false);
  });

  it('listOwned returns only the owner rows', () => {
    const repo = createConnectionRepo(db);
    repo.create(baseInput({ ownerId: 'alice', label: 'a1' }));
    repo.create(baseInput({ ownerId: 'alice', label: 'a2' }));
    repo.create(baseInput({ ownerId: 'bob', label: 'b1' }));
    expect(repo.listOwned('alice')).toHaveLength(2);
    expect(repo.listOwned('bob')).toHaveLength(1);
  });

  it('listAll returns all rows (admin)', () => {
    const repo = createConnectionRepo(db);
    repo.create(baseInput({ ownerId: 'alice' }));
    repo.create(baseInput({ ownerId: 'bob' }));
    repo.create(baseInput({ ownerId: null, label: 'global' }));
    expect(repo.listAll()).toHaveLength(3);
  });

  it('delete removes the row', () => {
    const repo = createConnectionRepo(db);
    const c = repo.create(baseInput());
    expect(repo.delete(c.id)).toBe(true);
    expect(repo.resolveConnection(c.id)).toBeNull();
    expect(repo.delete(c.id)).toBe(false);
  });
});

describe('ssh/connection-repo admin enable/disable', () => {
  let db: Database.Database;
  beforeEach(() => {
    db = bootstrapDb();
    db.prepare(`INSERT INTO users(id) VALUES (?)`).run('admin');
  });
  afterEach(() => {
    db.close();
    delete process.env.MCP_ENCRYPTION_KEY;
  });

  it('disableByAdmin sets enabled=0 with reason + actor', () => {
    const repo = createConnectionRepo(db);
    const c = repo.create(baseInput());
    const ok = repo.disableByAdmin(
      c.id,
      'compliance violation',
      'admin',
      '2026-05-12T12:00:00.000Z',
    );
    expect(ok).toBe(true);
    const after = repo.resolveConnection(c.id);
    expect(after?.enabled).toBe(false);
    expect(after?.disabledByAdmin).toBe(true);
    expect(after?.disabledByAdminReason).toBe('compliance violation');
    expect(after?.disabledByAdminUserId).toBe('admin');
    expect(after?.disabledByAdminAt).toBe('2026-05-12T12:00:00.000Z');
  });

  it('enableByAdmin clears admin-disable state', () => {
    const repo = createConnectionRepo(db);
    const c = repo.create(baseInput());
    repo.disableByAdmin(c.id, 'a test reason', 'admin');
    expect(repo.enableByAdmin(c.id)).toBe(true);
    const after = repo.resolveConnection(c.id);
    expect(after?.enabled).toBe(true);
    expect(after?.disabledByAdmin).toBe(false);
    expect(after?.disabledByAdminReason).toBeNull();
    expect(after?.disabledByAdminUserId).toBeNull();
  });
});

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

  it('setHostKeyPendingWithToken issues fresh UUID and stores pending state', () => {
    const repo = createConnectionRepo(db);
    const c = repo.create(baseInput());
    const b64 = wireFormatHostKeyB64('ssh-ed25519');
    const fp = 'SHA256:fp-1';
    const r = repo.setHostKeyPendingWithToken(c.id, b64, fp, 'tofu_record');
    expect(r).not.toBeNull();
    expect(r!.token).toMatch(/^[0-9a-f-]+$/);
    const after = repo.resolveConnection(c.id);
    expect(after?.hostKeyPending).toBe(true);
    expect(after?.hostKeyPendingB64).toBe(b64);
    expect(after?.hostKeyPendingFingerprint).toBe(fp);
    expect(after?.hostKeyPendingSource).toBe('tofu_record');
    expect(after?.hostKeyPendingToken).toBe(r!.token);
  });

  it('setHostKeyPendingWithToken returns null for missing connection', () => {
    const repo = createConnectionRepo(db);
    const r = repo.setHostKeyPendingWithToken('does-not-exist', 'x', 'fp', 'tofu_record');
    expect(r).toBeNull();
  });

  it('setHostKeyPendingWithToken replaces previous pending (token rotates)', () => {
    const repo = createConnectionRepo(db);
    const c = repo.create(baseInput());
    const r1 = repo.setHostKeyPendingWithToken(
      c.id,
      wireFormatHostKeyB64('ssh-ed25519'),
      'fp1',
      'tofu_record',
    );
    const r2 = repo.setHostKeyPendingWithToken(
      c.id,
      wireFormatHostKeyB64('ssh-rsa'),
      'fp2',
      'mismatch',
    );
    expect(r1!.token).not.toBe(r2!.token);
    const after = repo.resolveConnection(c.id);
    expect(after?.hostKeyPendingFingerprint).toBe('fp2');
    expect(after?.hostKeyPendingSource).toBe('mismatch');
    expect(after?.hostKeyPendingToken).toBe(r2!.token);
  });

  it('setHostKeyVerified returns not_pending when no pending key is set', () => {
    const repo = createConnectionRepo(db);
    const c = repo.create(baseInput());
    expect(repo.setHostKeyVerified(c.id, 'any', 'any')).toBe('not_pending');
  });

  it('setHostKeyVerified returns stale_token when token does not match', () => {
    const repo = createConnectionRepo(db);
    const c = repo.create(baseInput());
    const b64 = wireFormatHostKeyB64('ssh-ed25519');
    repo.setHostKeyPendingWithToken(c.id, b64, 'fp', 'tofu_record');
    expect(repo.setHostKeyVerified(c.id, 'wrong-token', 'fp')).toBe('stale_token');
  });

  it('setHostKeyVerified returns fingerprint_mismatch when fingerprint differs', () => {
    const repo = createConnectionRepo(db);
    const c = repo.create(baseInput());
    const b64 = wireFormatHostKeyB64('ssh-ed25519');
    const { token } = repo.setHostKeyPendingWithToken(c.id, b64, 'fp', 'tofu_record')!;
    expect(repo.setHostKeyVerified(c.id, token, 'WRONG_fp')).toBe('fingerprint_mismatch');
  });

  it('setHostKeyVerified promotes pending → primary atomically', () => {
    const repo = createConnectionRepo(db);
    const c = repo.create(baseInput());
    const b64 = wireFormatHostKeyB64('ssh-ed25519');
    const { token } = repo.setHostKeyPendingWithToken(c.id, b64, 'fp', 'tofu_record')!;
    expect(repo.setHostKeyVerified(c.id, token, 'fp', '2026-05-12T13:00:00.000Z')).toBe(
      'verified',
    );
    const after = repo.resolveConnection(c.id);
    expect(after?.hostKeyType).toBe('ssh-ed25519');
    expect(after?.hostKeyB64).toBe(b64);
    expect(after?.hostKeyFingerprint).toBe('fp');
    expect(after?.hostKeyVerifiedAt).toBe('2026-05-12T13:00:00.000Z');
    expect(after?.hostKeyRecordedAt).toBe('2026-05-12T13:00:00.000Z');
    expect(after?.hostKeyPending).toBe(false);
    expect(after?.hostKeyPendingB64).toBeNull();
    expect(after?.hostKeyPendingToken).toBeNull();
    expect(after?.hostKeyPendingSource).toBeNull();
  });

  it('replaceHostKey behaves like verified (atomic same-check)', () => {
    const repo = createConnectionRepo(db);
    const c = repo.create(baseInput());
    // First establish a verified key
    const b64old = wireFormatHostKeyB64('ssh-rsa');
    const t1 = repo.setHostKeyPendingWithToken(c.id, b64old, 'fp_old', 'tofu_record')!.token;
    repo.setHostKeyVerified(c.id, t1, 'fp_old');
    // Now observe a new pending (mismatch)
    const b64new = wireFormatHostKeyB64('ssh-ed25519');
    const t2 = repo.setHostKeyPendingWithToken(c.id, b64new, 'fp_new', 'mismatch')!.token;
    // Token + fp mismatch fails
    expect(repo.replaceHostKey(c.id, t2, 'WRONG')).toBe('fingerprint_mismatch');
    // Token + fp match → promoted
    expect(repo.replaceHostKey(c.id, t2, 'fp_new')).toBe('verified');
    const after = repo.resolveConnection(c.id);
    expect(after?.hostKeyType).toBe('ssh-ed25519');
    expect(after?.hostKeyFingerprint).toBe('fp_new');
    expect(after?.hostKeyPending).toBe(false);
  });

  it('verify after pending was replaced returns stale_token for the OLD token', () => {
    const repo = createConnectionRepo(db);
    const c = repo.create(baseInput());
    const oldB64 = wireFormatHostKeyB64('ssh-ed25519');
    const { token: oldToken } = repo.setHostKeyPendingWithToken(
      c.id,
      oldB64,
      'fp_old',
      'tofu_record',
    )!;
    // A second observation replaces pending — old token is now stale
    const newB64 = wireFormatHostKeyB64('ssh-rsa');
    repo.setHostKeyPendingWithToken(c.id, newB64, 'fp_new', 'mismatch');
    // Old token should NOT verify
    expect(repo.setHostKeyVerified(c.id, oldToken, 'fp_old')).toBe('stale_token');
  });
});