import type { CryptoProvider } from '@shade/core';
import { ValidationError } from '@shade/core';
import { decodeStreamChunk } from '@shade/proto';
import { StreamFinishedError, StreamProtocolError } from './errors.js';
import { StreamingSha256 } from './hash.js';
import { StreamReceiver } from './receiver.js';
import { StreamSender } from './sender.js';
import type { DecryptedChunk, EncryptedChunk, LaneInitSpec } from './types.js';

export interface MultiLaneSenderInit {
  crypto: CryptoProvider;
  subtle?: SubtleCrypto;
  streamId: Uint8Array;
  streamSecret: Uint8Array;
  lanes: LaneInitSpec[];
  /** Per-lane resume offsets; defaults to 0 for each lane. */
  startSeqByLane?: Map<number, number | bigint>;
}

export interface MultiLaneReceiverInit extends MultiLaneSenderInit {}

/** Per-lane sha256 fingerprint for the stream-finish envelope. */
export interface LaneFingerprint {
  laneId: number;
  /** sha256 over plaintext bytes carried by this lane in seq order. */
  sha256: Uint8Array;
  chunkCount: number;
  byteCount: number;
}

/**
 * Multi-lane stream sender.
 *
 * Wraps N independent `StreamSender`s sharing the same (streamSecret, streamId)
 * and tracks an `overallSha256` over the original plaintext as the consumer
 * dispatches bytes to lanes.
 *
 * The coordinator is partition-AGNOSTIC — the consumer (transfer layer)
 * decides which bytes go to which lane. The coordinator only enforces:
 *   - lane lookup by laneId
 *   - update of overallSha256 from `appendOverall(...)` (called BEFORE bytes
 *     are dispatched to lanes, in original byte order)
 */
export class MultiLaneSender {
  private readonly senders = new Map<number, StreamSender>();
  private readonly overallHasher = new StreamingSha256();
  private destroyed = false;

  private constructor(senders: StreamSender[]) {
    for (const s of senders) {
      this.senders.set(s.laneId, s);
    }
  }

  static async create(opts: MultiLaneSenderInit): Promise<MultiLaneSender> {
    if (opts.lanes.length === 0) {
      throw new ValidationError('at least one lane is required', 'lanes');
    }
    const seen = new Set<number>();
    for (const l of opts.lanes) {
      if (seen.has(l.laneId)) {
        throw new ValidationError(`duplicate laneId ${l.laneId}`, 'lanes');
      }
      seen.add(l.laneId);
    }
    const senders: StreamSender[] = [];
    for (const lane of opts.lanes) {
      const startSeq = opts.startSeqByLane?.get(lane.laneId);
      senders.push(
        await StreamSender.create({
          crypto: opts.crypto,
          ...(opts.subtle !== undefined ? { subtle: opts.subtle } : {}),
          streamId: opts.streamId,
          streamSecret: opts.streamSecret,
          laneId: lane.laneId,
          ...(startSeq !== undefined ? { startSeq } : {}),
        }),
      );
    }
    return new MultiLaneSender(senders);
  }

  /**
   * Update the overall (cross-lane, original-order) sha256 hasher.
   *
   * The consumer must call this once with each plaintext slice IN ORIGINAL
   * BYTE ORDER (regardless of which lane the slice ends up on). For range
   * partitioning this is straightforward (lanes are contiguous); for
   * round-robin the consumer hashes each chunk in the order it was read
   * from the input.
   */
  appendOverall(plaintext: Uint8Array): void {
    if (this.destroyed) throw new StreamFinishedError('MultiLaneSender: destroyed');
    this.overallHasher.update(plaintext);
  }

  /** Encrypt `plaintext` for the given lane. */
  async encryptForLane(
    laneId: number,
    plaintext: Uint8Array,
    isLast: boolean,
  ): Promise<EncryptedChunk> {
    const lane = this.senders.get(laneId);
    if (lane === undefined) {
      throw new StreamProtocolError(`Unknown laneId ${laneId}`);
    }
    return lane.encryptChunk(plaintext, isLast);
  }

  /** Resume helper: feed plaintext into a lane's hasher without advancing seq. */
  preHashForLane(laneId: number, plaintext: Uint8Array): void {
    const lane = this.senders.get(laneId);
    if (lane === undefined) {
      throw new StreamProtocolError(`Unknown laneId ${laneId}`);
    }
    lane.preHash(plaintext);
  }

  /** Snapshot per-lane fingerprints (call once when all lanes are finished). */
  getLaneFingerprints(): LaneFingerprint[] {
    const out: LaneFingerprint[] = [];
    for (const [laneId, sender] of this.senders) {
      out.push({
        laneId,
        sha256: sender.getLaneSha256Digest(),
        chunkCount: Number(sender.nextSequence),
        byteCount: Number(sender.bytesSent),
      });
    }
    out.sort((a, b) => a.laneId - b.laneId);
    return out;
  }

