import type { CryptoProvider } from '@shade/core';
import { ValidationError, constantTimeEqual } from '@shade/core';
import { decodeStreamChunk } from '@shade/proto';
import { aesGcmDecryptWithNonce } from './aead.js';
import {
  StreamFinishedError,
  StreamOutOfOrderError,
  StreamProtocolError,
  StreamReplayError,
} from './errors.js';
import { StreamingSha256 } from './hash.js';
import { deriveLaneKey, deriveStreamKey } from './kdf.js';
import { buildChunkAad, buildChunkNonce, MAX_SEQ } from './nonce.js';
import type { DecryptedChunk } from './types.js';

export interface StreamReceiverInit {
  crypto: CryptoProvider;
  subtle?: SubtleCrypto;
  streamId: Uint8Array;
  streamSecret: Uint8Array;
  laneId: number;
  /** First seq this receiver will accept; defaults to 0. Used for resume. */
  startSeq?: number | bigint;
}

/**
 * Per-lane stream receiver state machine.
 *
 * Verifies AEAD, enforces strict in-order seq (rejects replay + out-of-order),
 * and updates a running lane sha256 over the decrypted plaintext.
 */
export class StreamReceiver {
  private constructor(
    private readonly subtle: SubtleCrypto | undefined,
    private readonly crypto: CryptoProvider,
    private readonly streamIdBytes: Uint8Array,
    public readonly laneId: number,
    private laneKey: Uint8Array | null,
    private expectedSeq: bigint,
    private readonly hasher: StreamingSha256,
    private finished: boolean,
    private bytesReceivedInternal: bigint,
  ) {}

  /** 16-byte streamId this receiver decrypts under. Defensive copy. */
  get streamId(): Uint8Array {
    return this.streamIdBytes.slice();
  }

  static async create(opts: StreamReceiverInit): Promise<StreamReceiver> {
    const streamKey = await deriveStreamKey(opts.crypto, opts.streamSecret, opts.streamId);
    try {
      const laneKey = await deriveLaneKey(opts.crypto, streamKey, opts.streamId, opts.laneId);
      const startSeq =
        opts.startSeq === undefined
          ? 0n
          : typeof opts.startSeq === 'bigint'
            ? opts.startSeq
            : BigInt(opts.startSeq);
      if (startSeq < 0n || startSeq > MAX_SEQ) {
        throw new ValidationError(`startSeq out of range: ${opts.startSeq}`, 'startSeq');
      }
      return new StreamReceiver(
        opts.subtle,
        opts.crypto,
        opts.streamId.slice(),
        opts.laneId,
        laneKey,
        startSeq,
        new StreamingSha256(),
        false,
        0n,
      );
    } finally {
      opts.crypto.zeroize(streamKey);
    }
  }

  /**
   * Decrypt and authenticate a wire-level chunk envelope. Throws on:
   *  - mismatched streamId / laneId
   *  - replayed seq (already accepted)
   *  - out-of-order seq (gap or backwards)
   *  - tampered nonce / ciphertext / aad
   *  - any chunk after `isLast`
   */
  async decryptChunk(wireBytes: Uint8Array): Promise<DecryptedChunk> {
    if (this.finished) {
      throw new StreamFinishedError('StreamReceiver: lane already finished');
    }
    if (this.laneKey === null) {
      throw new StreamFinishedError('StreamReceiver: destroyed');
    }

    const env = decodeStreamChunk(wireBytes);

    if (!constantTimeEqual(env.streamId, this.streamIdBytes)) {
      throw new StreamProtocolError('stream-chunk streamId does not match this receiver');
    }
    if (env.laneId !== this.laneId) {
      throw new StreamProtocolError(
        `stream-chunk laneId=${env.laneId} routed to laneId=${this.laneId}`,
      );
    }

    const seqBig = typeof env.seq === 'bigint' ? env.seq : BigInt(env.seq);
    if (seqBig < this.expectedSeq) {
      throw new StreamReplayError(`Replay: seq=${seqBig}, already accepted < ${this.expectedSeq}`);
    }
    if (seqBig > this.expectedSeq) {
      throw new StreamOutOfOrderError(Number(this.expectedSeq), Number(seqBig));
    }

    // Defense-in-depth: the wire nonce MUST equal the deterministically derived one.
    const expectedNonce = buildChunkNonce(this.laneId, seqBig);
    if (!constantTimeEqual(env.nonce, expectedNonce)) {
      throw new StreamProtocolError('stream-chunk nonce mismatch');
    }
    if (env.aad.length !== 0) {
      throw new StreamProtocolError('stream-chunk aad must be empty in v0.2.0');
    }

    const aad = buildChunkAad(this.streamIdBytes, this.laneId, seqBig, env.isLast);
    const plaintext = await aesGcmDecryptWithNonce(
      this.laneKey,
      env.nonce,
      env.ciphertext,
      aad,
      this.subtle,
    );

    this.hasher.update(plaintext);
    this.bytesReceivedInternal += BigInt(plaintext.length);
    this.expectedSeq = seqBig + 1n;

    if (env.isLast) this.finished = true;

    return { plaintext, seq: Number(seqBig), isLast: env.isLast };
  }

  /** Snapshot the lane sha256 digest. Hasher is frozen after this call. */
  getLaneSha256Digest(): Uint8Array {
    return this.hasher.digest();
  }

  /** Total plaintext bytes accepted so far in this lane. */
  get bytesReceived(): bigint {
    return this.bytesReceivedInternal;
  }

  /** Next sequence number this receiver expects. */
  get nextExpectedSequence(): bigint {
    return this.expectedSeq;
  }

  /** Has this lane received its `isLast` chunk? */
  get isFinished(): boolean {
    return this.finished;
  }

  /** Zero the lane key in memory. After destroy, decrypt calls throw. */
  destroy(): void {
    if (this.laneKey !== null) {
      this.crypto.zeroize(this.laneKey);
      this.laneKey = null;
    }
    this.finished = true;
  }
}
