import "server-only";
import { createReadStream } from "node:fs";
import { stat } from "node:fs/promises";
import { Readable } from "node:stream";
import {
  appendUpload,
  cancelUpload,
  deleteManifest,
  deleteTag,
  findBlob,
  findManifest,
  findUpload,
  finishUpload,
  isDigest,
  isImagePath,
  listReferrers,
  listTags,
  MAX_MANIFEST_BYTES,
  MEDIA_TYPES,
  mountBlob,
  putManifest,
  RegistryError,
  startUpload,
  TAG_RE,
  unlinkBlob,
  type Upload,
} from "./registry";
import { canWrite, findRepo, repoAccess, type RepoRecord } from "./repos";
import { basicAuthToken, tokenCovers } from "./tokens";

// The OCI distribution API (Docker Registry HTTP API v2) under /v2. Clients log in with HTTP Basic
// auth, an API token as the password, like git over HTTPS: `docker login <host> -u <you>`.

const API_HEADERS = { "Docker-Distribution-API-Version": "registry/2.0" };

function reply(status: number, headers: Record<string, string> = {}, body: BodyInit | null = null) {
  return new Response(body, { status, headers: { ...API_HEADERS, ...headers } });
}

function json(status: number, data: unknown, headers: Record<string, string> = {}) {
  return reply(status, { "Content-Type": "application/json", ...headers }, JSON.stringify(data));
}

function error(status: number, code: string, message: string, headers: Record<string, string> = {}) {
  return json(status, { errors: [{ code, message, detail: null }] }, headers);
}

const unauthorized = (message: string) =>
  error(401, "UNAUTHORIZED", message, { "WWW-Authenticate": 'Basic realm="Vectrum"' });

type Target = { repo: RepoRecord; image: string; name: string };
type Credentials = Awaited<ReturnType<typeof basicAuthToken>>;

/**
 * Resolves an image name (<owner>/<repo>[/<path>]) and checks that the request may pull from or
 * push to it. Anonymous users can pull images of public repositories; everything else needs a
 * token that covers the repository, and pushing needs write access to it.
 */
async function authorize(creds: Credentials, name: string, need: "pull" | "push"): Promise<Target | Response> {
  if (creds === "invalid") return unauthorized("Invalid or expired token.");
  const [owner, repoName, ...rest] = name.split("/");
  const image = rest.join("/");
  if (!repoName || !isImagePath(name)) return error(400, "NAME_INVALID", `Invalid image name ${name}. Names are lowercase.`);
  const repo = await findRepo(owner, repoName);

  if (!creds) {
    if (repo && need === "pull" && repo.visibility === "public") return { repo, image, name };
    // Ask for credentials; don't reveal whether private repositories exist.
    return unauthorized("Authentication required. Log in with an API token as the password.");
  }

  const access = repo ? await repoAccess(repo, creds.user.id) : null;
  if (!repo || !access) {
    return error(404, "NAME_UNKNOWN", `No repository ${owner}/${repoName}. Images are pushed to an existing repository.`);
  }
  const pushing = need === "push";
  if (pushing && creds.token.permission !== "write") {
    return error(403, "DENIED", "This token is read-only. Create a read & write token to push.");
  }
  if (!tokenCovers(creds.token, repo.id, pushing)) {
    return error(403, "DENIED", "This token does not grant access to this repository.");
  }
  if (pushing && !canWrite(access)) return error(403, "DENIED", `Permission to ${repo.owner}/${repo.name} denied.`);
  return { repo, image, name };
}

const ROUTES = [
  { kind: "uploads", re: /^\/(.+)\/blobs\/uploads\/?$/ },
  { kind: "upload", re: /^\/(.+)\/blobs\/uploads\/([^/]+)$/ },
  { kind: "blob", re: /^\/(.+)\/blobs\/([^/]+)$/ },
  { kind: "manifest", re: /^\/(.+)\/manifests\/([^/]+)$/ },
  { kind: "tags", re: /^\/(.+)\/tags\/list$/ },
  { kind: "referrers", re: /^\/(.+)\/referrers\/([^/]+)$/ },
] as const;

type Kind = (typeof ROUTES)[number]["kind"];

const PUSH_METHODS = new Set(["POST", "PUT", "PATCH", "DELETE"]);
const ALLOWED: Record<Kind, string[]> = {
  uploads: ["POST"],
  upload: ["GET", "PATCH", "PUT", "DELETE"],
  blob: ["GET", "HEAD", "DELETE"],
  manifest: ["GET", "HEAD", "PUT", "DELETE"],
  tags: ["GET"],
  referrers: ["GET"],
};

