import "server-only";
import { createHash, randomUUID, type Hash } from "node:crypto";
import { createReadStream, createWriteStream, mkdirSync, readdirSync, renameSync, rmSync, statSync, writeFileSync } from "node:fs";
import { stat } from "node:fs/promises";
import path from "node:path";
import { Readable, Transform } from "node:stream";
import { pipeline } from "node:stream/promises";
import { db, transaction } from "./db";
import { DATA_ROOT } from "./paths";

// A container registry per repository. Images are named <owner>/<repo> or <owner>/<repo>/<path>;
// `image` below is that path ("" for the repository's own image). Layers and configs are
// content-addressed files shared by all repositories, and a repository may only serve the blobs
// that were pushed to or mounted into it. Manifests and tags live in the database.

const BLOBS_ROOT = path.join(DATA_ROOT, "registry", "blobs");
const UPLOADS_ROOT = path.join(DATA_ROOT, "registry", "uploads");

export const MAX_MANIFEST_BYTES = 4 * 1024 * 1024;

/** One path component of an image name, as defined by the OCI distribution spec. */
const COMPONENT_RE = /^[a-z0-9]+(?:(?:[._]|__|-+)[a-z0-9]+)*$/;
export const TAG_RE = /^[A-Za-z0-9_][A-Za-z0-9._-]{0,127}$/;
const DIGEST_RE = /^(sha256:[a-f0-9]{64}|sha512:[a-f0-9]{128})$/;

export const isDigest = (s: string) => DIGEST_RE.test(s);
export const isImagePath = (s: string) => s === "" || s.split("/").every((c) => COMPONENT_RE.test(c));

/** Whether owner/repo can be written in an image name (Docker requires lowercase components). */
export const isRegistryName = (owner: string, repo: string) =>
  COMPONENT_RE.test(owner.toLowerCase()) && COMPONENT_RE.test(repo.toLowerCase());

export class RegistryError extends Error {
  constructor(
    readonly status: number,
    readonly code: string,
    message: string,
  ) {
    super(message);
  }
}

function blobPath(digest: string) {
  const [algorithm, hex] = digest.split(":");
  return path.join(BLOBS_ROOT, algorithm, hex.slice(0, 2), hex);
}

const uploadPath = (id: string) => path.join(UPLOADS_ROOT, id);

// ---------------------------------------------------------------------------
// Blobs

export type BlobInfo = { digest: string; size: number; path: string };

/** A blob the repository may serve, touching its link so a push in progress keeps it. */
export function findBlob(repoId: number, digest: string): BlobInfo | null {
  const row = db
    .prepare("SELECT b.size FROM registry_repo_blobs l JOIN registry_blobs b ON b.digest = l.digest WHERE l.repo_id = ? AND l.digest = ?")
    .get(repoId, digest) as { size: number } | undefined;
  if (!row) return null;
  db.prepare("UPDATE registry_repo_blobs SET created_at = ? WHERE repo_id = ? AND digest = ?").run(Date.now(), repoId, digest);
  return { digest, size: row.size, path: blobPath(digest) };
}

function linkBlob(repoId: number, digest: string) {
  db.prepare(
    `INSERT INTO registry_repo_blobs (repo_id, digest, created_at) VALUES (?, ?, ?)
     ON CONFLICT (repo_id, digest) DO UPDATE SET created_at = excluded.created_at`,
  ).run(repoId, digest, Date.now());
}

/** Makes a blob of another repository available in this one without uploading it again. */
export function mountBlob(repoId: number, fromRepoId: number, digest: string) {
  if (!findBlob(fromRepoId, digest)) return false;
  linkBlob(repoId, digest);
  return true;
}

export function unlinkBlob(repoId: number, digest: string) {
  const res = db.prepare("DELETE FROM registry_repo_blobs WHERE repo_id = ? AND digest = ?").run(repoId, digest);
  pruneRegistry();
  return res.changes > 0;
}

