import { env } from "cloudflare:workers";
import { toHex } from "../../crypto";
import { guardedFetch } from "../../git/relay";
import { ActionSourceError } from "./github";

/**
 * Registry images for Docker actions. A tag is resolved to a manifest digest when the run is scheduled, so a
 * moving tag can't change what a run executes. Blobs are pulled through the forge on first use and kept in
 * R2 under `ci-images/` (content-addressed, sha256-checked by R2 on write, swept 30 days after upload), shared
 * by every account like action archives: public image bytes are identical for everyone and a digest can only
 * ever name one content.
 */
export const IMAGE_PREFIX = "ci-images/blobs/sha256/";
export const IMAGE_TTL_MS = 30 * 24 * 60 * 60 * 1000;
const MAX_IMAGE_BYTES = 2 * 1024 ** 3;
const MAX_LAYERS = 128;
const MAX_MANIFEST_BYTES = 4 * 1024 * 1024;
const TAG_CACHE_SECONDS = 300;
const TOKEN_CACHE_SECONDS = 240;
const FETCH_TIMEOUT_MS = 30_000;
const DOCKER_HUB = "registry-1.docker.io";
const HTTP_UNAUTHORIZED = 401;
const DIGEST = /^sha256:([0-9a-f]{64})$/;
const REPOSITORY = /^[a-z0-9]+(?:(?:[._]|__|-+)[a-z0-9]+)*(?:\/[a-z0-9]+(?:(?:[._]|__|-+)[a-z0-9]+)*)*$/;
const TAG = /^[A-Za-z0-9_][A-Za-z0-9_.-]{0,127}$/;
const HOST = /^[A-Za-z0-9.-]+(?::\d{1,5})?$/;
const INDEX_TYPES = new Set([
  "application/vnd.oci.image.index.v1+json",
  "application/vnd.docker.distribution.manifest.list.v2+json",
]);
const MANIFEST_TYPES = new Set([
  "application/vnd.oci.image.manifest.v1+json",
  "application/vnd.docker.distribution.manifest.v2+json",
]);
const ACCEPT = [...INDEX_TYPES, ...MANIFEST_TYPES].join(", ");

export type ImageRef = { registry: string; repository: string; reference: string };
export type BlobDescriptor = { digest: string; size: number; mediaType: string };

/** A registry image pinned to one linux/amd64 manifest. */
export type ResolvedImage = {
  /** As written in the workflow or action.yml. */
  ref: string;
  registry: string;
  repository: string;
  digest: string;
  mediaType: string;
  manifestSize: number;
  config: BlobDescriptor;
  layers: BlobDescriptor[];
};

/** `alpine`, `ghcr.io/owner/tool:1.2`, `docker.io/library/node@sha256:...`. Docker Hub is the default registry. */
export function parseImageRef(raw: string): ImageRef {
  const value = raw.trim();
  let rest = value;
  let reference = "latest";
  const at = rest.indexOf("@");
  if (at !== -1) {
    reference = rest.slice(at + 1);
    rest = rest.slice(0, at);
    if (!DIGEST.test(reference)) throw new ActionSourceError(`'${value}': image digests must be sha256:<64 hex>`);
  }
  const slash = rest.indexOf("/");
  const first = slash === -1 ? "" : rest.slice(0, slash);
  const hasRegistry = first.includes(".") || first.includes(":") || first === "localhost";
  let registry = hasRegistry ? first.toLowerCase() : DOCKER_HUB;
  let path = hasRegistry ? rest.slice(slash + 1) : rest;
  const colon = path.lastIndexOf(":");
  if (colon > path.lastIndexOf("/")) {
    if (at === -1) reference = path.slice(colon + 1);
    path = path.slice(0, colon);
  }
  if (registry === "docker.io" || registry === "index.docker.io") registry = DOCKER_HUB;
  if (registry === DOCKER_HUB && !path.includes("/")) path = `library/${path}`;
  if (!HOST.test(registry) || !REPOSITORY.test(path) || !(TAG.test(reference) || DIGEST.test(reference))) {
    throw new ActionSourceError(`'${value}' is not a valid image reference`);
  }
  return { registry, repository: path, reference };
}

export function isDigestPinned(raw: string): boolean {
  return /@sha256:[0-9a-f]{64}$/.test(raw.trim());
}

function blobKey(digest: string): string {
  const match = DIGEST.exec(digest);
  if (!match) throw new Error(`invalid digest ${digest}`);
  return `${IMAGE_PREFIX}${match[1]}`;
}

/** The local tag a loaded image gets on the runner (`docker load` keeps it; steps run it by this name). */
export function imageTag(digest: string): string {
  return `gitfor-actions/image:${digest.replace("sha256:", "")}`;
}