/** Handles every request under /v2. */
export async function handleRegistry(request: Request) {
  const url = new URL(request.url);
  const rest = url.pathname.replace(/^\/v2/, "");
  const creds = await basicAuthToken(request);

  if (rest === "" || rest === "/") {
    if (request.method !== "GET" && request.method !== "HEAD") return error(405, "UNSUPPORTED", "Method not allowed.");
    // `docker login` checks its credentials here, so anonymous requests are always challenged.
    if (!creds) return unauthorized("Authentication required. Log in with an API token as the password.");
    if (creds === "invalid") return unauthorized("Invalid or expired token.");
    return json(200, {});
  }

  for (const { kind, re } of ROUTES) {
    const match = re.exec(rest);
    if (!match) continue;
    const [name, param = ""] = match.slice(1).map(decodeURIComponent);
    if (!ALLOWED[kind].includes(request.method)) return error(405, "UNSUPPORTED", "Method not allowed.");
    const target = await authorize(creds, name, PUSH_METHODS.has(request.method) ? "push" : "pull");
    if (target instanceof Response) return target;
    try {
      switch (kind) {
        case "uploads":
          return await startBlobUpload(request, url, target, creds);
        case "upload":
          return await continueBlobUpload(request, url, target, param);
        case "blob":
          return await blob(request, target, param);
        case "manifest":
          return await manifest(request, target, param);
        case "tags":
          return tags(url, target);
        case "referrers":
          return referrers(url, target, param);
      }
    } catch (e) {
      if (e instanceof RegistryError) return error(e.status, e.code, e.message);
      throw e;
    }
  }
  return error(404, "NAME_UNKNOWN", "Not found.");
}

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

async function blob(request: Request, { repo }: Target, digest: string) {
  if (!isDigest(digest)) return error(400, "DIGEST_INVALID", "Invalid digest.");
  if (request.method === "DELETE") {
    return unlinkBlob(repo.id, digest) ? reply(202) : error(404, "BLOB_UNKNOWN", "Blob not found.");
  }
  const found = findBlob(repo.id, digest);
  const size = found && (await stat(found.path).catch(() => null))?.size;
  if (!found || size == null) return error(404, "BLOB_UNKNOWN", "Blob not found.");

  const headers: Record<string, string> = {
    "Content-Type": "application/octet-stream",
    "Docker-Content-Digest": digest,
    "Accept-Ranges": "bytes",
    // Content-addressed, so it never changes.
    "Cache-Control": "public, max-age=31536000, immutable",
  };
  let start = 0;
  let end = size - 1;
  const range = request.headers.get("range");
  if (range) {
    const m = /^bytes=(\d*)-(\d*)$/.exec(range.trim());
    if (m && (m[1] || m[2])) {
      start = m[1] ? Number(m[1]) : Math.max(0, size - Number(m[2]));
      end = m[1] && m[2] ? Math.min(Number(m[2]), size - 1) : size - 1;
    }
    if (!m || start > end || start >= size) return reply(416, { ...headers, "Content-Range": `bytes */${size}` });
    headers["Content-Range"] = `bytes ${start}-${end}/${size}`;
  }
  headers["Content-Length"] = String(end - start + 1);
  const status = range ? 206 : 200;
  if (request.method === "HEAD" || size === 0) return reply(status, headers);
  const stream = Readable.toWeb(createReadStream(found.path, { start, end })) as ReadableStream<Uint8Array>;
  return reply(status, headers, stream);
}

const uploadHeaders = ({ name }: Target, upload: Upload, size: number) => ({
  Location: `/v2/${name}/blobs/uploads/${upload.id}`,
  Range: `0-${Math.max(0, size - 1)}`,
  "Docker-Upload-UUID": upload.id,
});

const blobCreated = ({ name }: Target, digest: string) =>
  reply(201, { Location: `/v2/${name}/blobs/${digest}`, "Docker-Content-Digest": digest });

/** POST /v2/<name>/blobs/uploads/: mounts a blob, takes a whole blob, or starts a chunked upload. */
async function startBlobUpload(request: Request, url: URL, target: Target, creds: Credentials) {
  const mount = url.searchParams.get("mount");
  const from = url.searchParams.get("from");
  if (mount && from && isDigest(mount)) {
    const source = await authorize(creds, from, "pull");
    if (!(source instanceof Response) && mountBlob(target.repo.id, source.repo.id, mount)) return blobCreated(target, mount);
    // Not mountable: the client uploads the blob instead.
  }

  const upload = startUpload(target.repo.id);
  const digest = url.searchParams.get("digest");
  if (digest) {
    if (!isDigest(digest)) {
      cancelUpload(upload);
      return error(400, "DIGEST_INVALID", "Invalid digest.");
    }
    await appendOrCancel(upload, request.body);
    await finishUpload(upload, digest);
    return blobCreated(target, digest);
  }
  return reply(202, { ...uploadHeaders(target, upload, 0), "Content-Length": "0" });
}

async function appendOrCancel(upload: Upload, body: ReadableStream<Uint8Array> | null) {
  try {
    return await appendUpload(upload, body);
  } catch (e) {
    cancelUpload(upload);
    throw e;
  }
}