  /** Snapshot overall sha256 over original-order plaintext. */
  getOverallSha256(): Uint8Array {
    return this.overallHasher.digest();
  }

  /** All lanes have emitted their `isLast` chunks. */
  get allLanesFinished(): boolean {
    for (const lane of this.senders.values()) {
      if (!lane.isFinished) return false;
    }
    return true;
  }

  destroy(): void {
    if (this.destroyed) return;
    for (const sender of this.senders.values()) sender.destroy();
    this.destroyed = true;
  }
}

/**
 * Multi-lane stream receiver.
 *
 * Routes incoming wire envelopes to per-lane `StreamReceiver`s by laneId,
 * and exposes a hook for the consumer to update the overall sha256 in
 * original byte order (the consumer's reorder logic feeds plaintext here
 * after collecting it from individual lanes).
 */
export class MultiLaneReceiver {
  private readonly receivers = new Map<number, StreamReceiver>();
  private readonly overallHasher = new StreamingSha256();
  private destroyed = false;

  private constructor(receivers: StreamReceiver[]) {
    for (const r of receivers) {
      this.receivers.set(r.laneId, r);
    }
  }

  static async create(opts: MultiLaneReceiverInit): Promise<MultiLaneReceiver> {
    if (opts.lanes.length === 0) {
      throw new ValidationError('at least one lane is required', 'lanes');
    }
    const seen = new Set<number>();
    for (const l of opts.lanes) {
      if (seen.has(l.laneId)) {
        throw new ValidationError(`duplicate laneId ${l.laneId}`, 'lanes');
      }
      seen.add(l.laneId);
    }
    const receivers: StreamReceiver[] = [];
    for (const lane of opts.lanes) {
      const startSeq = opts.startSeqByLane?.get(lane.laneId);
      receivers.push(
        await StreamReceiver.create({
          crypto: opts.crypto,
          ...(opts.subtle !== undefined ? { subtle: opts.subtle } : {}),
          streamId: opts.streamId,
          streamSecret: opts.streamSecret,
          laneId: lane.laneId,
          ...(startSeq !== undefined ? { startSeq } : {}),
        }),
      );
    }
    return new MultiLaneReceiver(receivers);
  }

  /**
   * Decrypt one wire-level chunk. Routes to the lane indicated by the
   * envelope's laneId. Throws if no receiver exists for that laneId.
   */
  async decryptChunk(wireBytes: Uint8Array): Promise<DecryptedChunk & { laneId: number }> {
    if (this.destroyed) throw new StreamFinishedError('MultiLaneReceiver: destroyed');
    const env = decodeStreamChunk(wireBytes);
    const lane = this.receivers.get(env.laneId);
    if (lane === undefined) {
      throw new StreamProtocolError(`Unknown laneId ${env.laneId}`);
    }
    const dec = await lane.decryptChunk(wireBytes);
    return { ...dec, laneId: env.laneId };
  }

  /**
   * Update the overall sha256 hasher with plaintext IN ORIGINAL BYTE ORDER.
   * The consumer's reorder logic decides when to call this (as bytes leave
   * the reorder buffer in original order, OR for range-partitioning,
   * concatenated lane outputs in laneId order).
   */
  appendOverall(plaintext: Uint8Array): void {
    if (this.destroyed) throw new StreamFinishedError('MultiLaneReceiver: destroyed');
    this.overallHasher.update(plaintext);
  }

  /** Snapshot per-lane fingerprints (call once when all lanes are finished). */
  getLaneFingerprints(): LaneFingerprint[] {
    const out: LaneFingerprint[] = [];
    for (const [laneId, receiver] of this.receivers) {
      out.push({
        laneId,
        sha256: receiver.getLaneSha256Digest(),
        chunkCount: Number(receiver.nextExpectedSequence),
        byteCount: Number(receiver.bytesReceived),
      });
    }
    out.sort((a, b) => a.laneId - b.laneId);
    return out;
  }

  /** Snapshot overall sha256 over original-order plaintext. */
  getOverallSha256(): Uint8Array {
    return this.overallHasher.digest();
  }

  /** All lanes have received their `isLast` chunks. */
  get allLanesFinished(): boolean {
    for (const lane of this.receivers.values()) {
      if (!lane.isFinished) return false;
    }
    return true;
  }

  /** Get a specific lane receiver (for inspection/resume). */
  getLane(laneId: number): StreamReceiver | undefined {
    return this.receivers.get(laneId);
  }

  destroy(): void {
    if (this.destroyed) return;
    for (const receiver of this.receivers.values()) receiver.destroy();
    this.destroyed = true;
  }
}