function bearerChallenge(header: string | null): { realm: string; service: string | null } | null {
  if (!header || !/^bearer /i.test(header)) return null;
  const params = Object.fromEntries(
    [...header.slice("bearer ".length).matchAll(/(\w+)="([^"]*)"/g)].map((match) => [match[1].toLowerCase(), match[2]]),
  );
  return params.realm ? { realm: params.realm, service: params.service ?? null } : null;
}

/** Anonymous pull token for public images (Docker Hub, ghcr.io, quay.io and other token-auth registries). */
async function pullToken(
  image: Pick<ImageRef, "registry" | "repository">,
  challenge: { realm: string; service: string | null },
): Promise<string> {
  const cacheKey = `ci-image-token:v1:${image.registry}/${image.repository}`;
  const cached = await env.KV.get(cacheKey);
  if (cached) return cached;
  const url = new URL(challenge.realm);
  if (challenge.service) url.searchParams.set("service", challenge.service);
  url.searchParams.set("scope", `repository:${image.repository}:pull`);
  const response = await guardedFetch({ url: url.toString(), authorization: null }, url.toString(), {
    signal: AbortSignal.timeout(FETCH_TIMEOUT_MS),
  });
  if (!response.ok)
    throw new ActionSourceError(`registry ${image.registry} refused a pull token (HTTP ${response.status})`);
  const body = (await response.json()) as { token?: string; access_token?: string };
  const token = body.token ?? body.access_token;
  if (!token) throw new ActionSourceError(`registry ${image.registry} returned no pull token`);
  await env.KV.put(cacheKey, token, { expirationTtl: TOKEN_CACHE_SECONDS });
  return token;
}

/** GET against the registry API, answering a Bearer challenge once. Only public images are supported. */
export async function registryGet(
  image: Pick<ImageRef, "registry" | "repository">,
  path: string,
  accept?: string,
): Promise<Response> {
  const url = `https://${image.registry}/v2/${image.repository}/${path}`;
  const send = (token: string | null) => {
    const headers: Record<string, string> = { "user-agent": "gitfor" };
    if (accept) headers.accept = accept;
    if (token) headers.authorization = `Bearer ${token}`;
    return guardedFetch({ url, authorization: null }, url, { headers, signal: AbortSignal.timeout(FETCH_TIMEOUT_MS) });
  };
  let response = await send(await env.KV.get(`ci-image-token:v1:${image.registry}/${image.repository}`));
  if (response.status === HTTP_UNAUTHORIZED) {
    const challenge = bearerChallenge(response.headers.get("www-authenticate"));
    await response.body?.cancel();
    if (!challenge)
      throw new ActionSourceError(`registry ${image.registry} requires credentials; only public images are supported`);
    response = await send(await pullToken(image, challenge));
  }
  return response;
}

async function readManifest(
  image: ImageRef,
  reference: string,
): Promise<{ bytes: Uint8Array; mediaType: string; digest: string }> {
  const response = await registryGet(image, `manifests/${reference}`, ACCEPT);
  if (!response.ok) {
    await response.body?.cancel();
    throw new ActionSourceError(
      `image ${image.registry}/${image.repository}:${reference} was not found or is private (HTTP ${response.status})`,
    );
  }
  if (Number(response.headers.get("content-length") ?? "0") > MAX_MANIFEST_BYTES) {
    await response.body?.cancel();
    throw new ActionSourceError("the image manifest is too large");
  }
  const bytes = new Uint8Array(await response.arrayBuffer());
  if (bytes.byteLength > MAX_MANIFEST_BYTES) throw new ActionSourceError("the image manifest is too large");
  const digest = `sha256:${toHex(await crypto.subtle.digest("SHA-256", bytes))}`;
  if (DIGEST.test(reference) && digest !== reference) {
    throw new ActionSourceError(`the registry served a manifest that does not match ${reference}`);
  }
  const mediaType =
    (JSON.parse(new TextDecoder().decode(bytes)) as { mediaType?: string }).mediaType ??
    response.headers.get("content-type")?.split(";")[0] ??
    "";
  return { bytes, mediaType, digest };
}

type Index = { manifests?: Array<{ digest: string; platform?: { os?: string; architecture?: string } }> };
type Manifest = { config?: BlobDescriptor; layers?: BlobDescriptor[] };

function descriptor(value: BlobDescriptor | undefined, what: string): BlobDescriptor {
  if (!value || !DIGEST.test(value.digest) || !Number.isSafeInteger(value.size) || value.size < 0) {
    throw new ActionSourceError(`the image manifest has an invalid ${what}`);
  }
  return { digest: value.digest, size: value.size, mediaType: String(value.mediaType ?? "") };
}

async function storeManifest(digest: string, bytes: Uint8Array): Promise<void> {
  const key = blobKey(digest);
  if (await env.BUCKET.head(key)) return;
  await env.BUCKET.put(key, bytes, { sha256: digest.slice("sha256:".length) });
}

