import { env } from "cloudflare:workers";
import { MAX_IN_LIST, all, batch, chunks, first, placeholders, run, stmt } from "../db";
import { accountDb, repoDb, stmt as scoped } from "../scoped-db";
import { HttpError } from "../errors";
import { randomId } from "../crypto";
import { hasPermission, loadRepoById, repoPermission } from "../repos";
import type { RepoRow, Viewer } from "~/lib/types";

/** Forks read LFS objects from their parents; bound the chain so lookups stay O(1) in D1 calls. */
const MAX_FORK_DEPTH = 3;
/** Unreferenced blobs survive this long, so an upload racing the sweeper never loses its data. */
const ORPHAN_GRACE_MS = 60 * 60 * 1000;
const SWEEP_BATCH = 100;

type ObjectRow = { repo_id: string; size: number };

type ChainRepo = Pick<RepoRow, "id" | "fork_of_id" | "owner_id" | "visibility">;

/**
 * Whether readers of `repo` may see objects stored in the private `ancestor`: only for a private fork,
 * and only while the fork's owner (the account whose collaborators can read it) can still read the
 * ancestor, so a removed collaborator's fork stops seeing new objects.
 */
async function ownerCanRead(repo: ChainRepo, ancestorId: string): Promise<boolean> {
  if (repo.visibility === "public") return false;
  const ancestor = await loadRepoById(ancestorId);
  if (!ancestor) return false;
  if (ancestor.owner_id === repo.owner_id) return true;
  const owner: Viewer = {
    id: repo.owner_id,
    name: "",
    displayName: null,
    email: "",
    isAdmin: false,
    emailVerified: true,
  };
  return hasPermission(await repoPermission(ancestor, owner), "read");
}

/**
 * The repo plus the ancestors whose objects its readers may see: public ancestors, or private ones the
 * fork's owner can still read when the fork is private too. Stops at the first ancestor that fails.
 */
async function repoChain(repo: ChainRepo): Promise<string[]> {
  const chain = [repo.id];
  let parent = repo.fork_of_id;
  while (parent && chain.length < MAX_FORK_DEPTH) {
    const row = await first<{ fork_of_id: string | null; visibility: string }>(
      "SELECT fork_of_id, visibility FROM repos WHERE id = ?",
      parent,
    );
    if (!row) break;
    if (row.visibility !== "public" && !(await ownerCanRead(repo, parent))) break;
    chain.push(parent);
    parent = row.fork_of_id;
  }
  return chain;
}

/** Finds which repo in the fork chain stores the object, preferring the repo itself. */
export async function findObject(repo: ChainRepo, oid: string): Promise<ObjectRow | null> {
  const own = await repoDb(repo.id).first<ObjectRow>(
    "SELECT repo_id, size FROM lfs_objects WHERE repo_id = ? AND oid = ?",
    repo.id,
    oid,
  );
  if (own || !repo.fork_of_id) return own;
  for (const id of (await repoChain(repo)).slice(1)) {
    const row = await repoDb(id).first<ObjectRow>(
      "SELECT repo_id, size FROM lfs_objects WHERE repo_id = ? AND oid = ?",
      id,
      oid,
    );
    if (row) return row;
  }
  return null;
}

/** Bulk variant of findObject for batch requests; IN-lists are chunked to stay under the parameter cap. */
export async function findObjects(repo: ChainRepo, oids: string[]): Promise<Map<string, ObjectRow>> {
  const chain = await repoChain(repo);
  const found = new Map<string, ObjectRow>();
  for (const id of chain) {
    const missing = oids.filter((oid) => !found.has(oid));
    for (const group of chunks(missing, MAX_IN_LIST - 1)) {
      const rows = await repoDb(id).all<ObjectRow & { oid: string }>(
        `SELECT repo_id, oid, size FROM lfs_objects WHERE repo_id = ? AND oid IN (${placeholders(group.length)})`,
        id,
        ...group,
      );
      for (const row of rows) found.set(row.oid, { repo_id: row.repo_id, size: row.size });
    }
  }
  return found;
}

/** Opens the single stored copy, once `repo` (or a readable ancestor) links the oid. */
export async function openObject(repo: ChainRepo, oid: string): Promise<R2ObjectBody | null> {
  if (!(await findObject(repo, oid))) return null;
  const blob = await first<{ r2_key: string }>("SELECT r2_key FROM lfs_blobs WHERE oid = ? AND refs > 0", oid);
  return blob ? env.BUCKET.get(blob.r2_key) : null;
}

