import { createAtomicCheckpoint, type AtomicCheckpoint } from "./atomic-checkpoint";

export type StreamErrorCode = "STREAM_DISCONNECTED" | "STREAM_UTF8_INVALID" | "STREAM_CANCELLED" | "STREAM_CHECKPOINT_FAILED";

function boundedError(error: unknown): Error | unknown {
  const message = error instanceof Error ? error.message : String(error);
  if (message.length <= 512) return error;
  const result = new Error(`${message.slice(0, 509)}...`);
  if (error instanceof Error && error.name) result.name = error.name;
  return result;
}

export class StreamConsumeError extends Error {
  readonly code: StreamErrorCode;
  readonly byteCount: number;
  readonly chunkCount: number;

  constructor(code: StreamErrorCode, message: string, byteCount = 0, chunkCount = 0, options?: { cause?: unknown }) {
    const cause = options && Object.hasOwn(options, "cause") ? boundedError(options.cause) : undefined;
    super(`${code}: ${message}`, cause === undefined ? undefined : { cause });
    this.name = "StreamConsumeError";
    this.code = code;
    this.byteCount = byteCount;
    this.chunkCount = chunkCount;
  }
}

export interface StreamProgress {
  byteCount: number;
  chunkCount: number;
  sha256: string;
}

export interface StreamConsumerOptions {
  signal?: AbortSignal;
  requireCompletion?: boolean;
  isComplete?: (progress: StreamProgress) => boolean;
  onCheckpoint?: (checkpoint: AtomicCheckpoint) => Promise<void> | void;
  /**
   * Recreate the response after a transport-level read failure. The returned
   * stream must start at `progress.byteCount`; already consumed bytes are
   * retained and are never published twice.
   */
  reconnect?: (progress: StreamProgress, error: unknown) =>
    | (ReadableStream<Uint8Array> | AsyncIterable<Uint8Array>)
    | Promise<ReadableStream<Uint8Array> | AsyncIterable<Uint8Array>>;
  /** Maximum number of reconnect attempts after the initial stream. */
  maxRetries?: number;
  /** Optional delay between reconnect attempts. */
  retryDelayMs?: number;
}

export interface StreamConsumeResult extends StreamProgress {
  text: string;
}

async function sha256Hex(bytes: Uint8Array): Promise<string> {
  const subtle = globalThis.crypto?.subtle;
  if (!subtle) throw new Error("SHA-256 is unavailable in this runtime");
  const digest = await subtle.digest("SHA-256", bytes.slice().buffer);
  return [...new Uint8Array(digest)].map((value) => value.toString(16).padStart(2, "0")).join("");
}

function concatParts(parts: Uint8Array[]): Uint8Array {
  const result = new Uint8Array(parts.reduce((total, part) => total + part.byteLength, 0));
  let offset = 0;
  for (const part of parts) {
    result.set(part, offset);
    offset += part.byteLength;
  }
  return result;
}

function asStream(source: ReadableStream<Uint8Array> | AsyncIterable<Uint8Array>): ReadableStream<Uint8Array> {
  if (typeof (source as ReadableStream<Uint8Array>).getReader === "function") return source as ReadableStream<Uint8Array>;
  const iterable = source as AsyncIterable<Uint8Array>;
  const iterator = iterable[Symbol.asyncIterator]();
  let iteratorClosed = false;
  const closeIterator = async (reason?: unknown): Promise<void> => {
    if (iteratorClosed) return;
    iteratorClosed = true;
    await iterator.return?.(reason);
  };
  return new ReadableStream<Uint8Array>({
    async pull(controller) {
      try {
        const next = await iterator.next();
        if (next.done) controller.close(); else controller.enqueue(next.value);
      } catch (error) {
        await closeIterator(error).catch(() => undefined);
        controller.error(error);
      }
    },
    async cancel(reason) {
      // A fetch body is cancelled by the reader, but an async iterable needs
      // its return hook called explicitly so generators can release sockets or
      // other resources after a transport failure.
      await closeIterator(reason);
    },
  });
}

function errorMessage(error: unknown): string {
  const message = error instanceof Error ? error.message : String(error);
  // Provider/network errors occasionally contain a full response body. Keep
  // the diagnostic bounded so a failure cannot recreate the context problem.
  return message.length > 512 ? `${message.slice(0, 509)}...` : message;
}

async function checkpoint(options: StreamConsumerOptions, state: "STREAMING" | "COMPLETED" | "FAILED", bytes: Uint8Array, byteCount: number, chunkCount: number): Promise<void> {
  if (!options.onCheckpoint) return;
  const sha256 = await sha256Hex(bytes);
  try { await options.onCheckpoint(createAtomicCheckpoint({ state, byteCount, chunkCount, sha256 })); }
  catch (error) { throw new StreamConsumeError("STREAM_CHECKPOINT_FAILED", "checkpoint publication failed", byteCount, chunkCount, { cause: error }); }
}

function waitForRetry(delayMs: number, signal?: AbortSignal): Promise<boolean> {
  if (delayMs <= 0) return Promise.resolve(signal?.aborted !== true);
  return new Promise((resolve) => {
    let timer: ReturnType<typeof setTimeout> | undefined;
    const finish = (ready: boolean): void => {
      if (timer !== undefined) clearTimeout(timer);
      signal?.removeEventListener("abort", onAbort);
      resolve(ready);
    };
    const onAbort = (): void => finish(false);
    timer = setTimeout(() => finish(signal?.aborted !== true), delayMs);
    signal?.addEventListener("abort", onAbort, { once: true });
    if (signal?.aborted) finish(false);
  });
}

