import { describe, it, expect, beforeEach, afterEach } from 'vitest';
import { promises as fs } from 'node:fs';
import * as path from 'node:path';
import { tmpdir } from 'node:os';
import {
  sshExec,
  sshUpload,
  sshDownload,
  SshSessionError,
  type ResolvedConnection,
  type SessionHooks,
} from './session.js';
import {
  startTestServer,
  generateEd25519Pair,
  generateRsaPair,
  type RunningTestServer,
} from './session-test-server.js';

/** Build a SessionHooks mock that records observation calls. */
function makeHooks(): SessionHooks & { observed: Array<{ kind: 'first' | 'mismatch'; b64: string; fingerprint: string }> } {
  const observed: Array<{ kind: 'first' | 'mismatch'; b64: string; fingerprint: string }> = [];
  return {
    observed,
    onFirstObserve: async (o) => {
      observed.push({ kind: 'first', b64: o.b64, fingerprint: o.fingerprint });
      return { token: 'first-token-xxx' };
    },
    onMismatch: async (o) => {
      observed.push({ kind: 'mismatch', b64: o.b64, fingerprint: o.fingerprint });
      return { token: 'mismatch-token-xxx' };
    },
  };
}

function buildConnection(args: {
  server: RunningTestServer;
  clientPem: Buffer;
  hostKeyVerified?: boolean;
  hostKeyB64?: string | null;
}): ResolvedConnection {
  // Use `'hostKeyB64' in args` so we honor an explicit null (vs. omitted).
  const hostKeyB64 = 'hostKeyB64' in args ? args.hostKeyB64 ?? null : args.server.hostKeyOpenSshB64;
  return {
    id: 'conn-test',
    ownerId: null,
    host: '127.0.0.1',
    port: args.server.port,
    username: 'testuser',
    privateKeyPem: args.clientPem,
    hostKeyB64,
    hostKeyVerified: args.hostKeyVerified ?? true,
    allowPrivate: true,
  };
}

