import type { CryptoProvider } from './crypto.js';
import type { KeyPair, RatchetMessage, SessionState } from './types.js';
import {
  initSenderSession,
  initReceiverSession,
  ratchetEncrypt,
  ratchetDecrypt,
} from './ratchet.js';
import { StreamClosedError } from './errors.js';

/**
 * Streaming Double-Ratchet sub-sessions (V4.11).
 *
 * Wraps a long-lived, high-frequency, often-one-directional channel
 * (e.g. a server→client WebSocket log burst) in an *independent* Double
 * Ratchet that is derived from — but never mutates — an already
 * established parent Shade session.
 *
 * Why a sub-ratchet rather than reusing `ShadeSessionManager`:
 *
 *  - **Independence (R5).** A stream gets its own root key, chains, DH
 *    ratchet and op-mutex. Opening/closing it never touches the stored
 *    parent `SessionState` nor serialises against the HTTP send/receive
 *    queue.
 *  - **Performance (R7).** The stream ratchet lives only in memory and
 *    is *never* written to the keystore. There is therefore zero
 *    per-frame storage I/O — the cost is purely the symmetric KDF +
 *    AES-GCM, the same primitives the HTTP path uses.
 *  - **Forward secrecy.** Not persisting the evolving ratchet state is
 *    a feature, not a shortcut: writing per-frame secrets to disk would
 *    actively defeat the forward-secrecy guarantee. A dropped/reconnected
 *    stream is re-opened with a fresh handshake, not resumed.
 *
 * ## Seeding (no prekey-server round trip)
 *
 * The stream root key is derived from an identity-bound 3-DH exchange —
 * the X3DH pattern minus the signed / one-time prekeys, because the
 * peer's identity is *already* mutually pinned by the parent session's
 * TOFU. Two ephemeral keys are exchanged inside the transport itself
 * (`STREAM_OPEN` / `STREAM_OPEN_ACK`); no prekey server is involved.
 *
 *   slotA = DH(initiatorEphemeral, responderIdentity)   — auth of responder
 *   slotB = DH(initiatorIdentity,  responderEphemeral)   — auth of initiator
 *   slotC = DH(initiatorEphemeral, responderEphemeral)   — ephemeral FS
 *
 *   SK = HKDF(ikm = slotA‖slotB‖slotC, salt = streamId, info = "ShadeStream/v1")
 *
 * Both peers compute the identical three scalars regardless of role, so
 * `SK` agrees. An attacker lacking the responder's identity private key
 * cannot form slotA; one lacking the initiator's cannot form slotB —
 * the handshake is therefore mutually authenticated against the same
 * identities the parent session already trusts.
 *
 * `SK` then bootstraps a textbook Double Ratchet by handing the
 * responder's ephemeral to {@link initSenderSession} /
 * {@link initReceiverSession} exactly the way X3DH hands its signed
 * prekey to the ratchet — so `ratchetEncrypt` / `ratchetDecrypt` (and
 * thus every R1–R3 guarantee they already carry) apply unchanged.
 */

export type StreamRole = 'initiator' | 'responder';

/** Stream identifier length (bytes). 128 bits of collision resistance. */
export const STREAM_ID_BYTES = 16;

/** Ephemeral X25519 public-key length carried in the handshake. */
export const STREAM_EPHEMERAL_BYTES = 32;

const STREAM_KDF_INFO = new TextEncoder().encode('ShadeStream/v1');

/**
 * Derive the stream's independent root key from the identity-bound 3-DH
 * exchange. Pure: never reads or mutates any `SessionState`.
 *
 * @param role                 which end of the handshake we are
 * @param streamId             16-byte stream id (HKDF salt; binds the
 *                             derivation so two concurrent streams to the
 *                             same peer never share a root key)
 * @param selfIdentityDHPriv   our long-term identity X25519 private key
 * @param peerIdentityDHPub    peer's pinned identity X25519 public key
 *                             (the value the parent session pinned)
 * @param selfEphemeralPriv    our per-stream ephemeral X25519 private key
 * @param peerEphemeralPub     peer's per-stream ephemeral X25519 public key
 */
export async function deriveStreamRootKey(
  crypto: CryptoProvider,
  role: StreamRole,
  streamId: Uint8Array,
  selfIdentityDHPriv: Uint8Array,
  peerIdentityDHPub: Uint8Array,
  selfEphemeralPriv: Uint8Array,
  peerEphemeralPub: Uint8Array,
): Promise<Uint8Array> {
  // Each slot is pinned to a fixed semantic (not to local role) so both
  // sides feed HKDF the identical ikm:
  //   slotA = DH(initiatorEphemeral, responderIdentity)
  //   slotB = DH(initiatorIdentity,  responderEphemeral)
  //   slotC = DH(initiatorEphemeral, responderEphemeral)
  let slotA: Uint8Array;
  let slotB: Uint8Array;
  let slotC: Uint8Array;
  if (role === 'initiator') {
    slotA = await crypto.x25519(selfEphemeralPriv, peerIdentityDHPub);
    slotB = await crypto.x25519(selfIdentityDHPriv, peerEphemeralPub);
    slotC = await crypto.x25519(selfEphemeralPriv, peerEphemeralPub);
  } else {
    slotA = await crypto.x25519(selfIdentityDHPriv, peerEphemeralPub);
    slotB = await crypto.x25519(selfEphemeralPriv, peerIdentityDHPub);
    slotC = await crypto.x25519(selfEphemeralPriv, peerEphemeralPub);
  }

  const ikm = new Uint8Array(96);
  ikm.set(slotA, 0);
  ikm.set(slotB, 32);
  ikm.set(slotC, 64);
  const sk = await crypto.hkdf(ikm, streamId, STREAM_KDF_INFO, 32);

  crypto.zeroize(slotA);
  crypto.zeroize(slotB);
  crypto.zeroize(slotC);
  crypto.zeroize(ikm);
  return sk;
}