/** Consume a UTF-8 byte stream without replacement characters or partial durable output. */
export async function consumeUtf8Stream(source: ReadableStream<Uint8Array> | AsyncIterable<Uint8Array>, options: StreamConsumerOptions = {}): Promise<StreamConsumeResult> {
  const maxRetries = options.maxRetries ?? (options.reconnect ? 1 : 0);
  if (!Number.isSafeInteger(maxRetries) || maxRetries < 0) {
    throw new TypeError("STREAM_RETRY_INVALID: maxRetries must be a non-negative safe integer");
  }
  if (options.retryDelayMs !== undefined && (!Number.isFinite(options.retryDelayMs) || options.retryDelayMs < 0)) {
    throw new TypeError("STREAM_RETRY_INVALID: retryDelayMs must be a non-negative number");
  }
  if (maxRetries && !options.reconnect) {
    throw new TypeError("STREAM_RETRY_INVALID: reconnect is required when maxRetries is non-zero");
  }
  let reader = asStream(source).getReader();
  const decoder = new TextDecoder("utf-8", { fatal: true });
  const parts: Uint8Array[] = [];
  let text = "";
  let byteCount = 0;
  let chunkCount = 0;
  let completed = false;
  let retryCount = 0;
  let readerLockHeld = true;
  const aborted = (): boolean => options.signal?.aborted === true;
  const cancel = (): void => { void reader.cancel("STREAM_CANCELLED").catch(() => undefined); };
  options.signal?.addEventListener("abort", cancel, { once: true });
  const reconnect = async (error: unknown): Promise<boolean> => {
    if (!options.reconnect || retryCount >= maxRetries || aborted()) return false;
    const progress = { byteCount, chunkCount, sha256: await sha256Hex(concatParts(parts)) };
    const previousReader = reader;
    retryCount += 1;
    await previousReader.cancel("STREAM_RECONNECT").catch(() => undefined);
    previousReader.releaseLock();
    readerLockHeld = false;
    if (!await waitForRetry(options.retryDelayMs ?? 0, options.signal)) return false;
    if (aborted()) return false;
    let replacement;
    try { replacement = await options.reconnect(progress, error); }
    catch { return false; }
    if (!replacement) return false;
    try { reader = asStream(replacement).getReader(); }
    catch { return false; }
    if (aborted()) {
      await reader.cancel("STREAM_CANCELLED").catch(() => undefined);
      reader.releaseLock();
      readerLockHeld = false;
      return false;
    }
    readerLockHeld = true;
    return true;
  };
  try {
    while (true) {
      if (aborted()) throw new StreamConsumeError("STREAM_CANCELLED", "stream consumption was cancelled", byteCount, chunkCount);
      let next: ReadableStreamReadResult<Uint8Array>;
      try { next = await reader.read(); }
      catch (error) {
        if (aborted()) throw new StreamConsumeError("STREAM_CANCELLED", "stream consumption was cancelled", byteCount, chunkCount, { cause: error });
        if (await reconnect(error)) continue;
        if (aborted()) throw new StreamConsumeError("STREAM_CANCELLED", "stream consumption was cancelled", byteCount, chunkCount, { cause: error });
        throw new StreamConsumeError("STREAM_DISCONNECTED", `stream disconnected before completion: ${errorMessage(error)}`, byteCount, chunkCount, { cause: error });
      }
      if (aborted()) throw new StreamConsumeError("STREAM_CANCELLED", "stream consumption was cancelled", byteCount, chunkCount);
      if (next.done) break;
      if (!(next.value instanceof Uint8Array)) throw new StreamConsumeError("STREAM_DISCONNECTED", "stream yielded a non-byte chunk", byteCount, chunkCount);
      const bytes = new Uint8Array(next.value);
      parts.push(bytes);
      byteCount += bytes.byteLength;
      chunkCount += 1;
      try { text += decoder.decode(bytes, { stream: true }); }
      catch (error) { throw new StreamConsumeError("STREAM_UTF8_INVALID", "stream ended or yielded an invalid UTF-8 sequence", byteCount, chunkCount, { cause: error }); }
      await checkpoint(options, "STREAMING", concatParts(parts), byteCount, chunkCount);
    }
    try { text += decoder.decode(); }
    catch (error) { throw new StreamConsumeError("STREAM_UTF8_INVALID", "stream ended with an incomplete UTF-8 sequence", byteCount, chunkCount, { cause: error }); }
    const progress = { byteCount, chunkCount, sha256: await sha256Hex(concatParts(parts)) };
    if ((options.requireCompletion && !options.isComplete) || (options.isComplete && !options.isComplete(progress))) {
      throw new StreamConsumeError("STREAM_DISCONNECTED", "stream disconnected before completion", byteCount, chunkCount);
    }
    await checkpoint(options, "COMPLETED", concatParts(parts), byteCount, chunkCount);
    completed = true;
    return { text, ...progress };
  } catch (error) {
    if (!completed && error instanceof StreamConsumeError && options.onCheckpoint) {
      try {
        const bytes = concatParts(parts);
        await options.onCheckpoint(createAtomicCheckpoint({ state: "FAILED", byteCount, chunkCount, sha256: await sha256Hex(bytes) }));
      } catch {
        // Preserve the original transport/UTF-8 error; failed checkpoints are advisory.
      }
    }
    throw error;
  } finally {
    options.signal?.removeEventListener("abort", cancel);
    if (!completed) await reader.cancel("STREAM_CONSUMER_STOPPED").catch(() => undefined);
    if (readerLockHeld) reader.releaseLock();
  }
}

export const consumeModelStream = consumeUtf8Stream;
export const consumeStream = consumeUtf8Stream;