describe('ssh/session sshExec', () => {
  let server: RunningTestServer;
  let clientPem: Buffer;

  beforeEach(async () => {
    server = await startTestServer();
    clientPem = generateRsaPair().privatePem;
  });

  afterEach(async () => {
    await server.close();
  });

  it('runs a simple command against a verified host', async () => {
    const conn = buildConnection({ server, clientPem });
    const r = await sshExec(
      { connection: conn, command: 'whoami', timeoutMs: 5000 },
      makeHooks(),
    );
    const parsed = JSON.parse(r.outputJson) as { stdout: string; exit_code: number; untrusted: boolean };
    expect(parsed.stdout).toContain('echo: whoami');
    expect(parsed.exit_code).toBe(0);
    expect(parsed.untrusted).toBe(true);
    expect(r.exitCode).toBe(0);
    expect(r.hostFingerprint).toBe(server.hostKeyFingerprint);
  });

  it('rejects without making a connection when hostKeyVerified=false', async () => {
    const conn = buildConnection({ server, clientPem, hostKeyVerified: false });
    const hooks = makeHooks();
    await expect(
      sshExec({ connection: conn, command: 'noop', timeoutMs: 5000 }, hooks),
    ).rejects.toMatchObject({ code: 'host_key_not_verified' });
    expect(hooks.observed).toHaveLength(0);
  });

  it('reports first_observe when no host key was recorded', async () => {
    const conn = buildConnection({ server, clientPem, hostKeyB64: null });
    const hooks = makeHooks();
    let caught: SshSessionError | null = null;
    try {
      await sshExec({ connection: conn, command: 'x', timeoutMs: 5000 }, hooks);
    } catch (e) {
      caught = e as SshSessionError;
    }
    expect(caught?.code).toBe('host_key_first_observe');
    expect(caught?.observedFingerprint).toBe(server.hostKeyFingerprint);
    expect(caught?.pendingToken).toBe('first-token-xxx');
    expect(hooks.observed[0]).toMatchObject({ kind: 'first', fingerprint: server.hostKeyFingerprint });
  });

  it('reports mismatch when the recorded host key differs from the observed one', async () => {
    // Build a connection whose recorded key is from a DIFFERENT server key.
    const otherKey = generateEd25519Pair().privatePem;
    const otherServer = await startTestServer({ hostKeyPem: otherKey });
    const decoyHostB64 = otherServer.hostKeyOpenSshB64;
    await otherServer.close();

    const conn = buildConnection({ server, clientPem, hostKeyB64: decoyHostB64 });
    const hooks = makeHooks();
    let caught: SshSessionError | null = null;
    try {
      await sshExec({ connection: conn, command: 'x', timeoutMs: 5000 }, hooks);
    } catch (e) {
      caught = e as SshSessionError;
    }
    expect(caught?.code).toBe('host_key_mismatch');
    expect(caught?.observedFingerprint).toBe(server.hostKeyFingerprint);
    expect(caught?.pendingToken).toBe('mismatch-token-xxx');
    expect(hooks.observed[0]).toMatchObject({ kind: 'mismatch' });
  });

  it('captures stderr and exit code', async () => {
    const server2 = await startTestServer({
      exec: async () => ({ stdout: 'OK\n', stderr: 'WARN!\n', exit: 7 }),
    });
    try {
      const conn = buildConnection({ server: server2, clientPem });
      const r = await sshExec({ connection: conn, command: 'x', timeoutMs: 5000 }, makeHooks());
      const parsed = JSON.parse(r.outputJson) as { stdout: string; stderr: string; exit_code: number };
      expect(parsed.stdout).toContain('OK');
      expect(parsed.stderr).toContain('WARN!');
      expect(parsed.exit_code).toBe(7);
    } finally {
      await server2.close();
    }
  });

  it('truncates stdout that exceeds maxOutputBytes', async () => {
    const big = 'A'.repeat(100 * 1024);
    const server2 = await startTestServer({
      exec: async () => ({ stdout: big, exit: 0 }),
    });
    try {
      const conn = buildConnection({ server: server2, clientPem });
      const r = await sshExec(
        { connection: conn, command: 'big', timeoutMs: 5000, maxOutputBytes: 1024 },
        makeHooks(),
      );
      const parsed = JSON.parse(r.outputJson) as { stdout: string; truncated: boolean; stdout_bytes_raw: number };
      expect(parsed.truncated).toBe(true);
      expect(parsed.stdout_bytes_raw).toBeGreaterThanOrEqual(100 * 1024);
      expect(parsed.stdout.length).toBeLessThan(2048);
    } finally {
      await server2.close();
    }
  });

  it('times out long-running exec', async () => {
    const server2 = await startTestServer({
      exec: async () => ({ stdout: 'late\n', exit: 0, delayMs: 600 }),
    });
    try {
      const conn = buildConnection({ server: server2, clientPem });
      await expect(
        sshExec({ connection: conn, command: 'slow', timeoutMs: 200 }, makeHooks()),
      ).rejects.toMatchObject({ code: 'exec_timeout' });
    } finally {
      await server2.close();
    }
  });
});

describe('ssh/session sshUpload', () => {
  let server: RunningTestServer;
  let clientPem: Buffer;
  let scratch: string;

  beforeEach(async () => {
    server = await startTestServer();
    clientPem = generateRsaPair().privatePem;
    scratch = await fs.realpath(await fs.mkdtemp(path.join(tmpdir(), 'ssh-up-')));
  });

  afterEach(async () => {
    await server.close();
    await fs.rm(scratch, { recursive: true, force: true });
  });

  it('uploads a small file', async () => {
    const local = path.join(scratch, 'hello.txt');
    await fs.writeFile(local, 'hello world');
    const conn = buildConnection({ server, clientPem });
    const r = await sshUpload(
      { connection: conn, localPath: local, remotePath: '/remote/hello.txt', timeoutMs: 5000, maxBytes: 1024 * 1024 },
      makeHooks(),
    );
    expect(r.bytes).toBe(11);
    expect(r.hostFingerprint).toBe(server.hostKeyFingerprint);
    expect(server.getFile('/remote/hello.txt')?.toString()).toBe('hello world');
  });

  it('rejects a local file that exceeds maxBytes', async () => {
    const local = path.join(scratch, 'big.bin');
    await fs.writeFile(local, Buffer.alloc(2048));
    const conn = buildConnection({ server, clientPem });
    await expect(
      sshUpload(
        { connection: conn, localPath: local, remotePath: '/remote/big.bin', timeoutMs: 5000, maxBytes: 1024 },
        makeHooks(),
      ),
    ).rejects.toMatchObject({ code: 'output_too_large' });
    expect(server.getFile('/remote/big.bin')).toBeUndefined();
  });

  it('rejects when the local leaf is a symlink (O_NOFOLLOW)', async () => {
    const real = path.join(scratch, 'real.txt');
    const link = path.join(scratch, 'link.txt');
    await fs.writeFile(real, 'inside');
    await fs.symlink(real, link);
    const conn = buildConnection({ server, clientPem });
    await expect(
      sshUpload(
        { connection: conn, localPath: link, remotePath: '/remote/x', timeoutMs: 5000, maxBytes: 1024 },
        makeHooks(),
      ),
    ).rejects.toMatchObject({ code: 'local_io_failed' });
  });

  it('rejects when the local path is a directory, not a regular file', async () => {
    const dir = path.join(scratch, 'subdir');
    await fs.mkdir(dir);
    const conn = buildConnection({ server, clientPem });
    await expect(
      sshUpload(
        { connection: conn, localPath: dir, remotePath: '/remote/x', timeoutMs: 5000, maxBytes: 1024 },
        makeHooks(),
      ),
    ).rejects.toMatchObject({ code: 'local_io_failed' });
  });
});