/** /v2/<name>/blobs/uploads/<id>: status, a chunk, the final chunk with the digest, or cancel. */
async function continueBlobUpload(request: Request, url: URL, target: Target, id: string) {
  const upload = findUpload(target.repo.id, id);
  if (!upload) return error(404, "BLOB_UPLOAD_UNKNOWN", "Upload not found. It may have expired.");

  switch (request.method) {
    case "GET":
      return reply(204, uploadHeaders(target, upload, upload.size));
    case "DELETE":
      cancelUpload(upload);
      return reply(204);
    case "PATCH": {
      // Chunks must arrive in order: "Content-Range: <start>-<end>" has to continue where the upload ends.
      const range = /^(?:bytes )?(\d+)-(\d+)/.exec(request.headers.get("content-range") ?? "");
      if (range && Number(range[1]) !== upload.size) {
        return error(416, "BLOB_UPLOAD_INVALID", "Chunk out of order.", uploadHeaders(target, upload, upload.size));
      }
      const size = await appendUpload(upload, request.body);
      return reply(202, { ...uploadHeaders(target, upload, size), "Content-Length": "0" });
    }
    default: {
      const digest = url.searchParams.get("digest");
      if (!digest || !isDigest(digest)) return error(400, "DIGEST_INVALID", "A valid digest is required to finish an upload.");
      const size = await appendUpload(upload, request.body);
      await finishUpload({ ...upload, size }, digest);
      return blobCreated(target, digest);
    }
  }
}

// ---------------------------------------------------------------------------
// Manifests, tags, referrers

/** Reads a request body, refusing more than `limit` bytes. */
async function readBody(request: Request, limit: number) {
  if (Number(request.headers.get("content-length")) > limit) return null;
  const chunks: Uint8Array[] = [];
  let length = 0;
  if (!request.body) return Buffer.alloc(0);
  for await (const chunk of Readable.fromWeb(request.body as never) as AsyncIterable<Uint8Array>) {
    length += chunk.byteLength;
    if (length > limit) return null;
    chunks.push(chunk);
  }
  return Buffer.concat(chunks);
}

async function manifest(request: Request, target: Target, reference: string) {
  const { repo, image, name } = target;
  const valid = isDigest(reference) || TAG_RE.test(reference);

  if (request.method === "PUT") {
    if (!valid) return error(400, "TAG_INVALID", "Invalid tag.");
    const content = await readBody(request, MAX_MANIFEST_BYTES);
    if (!content) return error(413, "SIZE_INVALID", "Manifests can be at most 4 MB.");
    const { digest, subject } = putManifest(repo.id, image, reference, request.headers.get("content-type"), content);
    return reply(201, {
      Location: `/v2/${name}/manifests/${digest}`,
      "Docker-Content-Digest": digest,
      ...(subject ? { "OCI-Subject": subject } : {}),
    });
  }

  if (request.method === "DELETE") {
    // A digest deletes the manifest and its tags; a tag only removes the tag.
    const deleted = isDigest(reference) ? deleteManifest(repo.id, image, reference) : valid && deleteTag(repo.id, image, reference);
    return deleted ? reply(202) : error(404, "MANIFEST_UNKNOWN", "Manifest not found.");
  }

  const found = valid ? findManifest(repo.id, image, reference) : null;
  if (!found) return error(404, "MANIFEST_UNKNOWN", `Manifest ${reference} not found.`);
  const headers = {
    "Content-Type": found.mediaType,
    "Content-Length": String(found.content.byteLength),
    "Docker-Content-Digest": found.digest,
  };
  return reply(200, headers, request.method === "HEAD" ? null : Buffer.from(found.content));
}

/** GET /v2/<name>/tags/list?n=<count>&last=<tag> */
function tags(url: URL, { repo, image, name }: Target) {
  const nParam = url.searchParams.get("n");
  const n = nParam === null ? undefined : Math.max(0, Math.floor(Number(nParam)) || 0);
  const last = url.searchParams.get("last") ?? undefined;
  const list = n === 0 ? [] : listTags(repo.id, image, { n, last });
  const headers: Record<string, string> = {};
  if (n && list.length === n) {
    const next = new URLSearchParams({ n: String(n), last: list[list.length - 1] });
    headers.Link = `</v2/${name}/tags/list?${next}>; rel="next"`;
  }
  return json(200, { name, tags: list }, headers);
}

/** GET /v2/<name>/referrers/<digest>?artifactType=<type> */
function referrers(url: URL, { repo, image }: Target, digest: string) {
  if (!isDigest(digest)) return error(400, "DIGEST_INVALID", "Invalid digest.");
  const artifactType = url.searchParams.get("artifactType");
  const index = { schemaVersion: 2, mediaType: MEDIA_TYPES.ociIndex, manifests: listReferrers(repo.id, image, digest, artifactType) };
  return json(200, index, {
    "Content-Type": MEDIA_TYPES.ociIndex,
    ...(artifactType ? { "OCI-Filters-Applied": "artifactType" } : {}),
  });
}