/**
 * Bootstrap a fresh Double Ratchet `SessionState` from the derived
 * stream root key. The responder's ephemeral plays exactly the role
 * X3DH's signed prekey plays in {@link initSenderSession} /
 * {@link initReceiverSession}, so the ratchet handoff is identical to
 * the proven HTTP path.
 *
 * On the initiator only `responderEphemeral.publicKey` is needed; the
 * responder must pass its full ephemeral keypair.
 *
 * `peerIdentityDHPub` is recorded as the session's `remoteIdentityKey`
 * so stream fingerprints stay meaningful and consistent with the parent.
 */
export async function bootstrapStreamSession(
  crypto: CryptoProvider,
  role: StreamRole,
  sk: Uint8Array,
  peerIdentityDHPub: Uint8Array,
  responderEphemeral: KeyPair,
): Promise<SessionState> {
  if (role === 'initiator') {
    // initSenderSession derives a fresh root via kdfRootKey and does not
    // retain `sk`, so the caller may safely zeroize it afterwards.
    return initSenderSession(crypto, sk, peerIdentityDHPub, responderEphemeral.publicKey);
  }
  // initReceiverSession stores the root key BY REFERENCE. Hand it an
  // independent copy so the caller zeroizing its `sk` scratch buffer
  // can't wipe the live session root.
  return initReceiverSession(new Uint8Array(sk), peerIdentityDHPub, responderEphemeral);
}

/** Zeroize every secret a stream session holds, then drop the chains. */
function zeroizeSession(crypto: CryptoProvider, s: SessionState): void {
  crypto.zeroize(s.rootKey);
  if (s.sendChain.chainKey.length > 0) crypto.zeroize(s.sendChain.chainKey);
  if (s.receiveChain && s.receiveChain.chainKey.length > 0) {
    crypto.zeroize(s.receiveChain.chainKey);
  }
  if (s.dhSend.privateKey.length > 0) crypto.zeroize(s.dhSend.privateKey);
  for (const mk of s.skippedKeys.values()) crypto.zeroize(mk);
  s.skippedKeys.clear();
}

/**
 * In-memory holder for a stream's Double Ratchet. Serialises its own
 * `seal`/`open`/`close` on a private promise chain (independent of the
 * SDK's per-peer encrypt/decrypt queues — R5) so per-frame ratchet
 * mutations never interleave, while staying fully concurrent with the
 * parent session and with other streams.
 *
 * Never persisted: the ratchet exists only for the lifetime of the
 * stream and is zeroized on `close()`.
 */
export class StreamRatchet {
  private session: SessionState | null;
  private opChain: Promise<unknown> = Promise.resolve();

  constructor(
    private readonly crypto: CryptoProvider,
    session: SessionState,
    /** 16-byte stream id this ratchet is bound to. */
    public readonly streamId: Uint8Array,
  ) {
    this.session = session;
  }

  /** True once {@link close} has run; `seal`/`open` will throw. */
  get closed(): boolean {
    return this.session === null;
  }

  private run<T>(fn: (s: SessionState) => Promise<T>): Promise<T> {
    const next = this.opChain.catch(() => undefined).then(() => {
      if (!this.session) throw new StreamClosedError();
      return fn(this.session);
    });
    // Keep a never-rejecting tail so a failed frame doesn't poison the
    // next one (a single bad inbound frame must not wedge the stream).
    this.opChain = next.catch(() => undefined);
    return next;
  }

  /** Wrap one logical frame. Advances the sending chain by one step. */
  seal(plaintext: Uint8Array): Promise<RatchetMessage> {
    return this.run((s) => ratchetEncrypt(this.crypto, s, plaintext));
  }

  /**
   * Unwrap one inbound frame. Correct and memory-bounded across long
   * one-directional runs from the peer: ordered transport delivery
   * skips zero keys per frame, and out-of-order arrivals are still
   * capped by the ratchet's `MAX_SKIP` / `MAX_CACHED_SKIPPED_KEYS`.
   */
  open(message: RatchetMessage): Promise<Uint8Array> {
    return this.run((s) => ratchetDecrypt(this.crypto, s, message));
  }

  /** Zeroize and drop the ratchet. Idempotent. */
  close(): Promise<void> {
    return this.opChain
      .catch(() => undefined)
      .then(() => {
        if (this.session) {
          zeroizeSession(this.crypto, this.session);
          this.session = null;
        }
      });
  }
}
