import { PipeableTransformStream } from "../utils/stream/PipeableTransformStream";
import { GorpStreamAccumulator } from "./GorpStreamAccumulator";
import { SSEDecoder, SSEEncoder } from "../utils/stream/SSE";
import type { GorpStreamChunk, GorpStreamOperation } from "./types";

export class GorpStreamEncoder extends PipeableTransformStream<
  GorpStreamChunk,
  Uint8Array
> {
  constructor() {
    let isFirstChunk = true;
    super((readable) =>
      readable
        .pipeThrough(
          new TransformStream<GorpStreamChunk, readonly GorpStreamOperation[]>({
            transform(chunk, controller) {
              if (isFirstChunk) {
                controller.enqueue([
                  { type: "set", path: [], value: chunk.snapshot },
                ]);
              } else {
                controller.enqueue(chunk.operations);
              }
              isFirstChunk = false;
            },
          }),
        )
        .pipeThrough(new SSEEncoder()),
    );
  }
}

export class GorpStreamDecoder extends PipeableTransformStream<
  Uint8Array<ArrayBuffer>,
  GorpStreamChunk
> {
  constructor() {
    const accumulator = new GorpStreamAccumulator();
    super((readable) =>
      readable
        .pipeThrough(new SSEDecoder<readonly GorpStreamOperation[]>())
        .pipeThrough(
          new TransformStream<readonly GorpStreamOperation[], GorpStreamChunk>({
            transform(operations, controller) {
              accumulator.append(operations);
              controller.enqueue({
                snapshot: accumulator.state,
                operations,
              });
            },
          }),
        ),
    );
  }
}

export class GorpStreamResponse extends Response {
  constructor(body: ReadableStream<GorpStreamChunk>) {
    super(body.pipeThrough(new GorpStreamEncoder()), {
      headers: new Headers({
        "Content-Type": "text/event-stream",
        "Cache-Control": "no-cache",
        Connection: "keep-alive",
        "Assistant-Stream-Format": "object-stream/v0",
      }),
    });
  }
}

export const fromGorpStreamResponse = (
  response: Response,
): ReadableStream<GorpStreamChunk> => {
  if (!response.ok) {
    void response.body?.cancel().catch(() => undefined);
    throw new Error(`Response failed, status ${response.status}`);
  }
  if (!response.body) throw new Error("Response body is null");
  const mediaType = response.headers
    .get("Content-Type")
    ?.split(";", 1)[0]
    ?.trim()
    .toLowerCase();
  if (mediaType !== "text/event-stream") {
    void response.body?.cancel().catch(() => undefined);
    throw new Error("Response is not an event stream");
  }
  if (response.headers.get("Assistant-Stream-Format") !== "object-stream/v0") {
    void response.body?.cancel().catch(() => undefined);
    throw new Error("Unsupported Assistant-Stream-Format header");
  }
  return response.body.pipeThrough(new GorpStreamDecoder());
};