import type { CryptoProvider } from '@shade/core';
import { ValidationError } from '@shade/core';
import { encodeStreamChunk } from '@shade/proto';
import { aesGcmEncryptWithNonce } from './aead.js';
import { StreamFinishedError } from './errors.js';
import { StreamingSha256 } from './hash.js';
import { deriveLaneKey, deriveStreamKey } from './kdf.js';
import { buildChunkAad, buildChunkNonce, MAX_SEQ } from './nonce.js';
import type { EncryptedChunk, StreamChunkEnvelope } from './types.js';

export interface StreamSenderInit {
  crypto: CryptoProvider;
  /** Optional SubtleCrypto for AEAD; defaults to globalThis.crypto.subtle. */
  subtle?: SubtleCrypto;
  streamId: Uint8Array;
  /** 32-byte streamSecret. The lane key is derived internally and stays in this instance. */
  streamSecret: Uint8Array;
  laneId: number;
  /** First seq this sender will emit; defaults to 0. Used for resume. */
  startSeq?: number | bigint;
}

/**
 * Per-lane stream sender state machine.
 *
 * Encapsulates the lane-specific AEAD key + monotonic seq counter + running
 * lane-sha256. Multiple senders (one per lane) sharing the same
 * `(streamSecret, streamId)` make up a parallel transfer.
 */
export class StreamSender {
  private constructor(
    private readonly subtle: SubtleCrypto | undefined,
    private readonly crypto: CryptoProvider,
    private readonly streamIdBytes: Uint8Array,
    public readonly laneId: number,
    private laneKey: Uint8Array | null,
    private nextSeq: bigint,
    private readonly hasher: StreamingSha256,
    private finished: boolean,
    private bytesSentInternal: bigint,
  ) {}

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

  static async create(opts: StreamSenderInit): Promise<StreamSender> {
    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 StreamSender(
        opts.subtle,
        opts.crypto,
        opts.streamId.slice(),
        opts.laneId,
        laneKey,
        startSeq,
        new StreamingSha256(),
        false,
        0n,
      );
    } finally {
      opts.crypto.zeroize(streamKey);
    }
  }

  /**
   * Encrypt one plaintext chunk. Updates lane sha256 with the plaintext bytes
   * BEFORE encryption (so receiver can independently verify by hashing
   * decrypted plaintext in the same order).
   */
  async encryptChunk(plaintext: Uint8Array, isLast: boolean): Promise<EncryptedChunk> {
    if (this.finished) {
      throw new StreamFinishedError('StreamSender: lane already finished');
    }
    if (this.laneKey === null) {
      throw new StreamFinishedError('StreamSender: destroyed');
    }
    if (this.nextSeq > MAX_SEQ) {
      throw new ValidationError('seq overflow', 'seq');
    }

    const seq = this.nextSeq;
    const nonce = buildChunkNonce(this.laneId, seq);
    const aad = buildChunkAad(this.streamIdBytes, this.laneId, seq, isLast);
    const ciphertext = await aesGcmEncryptWithNonce(
      this.laneKey,
      nonce,
      plaintext,
      aad,
      this.subtle,
    );

    this.hasher.update(plaintext);
    this.bytesSentInternal += BigInt(plaintext.length);
    this.nextSeq = seq + 1n;

    const envelope: StreamChunkEnvelope = {
      streamId: this.streamIdBytes,
      laneId: this.laneId,
      seq,
      isLast,
      nonce,
      aad: new Uint8Array(0),
      ciphertext,
    };
    const bytes = encodeStreamChunk(envelope);

    if (isLast) this.finished = true;

    return { envelope, bytes, seq: Number(seq) };
  }

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

  /**
   * Feed plaintext into the lane sha256 WITHOUT advancing seq or
   * producing wire bytes. Used by resume flows to re-build the running
   * lane hash from already-shipped bytes (since `@noble/hashes` v2 doesn't
   * expose serializable state in v0.2.0).
   */
  preHash(plaintext: Uint8Array): void {
    if (this.finished) {
      throw new StreamFinishedError('StreamSender: lane already finished');
    }
    this.hasher.update(plaintext);
  }

  /** Number of plaintext bytes encrypted so far in this lane. */
  get bytesSent(): bigint {
    return this.bytesSentInternal;
  }

  /** Next sequence number this sender will emit. */
  get nextSequence(): bigint {
    return this.nextSeq;
  }

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

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