// ---------------------------------------------------------------------------
// Uploads

export type Upload = { id: string; repoId: number; size: number };

// Running hashes of uploads whose chunks arrived in order, so finishing a large layer doesn't read it again.
const g = globalThis as unknown as { __uploadHashes?: Map<string, { hash: Hash; size: number }>; __registryPrunedAt?: number };
const uploadHashes = (g.__uploadHashes ??= new Map());

export function startUpload(repoId: number): Upload {
  maybePrune();
  const id = randomUUID();
  const now = Date.now();
  db.prepare("INSERT INTO registry_uploads (id, repo_id, size, created_at, updated_at) VALUES (?, ?, 0, ?, ?)").run(id, repoId, now, now);
  mkdirSync(UPLOADS_ROOT, { recursive: true });
  writeFileSync(uploadPath(id), "");
  uploadHashes.set(id, { hash: createHash("sha256"), size: 0 });
  return { id, repoId, size: 0 };
}

export function findUpload(repoId: number, id: string): Upload | null {
  const row = db.prepare("SELECT size FROM registry_uploads WHERE id = ? AND repo_id = ?").get(id, repoId) as
    | { size: number }
    | undefined;
  return row ? { id, repoId, size: row.size } : null;
}

/** Appends a request body to an upload and returns its new size. */
export async function appendUpload(upload: Upload, body: ReadableStream<Uint8Array> | null) {
  if (!body) return upload.size;
  const running = uploadHashes.get(upload.id);
  const hash = running?.size === upload.size ? running.hash : null;
  uploadHashes.delete(upload.id);
  const tap = new Transform({
    transform(chunk, _enc, done) {
      hash?.update(chunk);
      done(null, chunk);
    },
  });
  await pipeline(Readable.fromWeb(body as never), tap, createWriteStream(uploadPath(upload.id), { flags: "a" }));
  const { size } = await stat(uploadPath(upload.id));
  db.prepare("UPDATE registry_uploads SET size = ?, updated_at = ? WHERE id = ?").run(size, Date.now(), upload.id);
  if (hash) uploadHashes.set(upload.id, { hash, size });
  return size;
}

async function hashFile(file: string, algorithm: string) {
  const hash = createHash(algorithm);
  await pipeline(createReadStream(file), hash);
  return hash.digest("hex");
}

/** Checks the uploaded content against `digest` and stores it as a blob of the upload's repository. */
export async function finishUpload(upload: Upload, digest: string) {
  const file = uploadPath(upload.id);
  const [algorithm, expected] = digest.split(":");
  const running = uploadHashes.get(upload.id);
  uploadHashes.delete(upload.id);
  const { size } = await stat(file);
  const actual =
    algorithm === "sha256" && running?.size === size ? running.hash.digest("hex") : await hashFile(file, algorithm);
  if (actual !== expected) {
    cancelUpload(upload);
    throw new RegistryError(400, "DIGEST_INVALID", "The uploaded content does not match the digest.");
  }
  // Synchronous from here on, so a concurrent prune can't remove the blob between moving and recording it.
  if (!db.prepare("SELECT 1 FROM registry_blobs WHERE digest = ?").get(digest)) {
    mkdirSync(path.dirname(blobPath(digest)), { recursive: true });
    renameSync(file, blobPath(digest));
  }
  transaction(() => {
    db.prepare("INSERT OR IGNORE INTO registry_blobs (digest, size, created_at) VALUES (?, ?, ?)").run(digest, size, Date.now());
    linkBlob(upload.repoId, digest);
    db.prepare("DELETE FROM registry_uploads WHERE id = ?").run(upload.id);
  });
  rmSync(file, { force: true });
  return size;
}

export function cancelUpload(upload: Upload) {
  uploadHashes.delete(upload.id);
  db.prepare("DELETE FROM registry_uploads WHERE id = ?").run(upload.id);
  rmSync(uploadPath(upload.id), { force: true });
}

