import { describe, expect, it, vi } from 'vitest'
import { readFileSync } from 'node:fs'
import { join } from 'node:path'
import type {
  AeadAlg,
  ClientHello,
  E2ESession,
  HandshakeResult,
  HostHello,
} from 'relay-contracts'
import { generateIdentity } from '../src/keys/identity.js'
import {
  MitmAbortError,
  createE2ETransform,
  makeHostHello,
  type CreateHostHandshake,
  type VerifyDeviceProof,
} from '../src/e2e/hostEndpoint.js'
import type { ReplaySealer } from '../src/e2e/replaySeal.js'

const ALG: AeadAlg = 'aes-256-gcm'

function clientHello(proof: string): ClientHello {
  return {
    clientEphPub: new Uint8Array([1, 2, 3]),
    clientNonce: new Uint8Array([4, 5, 6]),
    aeadOffer: [ALG],
    deviceAuthProof: proof,
  }
}

function fakeHandshake(keysByte: number): CreateHostHandshake {
  return () => ({
    async respond(): Promise<{ hello: HostHello; result: HandshakeResult }> {
      const result: HandshakeResult = {
        keys: {
          c2h: new Uint8Array([keysByte]) as never,
          h2c: new Uint8Array([keysByte + 1]) as never,
        },
        aead: ALG,
        transcript: new Uint8Array([0xff]),
      }
      const hello: HostHello = {
        hostEphPub: new Uint8Array([7]),
        hostNonce: new Uint8Array([8]),
        aeadChoice: ALG,
        enrollFpr: 'fpr',
        sig: new Uint8Array([9]),
      }
      return { hello, result }
    },
  })
}

// The load-bearing MITM verifier: only a specific proof passes.
const realVerifier: VerifyDeviceProof = async (proof) => proof === 'VALID-PROOF'

describe('makeHostHello (T15, anti-MITM)', () => {
  it('derives DirectionalKeys{c2h,h2c} on a valid proof (FIX 2 — no single sessionKey)', async () => {
    const { result } = await makeHostHello(clientHello('VALID-PROOF'), generateIdentity(), realVerifier, {
      createHostHandshake: fakeHandshake(0x10),
    })
    expect(result.keys).toHaveProperty('c2h')
    expect(result.keys).toHaveProperty('h2c')
    expect(result).not.toHaveProperty('sessionKey')
  })

  it('ABORTS on a forged proof: no keys, no HostHello (MitmAbortError)', async () => {
    const respond = vi.fn()
    const spyHandshake: CreateHostHandshake = () => ({ respond: respond as never })
    await expect(
      makeHostHello(clientHello('FORGED'), generateIdentity(), realVerifier, {
        createHostHandshake: spyHandshake,
      }),
    ).rejects.toBeInstanceOf(MitmAbortError)
    expect(respond).not.toHaveBeenCalled() // no key derivation reached
  })

  it('no-stub guard: swapping the verifier for always-true makes the MITM test FAIL', async () => {
    const stub: VerifyDeviceProof = async () => true
    // With the stubbed verifier, the forged proof WRONGLY completes — proving the verifier is
    // load-bearing (a real build forbids this stub via the import guard below).
    const out = await makeHostHello(clientHello('FORGED'), generateIdentity(), stub, {
      createHostHandshake: fakeHandshake(0x20),
    })
    expect(out.result.keys).toBeDefined()
  })

  it('no-stub CI guard: hostEndpoint.ts imports NO verifier from relay-e2e (FIX 6b)', () => {
    const raw = readFileSync(join(import.meta.dirname, '..', 'src', 'e2e', 'hostEndpoint.ts'), 'utf8')
    const src = raw.replace(/\/\*[\s\S]*?\*\//g, '').replace(/\/\/.*$/gm, '') // strip comments
    expect(src).not.toMatch(/import[^\n]*verifyDeviceAuthProof/)
    expect(src).not.toMatch(/from ['"]relay-e2e['"]/)
  })
})

describe('createE2ETransform (T15)', () => {
  function fakeSession(): E2ESession {
    // Reversible transform that HIDES the plaintext (XOR) so INV2 assertions are meaningful.
    return {
      role: 'host',
      seal: (pt) => Uint8Array.from([0xe0, ...pt.map((b) => b ^ 0x55)]),
      open: (ct) => Uint8Array.from(ct.slice(1)).map((b) => b ^ 0x55),
      rederive: () => {},
    }
  }
  function fakeReplay(): ReplaySealer & { calls: number } {
    const r = {
      calls: 0,
      epoch: 'test-epoch',
      seal(_pt: Uint8Array) {
        r.calls += 1
        return { seq: 0n, nonce: new Uint8Array(), ciphertext: new Uint8Array([0xde]), tag: new Uint8Array() }
      },
    }
    return r
  }

  it('outbound seals under BOTH the live session and the replay key (FIX 3, INV2)', () => {
    const replay = fakeReplay()
    const t = createE2ETransform(generateIdentity(), realVerifier, replay)
    t.openStream(1)
    t.seedSession(1, fakeSession(), new Uint8Array([0x48])) // HostHello control frame
    const marker = new TextEncoder().encode('PLAIN')
    const sealed = t.outbound(1, marker)
    expect(replay.calls).toBe(1) // replay-bound seal happened
    expect(Buffer.from(sealed).includes(Buffer.from(marker))).toBe(false) // ciphertext, not plaintext
    expect(t.takeControlFrames!(1)).toEqual([new Uint8Array([0x48])]) // HostHello flushed once
    expect(t.takeControlFrames!(1)).toEqual([])
  })

  it('inbound before seeding returns null (handshake pending); after, opens opaque', () => {
    const t = createE2ETransform(generateIdentity(), realVerifier, fakeReplay())
    const session = fakeSession()
    t.openStream(1)
    expect(t.inbound(1, new Uint8Array([1, 2]))).toBeNull()
    t.seedSession(1, session, new Uint8Array([0x48]))
    const ct = session.seal(new Uint8Array([9, 9])) // c2h ciphertext
    expect([...t.inbound(1, ct)!]).toEqual([9, 9])
  })

  it('outbound before the session is established aborts (no silent plaintext leak)', () => {
    const t = createE2ETransform(generateIdentity(), realVerifier, fakeReplay())
    t.openStream(1)
    expect(() => t.outbound(1, new Uint8Array([1]))).toThrow(MitmAbortError)
  })
})
