/**
 * T9 — session AEAD: frozen §4.4 `sealFrame`/`openFrame` (SYNCHRONOUS) + the stateful directional
 * `E2ESession`. Implements the frozen relay-contracts signatures verbatim via the T2 noble AEAD;
 * redefines none of them.
 *
 * Nonce discipline (INV13): the nonce is FULLY DETERMINISTIC — `nonceForSeq(seq, alg)` (no random
 * branch). Safe only because the direction split (T6) guarantees `seq` is unique per direction
 * subkey. `aad = directionLabel‖seq` binds direction + sequence into the tag, so a frame can't be
 * replayed under a different seq or reflected across directions.
 */
import type {
  AeadKey,
  E2EEnvelope,
  E2ESession,
  HandshakeResult,
  SessionRole,
} from 'relay-contracts'
import { writeUint64BE } from 'relay-contracts'
import { aeadOpen, aeadSeal } from './aead.js'
import { unwrapAeadKey } from './aead-key.js'
import { timingSafeEqual } from './crypto-provider.js'
import { decodeEnvelope, encodeEnvelope, nonceForSeq } from './envelope.js'
import { AeadOpenError, ReplayError } from './errors.js'
import { SequenceGuard } from './sequence.js'

const encoder = new TextEncoder()

/** aad = directionLabel‖seq(u64 BE) — binds both direction and sequence into the AEAD tag. */
function buildAad(label: string, seq: bigint): Uint8Array {
  const labelBytes = encoder.encode(label)
  const out = new Uint8Array(labelBytes.length + 8)
  out.set(labelBytes, 0)
  writeUint64BE(out, labelBytes.length, seq)
  return out
}

/** Frozen §4.4 seal: deterministic nonce = f(seq), aad = directionLabel‖seq. */
export function sealFrame(key: AeadKey, seq: bigint, plaintext: Uint8Array): E2EEnvelope {
  const { alg, aadLabel } = unwrapAeadKey(key)
  const nonce = nonceForSeq(seq, alg)
  const aad = buildAad(aadLabel, seq)
  const { ciphertext, tag } = aeadSeal(key, nonce, plaintext, aad)
  return { seq, nonce, ciphertext, tag }
}

/** Frozen §4.4 open: require env.seq === expectedSeq (INV13), recompute nonce, AEAD-verify. */
export function openFrame(key: AeadKey, env: E2EEnvelope, expectedSeq: bigint): Uint8Array {
  if (env.seq !== expectedSeq) {
    throw new ReplayError(`unexpected seq: expected ${expectedSeq}, got ${env.seq}`)
  }
  const { alg, aadLabel } = unwrapAeadKey(key)
  const nonce = nonceForSeq(expectedSeq, alg)
  if (!timingSafeEqual(nonce, env.nonce)) {
    throw new AeadOpenError('nonce does not match deterministic f(seq)')
  }
  const aad = buildAad(aadLabel, expectedSeq)
  return aeadOpen(key, nonce, env.ciphertext, env.tag, aad)
}

interface Directional {
  writeKey: AeadKey
  readKey: AeadKey
}

function directionFor(role: SessionRole, result: HandshakeResult): Directional {
  // c2h = client→host (client seals / host opens); h2c = host→client.
  return role === 'client'
    ? { writeKey: result.keys.c2h, readKey: result.keys.h2c }
    : { writeKey: result.keys.h2c, readKey: result.keys.c2h }
}

class E2ESessionImpl implements E2ESession {
  readonly role: SessionRole
  #dir: Directional
  #sendGuard = new SequenceGuard('send')
  #recvGuard = new SequenceGuard('recv')

  constructor(role: SessionRole, result: HandshakeResult) {
    this.role = role
    this.#dir = directionFor(role, result)
  }

  seal(plaintext: Uint8Array): Uint8Array {
    const seq = this.#sendGuard.next()
    return encodeEnvelope(sealFrame(this.#dir.writeKey, seq, plaintext))
  }

  open(dataPayload: Uint8Array): Uint8Array {
    const env = decodeEnvelope(dataPayload)
    this.#recvGuard.accept(env.seq)
    return openFrame(this.#dir.readKey, env, env.seq)
  }

  /** Re-key on refresh/reconnect: install new subkeys and reset BOTH seq guards to 0. */
  rederive(next: HandshakeResult): void {
    this.#dir = directionFor(this.role, next)
    this.#sendGuard = new SequenceGuard('send')
    this.#recvGuard = new SequenceGuard('recv')
  }
}

/** Frozen §4.4 factory — build a stateful directional session from a handshake result. */
export function createE2ESession(role: SessionRole, result: HandshakeResult): E2ESession {
  return new E2ESessionImpl(role, result)
}
