import { waitUntil } from "cloudflare:workers";
import { jobDb } from "../scope";
import { tarHeader, tarPadding, TAR_BLOCK } from "../../tar";
import { imageBlob, imageTag } from "./images";
import type { BlobDescriptor, ResolvedImage } from "./images";
import type { CompiledJob } from "../compile";

const OCI_INDEX = "application/vnd.oci.image.index.v1+json";
const FILE_MODE = 0o644;
const encoder = new TextEncoder();

function blobPath(digest: string): string {
  return `blobs/sha256/${digest.replace("sha256:", "")}`;
}

/**
 * The tar `docker load` takes: an OCI image layout (index.json + blobs) that also carries the classic
 * manifest.json, the same shape `docker save` writes, tagged so steps can run it by `imageTag(digest)`.
 */
function layoutFiles(image: ResolvedImage): Array<[string, Uint8Array]> {
  const tag = imageTag(image.digest);
  const index = {
    schemaVersion: 2,
    mediaType: OCI_INDEX,
    manifests: [
      {
        mediaType: image.mediaType,
        digest: image.digest,
        size: image.manifestSize,
        annotations: {
          "io.containerd.image.name": `docker.io/${tag}`,
          "org.opencontainers.image.ref.name": tag.split(":")[1],
        },
      },
    ],
  };
  const manifest = [
    {
      Config: blobPath(image.config.digest),
      RepoTags: [tag],
      Layers: image.layers.map((layer) => blobPath(layer.digest)),
    },
  ];
  return [
    ["oci-layout", encoder.encode(JSON.stringify({ imageLayoutVersion: "1.0.0" }))],
    ["index.json", encoder.encode(JSON.stringify(index))],
    ["manifest.json", encoder.encode(JSON.stringify(manifest))],
  ];
}

async function writeEntry(writer: WritableStreamDefaultWriter<Uint8Array>, path: string, size: number) {
  const header = tarHeader(path, size, FILE_MODE, "0", "", 0);
  if (!header) throw new Error(`cannot encode ${path} in a tar header`);
  await writer.write(header);
}

async function writeBlob(
  writer: WritableStreamDefaultWriter<Uint8Array>,
  image: ResolvedImage,
  blob: BlobDescriptor,
): Promise<void> {
  const { body, done } = await imageBlob(image, blob);
  await writeEntry(writer, blobPath(blob.digest), blob.size);
  const reader = body.getReader();
  let written = 0;
  for (;;) {
    const { value, done: finished } = await reader.read();
    if (finished) break;
    written += value.byteLength;
    await writer.write(value);
  }
  // A blob that doesn't match its digest fails R2's sha256 check; abort before the tar can complete.
  await done;
  if (written !== blob.size) throw new Error(`blob ${blob.digest} was ${written} bytes, expected ${blob.size}`);
  const padding = tarPadding(blob.size);
  if (padding > 0) await writer.write(new Uint8Array(padding));
}

async function streamLayout(writer: WritableStreamDefaultWriter<Uint8Array>, image: ResolvedImage): Promise<void> {
  for (const [path, bytes] of layoutFiles(image)) {
    await writeEntry(writer, path, bytes.byteLength);
    await writer.write(bytes);
    const padding = tarPadding(bytes.byteLength);
    if (padding > 0) await writer.write(new Uint8Array(padding));
  }
  const manifest = { digest: image.digest, size: image.manifestSize, mediaType: image.mediaType };
  for (const blob of [manifest, image.config, ...image.layers]) await writeBlob(writer, image, blob);
  await writer.write(new Uint8Array(TAR_BLOCK * 2));
}

/**
 * Streams a pinned image, as a `docker load` tar, to a job whose own spec runs it. Blobs come from the R2 cache
 * or are pulled from the registry through it; a blob that fails its digest aborts the stream, so `docker load`
 * sees a truncated archive and the step fails instead of running tampered bytes.
 */
export async function serveActionImage(jobId: string, digest: string): Promise<Response> {
  const db = await jobDb(jobId);
  const job = await db?.first<{ spec: string }>("SELECT spec FROM workflow_jobs WHERE id = ?", jobId);
  if (!job) return Response.json({ error: "job not found" }, { status: 404 });
  const spec: CompiledJob = JSON.parse(job.spec);
  const image = spec.steps
    .map((step) => step.action?.docker?.image)
    .find((candidate) => candidate?.kind === "registry" && candidate.digest === digest);
  if (!image || image.kind !== "registry") {
    return Response.json({ error: "this job does not use that image" }, { status: 403 });
  }
  const { readable, writable } = new TransformStream<Uint8Array, Uint8Array>();
  const writer = writable.getWriter();
  waitUntil(
    streamLayout(writer, image)
      .then(() => writer.close())
      .catch(async (error: unknown) => {
        console.error("action image stream failed", digest, error);
        await writer.abort(error).catch(() => undefined);
      }),
  );
  return new Response(readable, {
    headers: { "content-type": "application/x-tar", "cache-control": "private, no-store" },
  });
}