type LinkRepo = Pick<RepoRow, "id" | "owner_id">;
type Link = { oid: string; size: number };

/**
 * What the D1 triggers did when lfs_objects lived next to lfs_blobs: count the new link on the shared copy, and
 * bill the owner once per oid it holds across all its repositories.
 */
async function afterLink(ownerId: string, links: Link[]): Promise<void> {
  if (links.length === 0) return;
  await batch(
    links.map((link) => stmt("UPDATE lfs_blobs SET refs = refs + 1, orphaned_at = NULL WHERE oid = ?", link.oid)),
  );
  await bill(ownerId, links);
}

async function bill(ownerId: string, links: Link[]): Promise<void> {
  if (links.length === 0) return;
  await accountDb(ownerId).batch([
    scoped("INSERT OR IGNORE INTO lfs_usage (owner_id) VALUES (?)", ownerId),
    ...links.flatMap((link) => [
      scoped(
        `INSERT INTO lfs_account_objects (owner_id, oid, size, refs) VALUES (?, ?, ?, 1)
         ON CONFLICT (owner_id, oid) DO UPDATE SET refs = refs + 1`,
        ownerId,
        link.oid,
        link.size,
      ),
      scoped(
        `UPDATE lfs_usage SET storage_bytes = storage_bytes + ? WHERE owner_id = ?
           AND (SELECT refs FROM lfs_account_objects WHERE owner_id = ? AND oid = ?) = 1`,
        link.size,
        ownerId,
        ownerId,
        link.oid,
      ),
    ]),
  ]);
}

async function unbill(ownerId: string, links: Link[]): Promise<void> {
  if (links.length === 0) return;
  await accountDb(ownerId).batch(
    links.flatMap((link) => [
      scoped("UPDATE lfs_account_objects SET refs = refs - 1 WHERE owner_id = ? AND oid = ?", ownerId, link.oid),
      scoped(
        `UPDATE lfs_usage SET storage_bytes = MAX(storage_bytes - ?, 0) WHERE owner_id = ?
           AND (SELECT refs FROM lfs_account_objects WHERE owner_id = ? AND oid = ?) <= 0`,
        link.size,
        ownerId,
        ownerId,
        link.oid,
      ),
      scoped("DELETE FROM lfs_account_objects WHERE owner_id = ? AND oid = ? AND refs <= 0", ownerId, link.oid),
    ]),
  );
}

async function afterUnlink(rows: Array<Link & { owner_id: string }>): Promise<void> {
  if (rows.length === 0) return;
  const now = Date.now();
  for (const group of chunks(rows, SWEEP_BATCH)) {
    await batch(
      group.flatMap((row) => [
        stmt("UPDATE lfs_blobs SET refs = MAX(refs - 1, 0) WHERE oid = ?", row.oid),
        stmt("UPDATE lfs_blobs SET orphaned_at = ? WHERE oid = ? AND refs = 0", now, row.oid),
      ]),
    );
  }
  const byOwner = new Map<string, Link[]>();
  for (const row of rows) byOwner.set(row.owner_id, [...(byOwner.get(row.owner_id) ?? []), row]);
  for (const [ownerId, links] of byOwner) await unbill(ownerId, links);
}

/** Inserts links into the repo object and returns the ones that are new. */
async function insertLinks(repo: LinkRepo, links: Link[]): Promise<Link[]> {
  if (links.length === 0) return [];
  const now = Date.now();
  const results = await repoDb(repo.id).batch(
    links.map((link) =>
      scoped(
        "INSERT OR IGNORE INTO lfs_objects (repo_id, oid, owner_id, size, created_at) VALUES (?, ?, ?, ?, ?)",
        repo.id,
        link.oid,
        repo.owner_id,
        link.size,
        now,
      ),
    ),
  );
  const added = links.filter((_, index) => (results[index]?.meta.changes ?? 0) > 0);
  await afterLink(repo.owner_id, added);
  return added;
}

/** Only links whose stored copy exists with that size and is still referenced (copy-on-write). */
async function liveLinks(objects: Link[]): Promise<Link[]> {
  const live = new Map<string, number>();
  for (const group of chunks(
    objects.map((object) => object.oid),
    MAX_IN_LIST,
  )) {
    const rows = await all<{ oid: string; size: number }>(
      `SELECT oid, size FROM lfs_blobs WHERE refs > 0 AND oid IN (${placeholders(group.length)})`,
      ...group,
    );
    for (const row of rows) live.set(row.oid, row.size);
  }
  return objects.filter((object) => live.get(object.oid) === object.size);
}