// ---------------------------------------------------------------------------
// Manifests and tags

export const MEDIA_TYPES = {
  ociManifest: "application/vnd.oci.image.manifest.v1+json",
  ociIndex: "application/vnd.oci.image.index.v1+json",
  dockerManifest: "application/vnd.docker.distribution.manifest.v2+json",
  dockerList: "application/vnd.docker.distribution.manifest.list.v2+json",
};
const INDEX_TYPES = new Set([MEDIA_TYPES.ociIndex, MEDIA_TYPES.dockerList]);
const SCHEMA1_TYPES = new Set([
  "application/vnd.docker.distribution.manifest.v1+json",
  "application/vnd.docker.distribution.manifest.v1+prettyjws",
]);

type Descriptor = {
  mediaType?: string;
  digest?: string;
  size?: number;
  urls?: string[];
  platform?: { os?: string; architecture?: string; variant?: string };
  annotations?: Record<string, string>;
};
type ManifestJson = {
  mediaType?: string;
  artifactType?: string;
  config?: Descriptor;
  layers?: Descriptor[];
  manifests?: Descriptor[];
  subject?: Descriptor;
  annotations?: Record<string, string>;
};

export type Manifest = { digest: string; mediaType: string; content: Uint8Array };

type ManifestRow = { digest: string; media_type: string; content: Uint8Array; artifact_type: string | null };

const toManifest = (r: ManifestRow): Manifest => ({ digest: r.digest, mediaType: r.media_type, content: r.content });
const parse = (content: Uint8Array) => JSON.parse(Buffer.from(content).toString("utf8")) as ManifestJson;

/** Looks up a manifest by tag or digest. */
export function findManifest(repoId: number, image: string, reference: string): Manifest | null {
  const row = (
    isDigest(reference)
      ? db.prepare("SELECT * FROM registry_manifests WHERE repo_id = ? AND image = ? AND digest = ?").get(repoId, image, reference)
      : db
          .prepare(
            `SELECT m.* FROM registry_tags t JOIN registry_manifests m
             ON m.repo_id = t.repo_id AND m.image = t.image AND m.digest = t.digest
             WHERE t.repo_id = ? AND t.image = ? AND t.tag = ?`,
          )
          .get(repoId, image, reference)
  ) as ManifestRow | undefined;
  return row ? toManifest(row) : null;
}

const isDescriptor = (d: unknown): d is Descriptor & { digest: string } =>
  !!d && typeof d === "object" && typeof (d as Descriptor).digest === "string" && isDigest((d as Descriptor).digest!);

/**
 * Stores a pushed manifest under its digest and, if `reference` is a tag, points the tag at it.
 * The blobs an image manifest references must already be in the repository.
 */
