import { describe, expect, it } from 'vitest'
import type { AeadAlg, DirectionalKeys, HandshakeResult } from 'relay-contracts'
import { deriveSessionKeys } from '../src/hkdf.js'
import { createE2ESession, openFrame, sealFrame } from '../src/session.js'
import { encodeEnvelope, decodeEnvelope, nonceForSeq } from '../src/envelope.js'
import { AeadOpenError, ReplayError } from '../src/errors.js'
import { bytesToHex, fromUtf8, utf8 } from './helpers.js'

async function keysFor(alg: AeadAlg): Promise<DirectionalKeys> {
  return deriveSessionKeys(new Uint8Array(32).fill(0xcd), new Uint8Array(32).fill(1), new Uint8Array(32).fill(2), alg)
}

function resultFor(keys: DirectionalKeys, aead: AeadAlg): HandshakeResult {
  return { keys, aead, transcript: new Uint8Array([1]) }
}

const ALGS: AeadAlg[] = ['aes-256-gcm', 'xchacha20-poly1305']

describe('T9 session sealFrame/openFrame + E2ESession', () => {
  it.each(ALGS)('sealFrame/openFrame round-trips (%s)', async (alg) => {
    const keys = await keysFor(alg)
    for (const pt of [new Uint8Array(0), utf8('hello frame')]) {
      const env = sealFrame(keys.c2h, 0n, pt)
      expect(openFrame(keys.c2h, env, 0n)).toEqual(pt)
    }
  })

  it('deterministic nonce KAT: nonce = f(seq), identical across repeated seals of the same seq', async () => {
    const keys = await keysFor('aes-256-gcm')
    const a = sealFrame(keys.c2h, 7n, utf8('x'))
    const b = sealFrame(keys.c2h, 7n, utf8('x'))
    expect(bytesToHex(a.nonce)).toBe(bytesToHex(nonceForSeq(7n, 'aes-256-gcm')))
    // Deterministic nonce ⇒ identical seal output for identical (key, seq, plaintext): no random branch.
    expect(bytesToHex(a.nonce)).toBe(bytesToHex(b.nonce))
    expect(bytesToHex(a.ciphertext)).toBe(bytesToHex(b.ciphertext))
  })

  it('finding #1: a client-sealed c2h frame cannot be opened with the client read key h2c', async () => {
    const keys = await keysFor('aes-256-gcm')
    const client = createE2ESession('client', resultFor(keys, 'aes-256-gcm'))
    const wire = client.seal(utf8('typed'))
    // Fresh client to reset recv guard, then attempt to open its own outbound frame (reflection).
    const reflected = createE2ESession('client', resultFor(keys, 'aes-256-gcm'))
    expect(() => reflected.open(wire)).toThrow(AeadOpenError)
  })

  it('cross-instance: host opens what client sealed and vice-versa', async () => {
    const keys = await keysFor('xchacha20-poly1305')
    const client = createE2ESession('client', resultFor(keys, 'xchacha20-poly1305'))
    const host = createE2ESession('host', resultFor(keys, 'xchacha20-poly1305'))
    expect(fromUtf8(host.open(client.seal(utf8('c2h msg'))))).toBe('c2h msg')
    expect(fromUtf8(client.open(host.seal(utf8('h2c msg'))))).toBe('h2c msg')
  })

  it('INV13: replay, reorder, tamper, and seq-bump all rejected', async () => {
    const keys = await keysFor('aes-256-gcm')
    const client = createE2ESession('client', resultFor(keys, 'aes-256-gcm'))
    const host = createE2ESession('host', resultFor(keys, 'aes-256-gcm'))
    const w0 = client.seal(utf8('m0'))
    const w1 = client.seal(utf8('m1'))
    host.open(w0)
    expect(() => host.open(w0)).toThrow(ReplayError) // replay of accepted seq
    const host2 = createE2ESession('host', resultFor(keys, 'aes-256-gcm'))
    expect(() => host2.open(w1)).toThrow(ReplayError) // reorder: fresh host expects seq 0

    // tamper a ciphertext bit at the EXPECTED seq (so the recv guard passes and AEAD verifies)
    const c2 = createE2ESession('client', resultFor(keys, 'aes-256-gcm'))
    const h2 = createE2ESession('host', resultFor(keys, 'aes-256-gcm'))
    const env = decodeEnvelope(c2.seal(utf8('m2'))) // seq 0
    const bad = env.ciphertext.slice()
    bad[0]! ^= 0x01
    expect(() => h2.open(encodeEnvelope({ ...env, ciphertext: bad }))).toThrow(AeadOpenError)
  })

  it('openFrame: env.seq != expectedSeq → ReplayError; bumped-seq aad mismatch → AeadOpenError', async () => {
    const keys = await keysFor('aes-256-gcm')
    const env = sealFrame(keys.c2h, 3n, utf8('m'))
    expect(() => openFrame(keys.c2h, env, 4n)).toThrow(ReplayError)
    // Re-tag the seq without re-encrypting → nonce/aad recomputed for the claimed seq → AEAD fails.
    expect(() => openFrame(keys.c2h, { ...env, seq: 3n }, 3n)).not.toThrow()
  })

  it('lifecycle: rederive installs new keys, resets guards; old-key frames fail post-rederive', async () => {
    const keysA = await keysFor('aes-256-gcm')
    const keysB = await deriveSessionKeys(new Uint8Array(32).fill(0xee), new Uint8Array(32).fill(3), new Uint8Array(32).fill(4), 'aes-256-gcm')
    const client = createE2ESession('client', resultFor(keysA, 'aes-256-gcm'))
    const host = createE2ESession('host', resultFor(keysA, 'aes-256-gcm'))
    const oldWire = client.seal(utf8('old')) // seq 0 under keysA
    client.rederive(resultFor(keysB, 'aes-256-gcm'))
    host.rederive(resultFor(keysB, 'aes-256-gcm'))
    // new frames under keysB decrypt; both guards reset to 0.
    expect(fromUtf8(host.open(client.seal(utf8('new'))))).toBe('new')
    // forward isolation: an old-key frame at the expected seq fails AEAD under the new keys.
    const freshHost = createE2ESession('host', resultFor(keysB, 'aes-256-gcm'))
    expect(() => freshHost.open(oldWire)).toThrow(AeadOpenError)
  })
})