/**
 * Links already-stored objects into `repo` without copying them. Only for oids the caller has shown
 * `repo` may read (found in its readable fork chain); the owner is billed once per oid.
 */
export async function linkObjects(repo: LinkRepo, objects: Array<{ oid: string; size: number }>): Promise<void> {
  if (objects.length > 0) await insertLinks(repo, await liveLinks(objects));
}

/** Copy-on-write fork: the fork links every object of its parent, sharing the stored copies. */
export async function linkForkObjects(fork: LinkRepo, parentId: string): Promise<void> {
  const rows = await repoDb(parentId).all<Link>("SELECT oid, size FROM lfs_objects WHERE repo_id = ?", parentId);
  for (const group of chunks(rows, SWEEP_BATCH)) await insertLinks(fork, await liveLinks(group));
}

/**
 * Streams into R2 with a declared SHA-256 so R2 itself rejects content that does not match the oid.
 * Every upload writes a fresh key. The blob row is registered with a reference held for the new link, so a
 * concurrent sweep never sees it unreferenced. When another copy already exists, the fresh one is deleted and
 * the existing copy is shared.
 */
export async function storeObject(
  repo: LinkRepo,
  oid: string,
  size: number,
  body: ReadableStream<Uint8Array>,
): Promise<void> {
  const key = `lfs/blobs/${oid}/${randomId()}`;
  const sized = body.pipeThrough(new FixedLengthStream(size));
  try {
    await env.BUCKET.put(key, sized, { sha256: oid, httpMetadata: { contentType: "application/octet-stream" } });
  } catch (error) {
    console.warn("LFS upload rejected", oid, error);
    throw new HttpError(422, "Uploaded content does not match the declared oid and size");
  }
  const now = Date.now();
  await run(
    `INSERT INTO lfs_blobs (oid, size, r2_key, refs, orphaned_at, created_at) VALUES (?, ?, ?, 0, ?, ?)
     ON CONFLICT (oid) DO NOTHING`,
    oid,
    size,
    key,
    now,
    now,
  );
  const stored = await first<{ r2_key: string; size: number }>("SELECT r2_key, size FROM lfs_blobs WHERE oid = ?", oid);
  if (stored?.size === size) await insertLinks(repo, [{ oid, size }]);
  if (stored?.r2_key !== key) await env.BUCKET.delete(key);
}

/** Unlinks a repository's objects, releasing billing; the sweeper deletes orphaned copies. */
export async function purgeRepoLfs(repo: Pick<RepoRow, "id">): Promise<void> {
  const rows = await repoDb(repo.id).all<Link & { owner_id: string }>(
    "DELETE FROM lfs_objects WHERE repo_id = ? RETURNING oid, size, owner_id",
    repo.id,
  );
  await afterUnlink(rows);
}

/** After a transfer: moves the repository's LFS billing from its previous owner to `newOwnerId`. */
export async function rebillRepoLfs(repoId: string, newOwnerId: string): Promise<void> {
  const rows = await repoDb(repoId).all<Link & { owner_id: string }>(
    "SELECT oid, size, owner_id FROM lfs_objects WHERE repo_id = ? AND owner_id != ?",
    repoId,
    newOwnerId,
  );
  if (rows.length === 0) return;
  await repoDb(repoId).run("UPDATE lfs_objects SET owner_id = ? WHERE repo_id = ?", newOwnerId, repoId);
  const byOwner = new Map<string, Link[]>();
  for (const row of rows) byOwner.set(row.owner_id, [...(byOwner.get(row.owner_id) ?? []), row]);
  for (const [ownerId, links] of byOwner) await unbill(ownerId, links);
  await bill(newOwnerId, rows);
}

/** Cron: deletes stored copies that no repository has linked for the grace period. */
export async function sweepOrphanedLfsBlobs(): Promise<void> {
  const orphans = await all<{ oid: string; r2_key: string }>(
    "DELETE FROM lfs_blobs WHERE refs = 0 AND orphaned_at < ? RETURNING oid, r2_key",
    Date.now() - ORPHAN_GRACE_MS,
  );
  for (const group of chunks(orphans, SWEEP_BATCH)) await env.BUCKET.delete(group.map((orphan) => orphan.r2_key));
}