export function putManifest(repoId: number, image: string, reference: string, contentType: string | null, content: Uint8Array) {
  const algorithm = isDigest(reference) ? reference.split(":")[0] : "sha256";
  const digest = `${algorithm}:${createHash(algorithm).update(content).digest("hex")}`;
  if (isDigest(reference) && reference !== digest) {
    throw new RegistryError(400, "DIGEST_INVALID", "The manifest does not match the digest in the URL.");
  }

  let json: ManifestJson;
  try {
    json = JSON.parse(Buffer.from(content).toString("utf8"));
    if (!json || typeof json !== "object" || Array.isArray(json)) throw new Error();
  } catch {
    throw new RegistryError(400, "MANIFEST_INVALID", "The manifest is not valid JSON.");
  }
  const mediaType = contentType?.split(";")[0].trim() || json.mediaType;
  if (!mediaType) throw new RegistryError(400, "MANIFEST_INVALID", "The manifest has no media type.");
  if (SCHEMA1_TYPES.has(mediaType)) {
    throw new RegistryError(400, "MANIFEST_INVALID", "Docker schema 1 manifests are not supported; push with a current client.");
  }
  if (json.mediaType && json.mediaType !== mediaType) {
    throw new RegistryError(400, "MANIFEST_INVALID", "The Content-Type does not match the manifest's mediaType.");
  }

  const isIndex = INDEX_TYPES.has(mediaType) || (Array.isArray(json.manifests) && !json.layers);
  const children = (isIndex ? (json.manifests ?? []) : []).filter(isDescriptor);
  // Foreign layers (with URLs) are fetched from elsewhere and aren't uploaded to the registry.
  const blobs = isIndex ? [] : [json.config, ...(json.layers ?? [])].filter(isDescriptor).filter((d) => !d.urls?.length);
  for (const blob of blobs) {
    if (!db.prepare("SELECT 1 FROM registry_repo_blobs WHERE repo_id = ? AND digest = ?").get(repoId, blob.digest)) {
      throw new RegistryError(400, "MANIFEST_BLOB_UNKNOWN", `Blob ${blob.digest} is not in this repository. Push it first.`);
    }
  }
  const size = isIndex ? 0 : [json.config, ...(json.layers ?? [])].reduce((sum, d) => sum + (Number(d?.size) || 0), 0);
  const subject = isDescriptor(json.subject) ? json.subject.digest : null;
  // As reported by the referrers API: the declared artifact type, else an image manifest's config type.
  const artifactType = json.artifactType ?? (isIndex ? null : (json.config?.mediaType ?? null));

  transaction(() => {
    db.prepare(
      `INSERT OR IGNORE INTO registry_manifests (repo_id, image, digest, media_type, content, size, subject, artifact_type, created_at)
       VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`,
    ).run(repoId, image, digest, mediaType, content, size, subject, artifactType, Date.now());
    for (const ref of [...children, ...blobs]) {
      db.prepare("INSERT OR IGNORE INTO registry_manifest_refs (repo_id, image, manifest_digest, digest) VALUES (?, ?, ?, ?)").run(
        repoId,
        image,
        digest,
        ref.digest,
      );
    }
    if (!isDigest(reference)) {
      db.prepare(
        `INSERT INTO registry_tags (repo_id, image, tag, digest, updated_at) VALUES (?, ?, ?, ?, ?)
         ON CONFLICT (repo_id, image, tag) DO UPDATE SET digest = excluded.digest, updated_at = excluded.updated_at`,
      ).run(repoId, image, reference, digest, Date.now());
    }
  });
  return { digest, subject };
}

/** Deletes one manifest (and the tags pointing at it). */
export function deleteManifest(repoId: number, image: string, digest: string) {
  const res = db.prepare("DELETE FROM registry_manifests WHERE repo_id = ? AND image = ? AND digest = ?").run(repoId, image, digest);
  pruneRegistry();
  return res.changes > 0;
}

export function deleteTag(repoId: number, image: string, tag: string) {
  return db.prepare("DELETE FROM registry_tags WHERE repo_id = ? AND image = ? AND tag = ?").run(repoId, image, tag).changes > 0;
}

/** Tags in lexical order, paginated as in the distribution spec (`n` per page, after `last`). */
export function listTags(repoId: number, image: string, opts: { n?: number; last?: string } = {}) {
  const rows = db
    .prepare(
      `SELECT tag FROM registry_tags WHERE repo_id = ? AND image = ? AND tag > ? ORDER BY tag COLLATE BINARY LIMIT ?`,
    )
    .all(repoId, image, opts.last ?? "", opts.n ?? -1) as { tag: string }[];
  return rows.map((r) => r.tag);
}

