import { describe, expect, it } from 'vitest'
import type { ClientHello, HostHello } from 'relay-contracts'
import { createClientHandshake, createHostHandshake } from '../src/handshake.js'
import { nobleEd25519Signer, nobleEd25519Verifier } from '../src/ed25519.js'
import { computeEnrollFpr } from '../src/fingerprint.js'
import { MemoryDevicePinStore } from '../src/keystore.js'
import { unwrapAeadKey } from '../src/aead-key.js'
import { FingerprintMismatchError, HandshakeStateError } from '../src/errors.js'
import { boundProofProvider, makeHost, makeHostIdentity, verifyBoundProof, bytesToHex } from './helpers.js'

function makeClient(hostId: string, pinStore = new MemoryDevicePinStore()) {
  return createClientHandshake({
    aeadOffer: ['xchacha20-poly1305', 'aes-256-gcm'],
    deviceAuthProofProvider: boundProofProvider(),
    verifier: nobleEd25519Verifier(),
    pinStore,
    hostId,
  })
}

describe('T8 handshake', () => {
  it('happy path: both sides derive the SAME DirectionalKeys + aead; relay sees no key', async () => {
    const id = makeHostIdentity()
    const client = makeClient('h1')
    const host = makeHost(id.privateKey, id.agentPubkey)
    const ch = await client.start()
    const hh = await host.onClientHello(ch)
    const rc = await client.onHostHello(hh, id.agentPubkey)
    const rh = host.result!
    expect(rc.aead).toBe('xchacha20-poly1305')
    expect(rc.aead).toBe(rh.aead)
    expect(bytesToHex(unwrapAeadKey(rc.keys.c2h).raw)).toBe(bytesToHex(unwrapAeadKey(rh.keys.c2h).raw))
    expect(bytesToHex(unwrapAeadKey(rc.keys.h2c).raw)).toBe(bytesToHex(unwrapAeadKey(rh.keys.h2c).raw))
    // The relay only ever sees the ClientHello/HostHello structs — no secret field present.
    expect(Object.keys(ch)).toEqual(['clientEphPub', 'clientNonce', 'aeadOffer', 'deviceAuthProof'])
    expect(Object.keys(hh)).toEqual(['hostEphPub', 'hostNonce', 'aeadChoice', 'enrollFpr', 'sig'])
  })

  it('AEAD negotiation: strongest common alg; empty intersection → HandshakeStateError', async () => {
    const id = makeHostIdentity()
    const client = createClientHandshake({
      aeadOffer: ['aes-256-gcm'],
      deviceAuthProofProvider: boundProofProvider(),
      verifier: nobleEd25519Verifier(),
      pinStore: new MemoryDevicePinStore(),
      hostId: 'h1',
    })
    const host = makeHost(id.privateKey, id.agentPubkey, ['aes-256-gcm'])
    const hh = await host.onClientHello(await client.start())
    expect(hh.aeadChoice).toBe('aes-256-gcm')

    const client2 = createClientHandshake({
      aeadOffer: ['xchacha20-poly1305'],
      deviceAuthProofProvider: boundProofProvider(),
      verifier: nobleEd25519Verifier(),
      pinStore: new MemoryDevicePinStore(),
      hostId: 'h1',
    })
    const hostGcmOnly = makeHost(id.privateKey, id.agentPubkey, ['aes-256-gcm'])
    await expect(hostGcmOnly.onClientHello(await client2.start())).rejects.toBeInstanceOf(
      HandshakeStateError,
    )
  })

  it('MITM abort: relay swaps hostEph/sig → sig fails vs registry pubkey → FingerprintMismatchError, no key', async () => {
    const id = makeHostIdentity()
    const evil = makeHostIdentity()
    const client = makeClient('h1')
    const host = makeHost(id.privateKey, id.agentPubkey)
    const ch = await client.start()
    const hh = await host.onClientHello(ch)
    // Relay forges the signature with its own key while keeping the honest enrollFpr.
    const forged: HostHello = { ...hh, sig: await nobleForge(evil.privateKey, hh) }
    await expect(client.onHostHello(forged, id.agentPubkey)).rejects.toBeInstanceOf(
      FingerprintMismatchError,
    )
    expect(client.phase).toBe('aborted')
  })

  it('finding #3: enrollFpr matches pin but supplied pubkey hashes differently → abort, no string-only match', async () => {
    const id = makeHostIdentity()
    const other = makeHostIdentity()
    const client = makeClient('h1')
    const host = makeHost(id.privateKey, id.agentPubkey)
    const hh = await host.onClientHello(await client.start())
    // hh.enrollFpr is the honest host's; but we hand onHostHello a DIFFERENT registry pubkey.
    await expect(client.onHostHello(hh, other.agentPubkey)).rejects.toBeInstanceOf(
      FingerprintMismatchError,
    )
  })

  it('TOFU first use records the recomputed pin; a later different registry pubkey aborts (no auto-repin)', async () => {
    const id = makeHostIdentity()
    const pinStore = new MemoryDevicePinStore()
    const client = makeClient('h1', pinStore)
    const host = makeHost(id.privateKey, id.agentPubkey)
    await client.onHostHello(await host.onClientHello(await client.start()), id.agentPubkey)
    expect(await pinStore.get('h1')).toBe(computeEnrollFpr(id.agentPubkey))

    const rotated = makeHostIdentity()
    const client2 = makeClient('h1', pinStore)
    const host2 = makeHost(rotated.privateKey, rotated.agentPubkey)
    await expect(
      client2.onHostHello(await host2.onClientHello(await client2.start()), rotated.agentPubkey),
    ).rejects.toBeInstanceOf(FingerprintMismatchError)
    // pin untouched
    expect(await pinStore.get('h1')).toBe(computeEnrollFpr(id.agentPubkey))
  })

  it('finding #2: a proof captured from handshake A cannot be replayed into handshake B', async () => {
    const id = makeHostIdentity()
    const clientA = makeClient('h1')
    const chA = await clientA.start()
    const clientB = makeClient('h1')
    const chB = await clientB.start()
    // Splice A's proof into B's (different clientEphPub) client_hello.
    const spliced: ClientHello = { ...chB, deviceAuthProof: chA.deviceAuthProof }
    const host = makeHost(id.privateKey, id.agentPubkey)
    await expect(host.onClientHello(spliced)).rejects.toBeInstanceOf(HandshakeStateError)
    // Sanity: the unbound proof really was bound to A, not B.
    expect(verifyBoundProof(chA.deviceAuthProof, chB)).toBe(false)
  })

  it('deviceAuthProof gate (INV3): forged/absent proof → HandshakeStateError, no session', async () => {
    const id = makeHostIdentity()
    const client = makeClient('h1')
    const ch = await client.start()
    const host = makeHost(id.privateKey, id.agentPubkey)
    await expect(host.onClientHello({ ...ch, deviceAuthProof: 'bogus' })).rejects.toBeInstanceOf(
      HandshakeStateError,
    )
    expect(host.result).toBeNull()
  })

  it('state machine: out-of-order / double calls → HandshakeStateError; phases terminal', async () => {
    const id = makeHostIdentity()
    const client = makeClient('h1')
    const host = makeHost(id.privateKey, id.agentPubkey)
    await expect(client.onHostHello({} as HostHello, id.agentPubkey)).rejects.toBeInstanceOf(
      HandshakeStateError,
    )
    const ch = await client.start()
    await expect(client.start()).rejects.toBeInstanceOf(HandshakeStateError)
    await host.onClientHello(ch)
    await expect(host.onClientHello(ch)).rejects.toBeInstanceOf(HandshakeStateError)
    expect(host.phase).toBe('established')
  })

  it('malformed input: bad HostHello shape → parse error before any crypto', async () => {
    const id = makeHostIdentity()
    const client = makeClient('h1')
    await client.start()
    await expect(
      client.onHostHello({ hostEphPub: 'nope' } as unknown as HostHello, id.agentPubkey),
    ).rejects.toBeTruthy()
  })
})

// Forge a signature over the honest transcript using a different key (relay MITM attempt).
import { buildTranscript } from '../src/transcript.js'
async function nobleForge(evilPriv: Uint8Array, hh: HostHello): Promise<Uint8Array> {
  // The forger cannot know the real transcript's client fields here; sign an arbitrary transcript.
  // What matters for the test is that the sig does NOT verify against the honest agentPubkey.
  const t = buildTranscript({
    clientEphPub: new Uint8Array(32),
    clientNonce: new Uint8Array(32),
    aeadOffer: ['xchacha20-poly1305'],
    hostEphPub: hh.hostEphPub,
    hostNonce: hh.hostNonce,
    aeadChoice: hh.aeadChoice,
    enrollFpr: hh.enrollFpr,
  })
  return nobleEd25519Signer(evilPriv).sign(t)
}