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;
}

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 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 }); }
}

/** 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 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;
  const aborted = (): boolean => options.signal?.aborted === true;
  const cancel = (): void => { void reader.cancel("STREAM_CANCELLED").catch(() => undefined); };
  options.signal?.addEventListener("abort", cancel, { once: 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 });
        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", Uint8Array.from(parts.flatMap((part) => [...part])), 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(Uint8Array.from(parts.flatMap((part) => [...part]))) };
    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", Uint8Array.from(parts.flatMap((part) => [...part])), byteCount, chunkCount);
    completed = true;
    return { text, ...progress };
  } catch (error) {
    if (!completed && error instanceof StreamConsumeError && options.onCheckpoint) {
      try {
        const bytes = Uint8Array.from(parts.flatMap((part) => [...part]));
        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);
    reader.releaseLock();
  }
}

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