/** Manifests whose subject is `digest` (signatures, SBOMs, attestations), for the referrers API. */
export function listReferrers(repoId: number, image: string, digest: string, artifactType: string | null) {
  const rows = db
    .prepare(
      `SELECT digest, media_type, content, artifact_type FROM registry_manifests
       WHERE repo_id = ? AND image = ? AND subject = ? AND (? IS NULL OR artifact_type = ?) ORDER BY created_at`,
    )
    .all(repoId, image, digest, artifactType, artifactType) as ManifestRow[];
  return rows.map((r) => {
    const { annotations } = parse(r.content);
    return {
      mediaType: r.media_type,
      digest: r.digest,
      size: r.content.byteLength,
      ...(r.artifact_type ? { artifactType: r.artifact_type } : {}),
      ...(annotations ? { annotations } : {}),
    };
  });
}

// ---------------------------------------------------------------------------
// Garbage collection

/** Blobs pushed this recently are kept even if no manifest references them yet: the push may still be running. */
const BLOB_GRACE_MS = 24 * 3600e3;
const UPLOAD_TTL_MS = 24 * 3600e3;
const PRUNE_INTERVAL_MS = 3600e3;

/**
 * Drops blob links no manifest of the repository needs, then deletes blobs no repository links to,
 * and abandoned uploads. Runs synchronously so it can't interleave with an upload being stored.
 */
export function pruneRegistry() {
  g.__registryPrunedAt = Date.now();
  const now = Date.now();
  const orphans = transaction(() => {
    db.prepare(
      `DELETE FROM registry_repo_blobs WHERE created_at < ? AND NOT EXISTS (
         SELECT 1 FROM registry_manifest_refs r WHERE r.repo_id = registry_repo_blobs.repo_id AND r.digest = registry_repo_blobs.digest)`,
    ).run(now - BLOB_GRACE_MS);
    return db
      .prepare("DELETE FROM registry_blobs WHERE digest NOT IN (SELECT digest FROM registry_repo_blobs) RETURNING digest")
      .all() as { digest: string }[];
  });
  for (const { digest } of orphans) rmSync(blobPath(digest), { force: true });

  const stale = db.prepare("DELETE FROM registry_uploads WHERE updated_at < ? RETURNING id").all(now - UPLOAD_TTL_MS) as { id: string }[];
  for (const { id } of stale) uploadHashes.delete(id);
  // Upload files without a row: abandoned, or their repository was deleted.
  const live = new Set((db.prepare("SELECT id FROM registry_uploads").all() as { id: string }[]).map((r) => r.id));
  let files: string[] = [];
  try {
    files = readdirSync(UPLOADS_ROOT);
  } catch {}
  for (const file of files) {
    if (live.has(file)) continue;
    try {
      if (statSync(uploadPath(file)).mtimeMs < now - 60e3) rmSync(uploadPath(file), { force: true });
    } catch {}
  }
}

function maybePrune() {
  if (Date.now() - (g.__registryPrunedAt ?? 0) > PRUNE_INTERVAL_MS) pruneRegistry();
}

// ---------------------------------------------------------------------------
// Listing for the web UI

export type Platform = { os: string; architecture: string; variant?: string };

export type ImageVersion = {
  digest: string;
  mediaType: string;
  tags: { name: string; updatedAt: number }[];
  /** Compressed size of config and layers; for multi-platform images, of the largest platform. */
  size: number;
  platforms: Platform[];
  createdAt: number;
};

export type ImageSummary = { image: string; versions: ImageVersion[]; updatedAt: number };

export function countImages(repoId: number) {
  return (db.prepare("SELECT COUNT(DISTINCT image) AS n FROM registry_manifests WHERE repo_id = ?").get(repoId) as { n: number }).n;
}

/**
 * The repository's images with their versions: manifests that aren't part of a multi-platform
 * index and aren't attached to another manifest (signatures, attestations), newest first.
 */