describe('ssh/session sshDownload', () => {
  let server: RunningTestServer;
  let clientPem: Buffer;
  let scratch: string;

  beforeEach(async () => {
    clientPem = generateRsaPair().privatePem;
    scratch = await fs.realpath(await fs.mkdtemp(path.join(tmpdir(), 'ssh-dl-')));
    server = await startTestServer({
      files: { '/remote/data.txt': Buffer.from('payload contents', 'utf-8') },
    });
  });

  afterEach(async () => {
    await server.close();
    await fs.rm(scratch, { recursive: true, force: true });
  });

  it('downloads a remote file to a new local path', async () => {
    const local = path.join(scratch, 'out.txt');
    const conn = buildConnection({ server, clientPem });
    const r = await sshDownload(
      { connection: conn, remotePath: '/remote/data.txt', localPath: local, timeoutMs: 5000, maxBytes: 1024 * 1024 },
      makeHooks(),
    );
    expect(r.bytes).toBe('payload contents'.length);
    expect(r.hostFingerprint).toBe(server.hostKeyFingerprint);
    const got = await fs.readFile(local, 'utf-8');
    expect(got).toBe('payload contents');
  });

  it('refuses to overwrite an existing local target', async () => {
    const local = path.join(scratch, 'out.txt');
    await fs.writeFile(local, 'old');
    const conn = buildConnection({ server, clientPem });
    await expect(
      sshDownload(
        { connection: conn, remotePath: '/remote/data.txt', localPath: local, timeoutMs: 5000, maxBytes: 1024 * 1024 },
        makeHooks(),
      ),
    ).rejects.toMatchObject({ code: 'local_target_exists' });
    // Existing file untouched.
    expect(await fs.readFile(local, 'utf-8')).toBe('old');
  });

  it('rejects when the remote file exceeds maxBytes', async () => {
    const local = path.join(scratch, 'out.txt');
    const conn = buildConnection({ server, clientPem });
    await expect(
      sshDownload(
        { connection: conn, remotePath: '/remote/data.txt', localPath: local, timeoutMs: 5000, maxBytes: 1 },
        makeHooks(),
      ),
    ).rejects.toMatchObject({ code: 'remote_too_large' });
    // Partial files cleaned up.
    const after = await fs.readdir(scratch);
    expect(after.filter((f) => f.endsWith('.partial') || f.includes('.partial-'))).toHaveLength(0);
    expect(after.filter((f) => f === 'out.txt')).toHaveLength(0);
  });

  it('rejects when the remote target does not exist', async () => {
    const local = path.join(scratch, 'out.txt');
    const conn = buildConnection({ server, clientPem });
    await expect(
      sshDownload(
        { connection: conn, remotePath: '/remote/missing.txt', localPath: local, timeoutMs: 5000, maxBytes: 1024 * 1024 },
        makeHooks(),
      ),
    ).rejects.toMatchObject({ code: 'remote_io_failed' });
    // No partial file remains.
    const after = await fs.readdir(scratch);
    expect(after.some((f) => f.startsWith('out.txt'))).toBe(false);
  });
});