/** Resolves a tag to the linux/amd64 manifest and stores the manifest, so jobs pull by an immutable digest. */
export async function resolveImage(raw: string): Promise<ResolvedImage> {
  const image = parseImageRef(raw);
  const cacheKey = `ci-image:v1:${image.registry}/${image.repository}@${image.reference}`;
  const cached = await env.KV.get<ResolvedImage>(cacheKey, "json");
  if (cached) return { ...cached, ref: raw.trim() };
  let manifest = await readManifest(image, image.reference);
  if (INDEX_TYPES.has(manifest.mediaType)) {
    const index = JSON.parse(new TextDecoder().decode(manifest.bytes)) as Index;
    const entry = index.manifests?.find(
      (candidate) => candidate.platform?.os === "linux" && candidate.platform.architecture === "amd64",
    );
    if (!entry || !DIGEST.test(entry.digest)) throw new ActionSourceError(`'${raw}' has no linux/amd64 image`);
    manifest = await readManifest(image, entry.digest);
  }
  if (!MANIFEST_TYPES.has(manifest.mediaType)) {
    throw new ActionSourceError(`'${raw}' uses an unsupported manifest type '${manifest.mediaType}'`);
  }
  const parsed = JSON.parse(new TextDecoder().decode(manifest.bytes)) as Manifest;
  const config = descriptor(parsed.config, "config");
  const layers = (parsed.layers ?? []).map((layer) => descriptor(layer, "layer"));
  if (layers.length === 0 || layers.length > MAX_LAYERS)
    throw new ActionSourceError(`'${raw}' has ${layers.length} layers`);
  const total = layers.reduce((sum, layer) => sum + layer.size, config.size);
  if (total > MAX_IMAGE_BYTES) {
    throw new ActionSourceError(`'${raw}' is ${Math.round(total / 1024 ** 2)} MiB; images may be at most 2 GiB`);
  }
  await storeManifest(manifest.digest, manifest.bytes);
  const resolved: ResolvedImage = {
    ref: raw.trim(),
    registry: image.registry,
    repository: image.repository,
    digest: manifest.digest,
    mediaType: manifest.mediaType,
    manifestSize: manifest.bytes.byteLength,
    config,
    layers,
  };
  const ttl = DIGEST.test(image.reference) ? IMAGE_TTL_MS / 1000 : TAG_CACHE_SECONDS;
  await env.KV.put(cacheKey, JSON.stringify(resolved), { expirationTtl: ttl });
  return resolved;
}

/**
 * A blob from the cache, or pulled from the registry while it is written to the cache. R2 checks the sha256
 * on write, and `done` rejects when it doesn't match, so callers must await it before trusting the bytes.
 */
export async function imageBlob(
  image: ResolvedImage,
  blob: BlobDescriptor,
): Promise<{ body: ReadableStream<Uint8Array>; done: Promise<void> }> {
  const key = blobKey(blob.digest);
  const cached = await env.BUCKET.get(key);
  if (cached && cached.size === blob.size) return { body: cached.body, done: Promise.resolve() };
  if (cached) await cached.body.cancel();
  const isManifest = blob.digest === image.digest;
  const response = await registryGet(
    image,
    isManifest ? `manifests/${blob.digest}` : `blobs/${blob.digest}`,
    isManifest ? ACCEPT : undefined,
  );
  if (!response.ok || !response.body) {
    await response.body?.cancel();
    throw new Error(`registry returned HTTP ${response.status} for ${blob.digest}`);
  }
  const [forClient, forStore] = response.body.tee();
  const sized = new FixedLengthStream(blob.size);
  const writing = forStore.pipeTo(sized.writable);
  const stored = env.BUCKET.put(key, sized.readable, { sha256: blob.digest.slice("sha256:".length) });
  const done = Promise.all([writing, stored]).then(() => undefined);
  return { body: forClient, done };
}

/** Drops cached image blobs uploaded more than IMAGE_TTL_MS ago; a run that still needs one re-pulls it. */
export async function sweepImageBlobs(limit = 1000): Promise<number> {
  const cutoff = Date.now() - IMAGE_TTL_MS;
  const cursorKey = "ci-images:sweep-cursor";
  const cursor = (await env.KV.get(cursorKey)) ?? undefined;
  const listing = await env.BUCKET.list({ prefix: IMAGE_PREFIX, cursor, limit });
  const stale = listing.objects.filter((object) => object.uploaded.getTime() < cutoff).map((object) => object.key);
  if (stale.length > 0) await env.BUCKET.delete(stale);
  if (listing.truncated) await env.KV.put(cursorKey, listing.cursor);
  else await env.KV.delete(cursorKey);
  return stale.length;
}