export function listImages(repoId: number): ImageSummary[] {
  const manifests = db
    .prepare("SELECT image, digest, media_type, content, size, subject, created_at FROM registry_manifests WHERE repo_id = ? ORDER BY created_at DESC")
    .all(repoId) as (ManifestRow & { image: string; size: number; subject: string | null; created_at: number })[];
  const tags = db.prepare("SELECT image, tag, digest, updated_at FROM registry_tags WHERE repo_id = ? ORDER BY updated_at DESC").all(repoId) as {
    image: string;
    tag: string;
    digest: string;
    updated_at: number;
  }[];
  const refs = db.prepare("SELECT image, digest FROM registry_manifest_refs WHERE repo_id = ?").all(repoId) as { image: string; digest: string }[];

  const key = (image: string, digest: string) => `${image}\n${digest}`;
  const referenced = new Set(refs.map((r) => key(r.image, r.digest)));
  const sizes = new Map(manifests.map((m) => [key(m.image, m.digest), m.size]));
  const images = new Map<string, ImageSummary>();

  for (const m of manifests) {
    const k = key(m.image, m.digest);
    const versionTags = tags.filter((t) => t.image === m.image && t.digest === m.digest);
    if (versionTags.length === 0 && (m.subject || referenced.has(k))) continue;

    const json = parse(m.content);
    const platforms = (json.manifests ?? [])
      .filter((d) => d.platform?.os && d.platform.os !== "unknown" && d.platform.architecture !== "unknown")
      .map((d) => ({ os: d.platform!.os!, architecture: d.platform!.architecture!, ...(d.platform!.variant ? { variant: d.platform!.variant } : {}) }));
    const size = json.manifests
      ? Math.max(0, ...json.manifests.filter(isDescriptor).map((d) => sizes.get(key(m.image, d.digest)) ?? 0))
      : m.size;

    const summary = images.get(m.image) ?? { image: m.image, versions: [], updatedAt: 0 };
    images.set(m.image, summary);
    const updatedAt = Math.max(m.created_at, ...versionTags.map((t) => t.updated_at));
    summary.updatedAt = Math.max(summary.updatedAt, updatedAt);
    summary.versions.push({
      digest: m.digest,
      mediaType: m.media_type,
      tags: versionTags.map((t) => ({ name: t.tag, updatedAt: t.updated_at })),
      size,
      platforms,
      createdAt: m.created_at,
    });
  }
  for (const summary of images.values()) {
    summary.versions.sort((a, b) => Number(b.tags.length > 0) - Number(a.tags.length > 0) || b.createdAt - a.createdAt);
  }
  return [...images.values()].sort((a, b) => a.image.localeCompare(b.image));
}

/**
 * Deletes a version from the web UI: the manifest, the platform manifests only it used, and
 * the signatures and attestations attached to any of them.
 */
export function deleteImageVersion(repoId: number, image: string, digest: string) {
  transaction(() => {
    const doomed = new Set([digest]);
    const children = db
      .prepare("SELECT digest FROM registry_manifest_refs WHERE repo_id = ? AND image = ? AND manifest_digest = ?")
      .all(repoId, image, digest) as { digest: string }[];
    for (const { digest: child } of children) {
      const tagged = db.prepare("SELECT 1 FROM registry_tags WHERE repo_id = ? AND image = ? AND digest = ?").get(repoId, image, child);
      const sharedWith = db
        .prepare("SELECT 1 FROM registry_manifest_refs WHERE repo_id = ? AND image = ? AND digest = ? AND manifest_digest != ?")
        .get(repoId, image, child, digest);
      if (!tagged && !sharedWith) doomed.add(child);
    }
    const attached = db.prepare("SELECT digest FROM registry_manifests WHERE repo_id = ? AND image = ? AND subject = ?");
    for (const d of doomed) for (const r of attached.all(repoId, image, d) as { digest: string }[]) doomed.add(r.digest);
    const del = db.prepare("DELETE FROM registry_manifests WHERE repo_id = ? AND image = ? AND digest = ?");
    for (const d of doomed) del.run(repoId, image, d);
  });
  pruneRegistry();
}