import { mergeBudget, mergeRootTrees } from "./git/merge";
import type { MergeSource } from "./git/merge";
import { MergeTooComplexError } from "./git/merge3";
import { makeCommit, signatureNow } from "./git/objects";
import type { RawObject } from "./git/objects";
import { compareCommits } from "./git/history";
import type { GitStore } from "./git/store";
import { pushToArtifacts, PushRejectedError } from "./git/writer";
import { currentBaseSha, loadPullRepos, withPullStore } from "./pull-merge";
import type { Landed } from "./pull-merge";
import { MergeError } from "./pulls";
import type { PullRow } from "./pulls";
import type { CommitSummary } from "~/lib/types";

const MAX_REBASE_COMMITS = 100;
const MAX_REBASE_TREE_READS = 3000;
const MAX_REBASE_CONTENT_MERGES = 400;
const SHA_BYTES = 20;
const decoder = new TextDecoder();

type Entry = { name: string; mode: string; hash: string };

function hex(bytes: Uint8Array): string {
  return Array.from(bytes, (byte) => byte.toString(16).padStart(2, "0")).join("");
}

function parseTree(content: Uint8Array): Entry[] {
  const entries: Entry[] = [];
  let offset = 0;
  while (offset < content.byteLength) {
    const space = content.indexOf(0x20, offset);
    const nul = content.indexOf(0, space);
    const mode = decoder.decode(content.subarray(offset, space));
    const name = decoder.decode(content.subarray(space + 1, nul));
    entries.push({ name, mode, hash: hex(content.subarray(nul + 1, nul + 1 + SHA_BYTES)) });
    offset = nul + 1 + SHA_BYTES;
  }
  return entries;
}

/** Reads objects created earlier in the rebase from memory before falling back to the repository. */
function overlay(store: GitStore, created: Map<string, RawObject>): MergeSource {
  return {
    async tree(sha) {
      const object = created.get(sha);
      return object ? parseTree(object.content) : store.tree(sha);
    },
    async blob(sha) {
      return created.get(sha)?.content ?? store.blob(sha);
    },
  };
}

export class RebaseConflictError extends Error {}

export type RebasePlan =
  | { kind: "up_to_date" }
  | { kind: "fast_forward" }
  | { kind: "rebased"; sha: string; objects: RawObject[] };

/**
 * Replays the head branch's commits onto `baseSha` with the in-Worker three-way merge. Fails closed:
 * any conflict aborts, and the result must match a direct merge of base and head or it is discarded.
 * Returns null for histories git should handle (merge commits, too many commits, merge budgets).
 *
 * With `upstreamSha` this is `git rebase --onto baseSha upstreamSha head`: only commits not reachable
 * from upstream are replayed, which is how a stacked branch drops its parent's pre-squash commits.
 */
export async function planRebase(
  store: GitStore,
  baseSha: string,
  headSha: string,
  /** `timezone` is the merging user's browser offset; replayed commits use it since the original is unknown. */
  committer: { name: string; email: string; timezone?: string },
  labels: { ours: string; theirs: string },
  upstreamSha: string = baseSha,
): Promise<RebasePlan | null> {
  const comparison = await compareCommits(store, upstreamSha, store, headSha);
  if (comparison.mergeBase === null || comparison.truncated) return null;
  if (comparison.mergeBase === headSha || comparison.commits.length === 0) return { kind: "up_to_date" };
  const commits: CommitSummary[] = [...comparison.commits].reverse();
  if (upstreamSha === baseSha) {
    if (comparison.mergeBase === baseSha) return { kind: "fast_forward" };
  } else if (commits[0]?.parents[0] === baseSha) {
    return { kind: "up_to_date" };
  }
  if (commits.length > MAX_REBASE_COMMITS || commits.some((commit) => commit.parents.length !== 1)) return null;
  const [base, head, mergeBase] = await Promise.all([
    store.commit(baseSha),
    store.commit(headSha),
    store.commit(comparison.mergeBase),
  ]);
  if (!base || !head || !mergeBase) return null;
  const created = new Map<string, RawObject>();
  const source = overlay(store, created);
  const now = signatureNow(committer.name, committer.email, committer.timezone);
  let current = baseSha;
  let currentTree = base.treeSha;
  const budget = mergeBudget(MAX_REBASE_TREE_READS, MAX_REBASE_CONTENT_MERGES);
  try {
    for (const commit of commits) {
      const parent = await store.commit(commit.parents[0] ?? "");
      if (!parent) return null;
      const result = await mergeRootTrees({
        source,
        baseTree: parent.treeSha,
        oursTree: currentTree,
        theirsTree: commit.treeSha,
        labels,
        budget,
      });
      if (result.conflicts.length > 0) throw new RebaseConflictError(`Conflict replaying ${commit.shortSha}`);
      for (const object of result.objects) created.set(object.sha, object);
      const rewritten = await makeCommit({
        tree: result.treeSha,
        parents: [current],
        author: {
          name: commit.authorName,
          email: commit.authorEmail,
          timestamp: Math.floor(commit.authoredAt / 1000),
          timezone: now.timezone,
        },
        committer: now,
        message: commit.message,
      });
      created.set(rewritten.sha, rewritten);
      current = rewritten.sha;
      currentTree = result.treeSha;
    }
    const direct = await mergeRootTrees({
      source: store,
      baseTree: mergeBase.treeSha,
      oursTree: base.treeSha,
      theirsTree: head.treeSha,
      labels,
      budget,
    });
    if (direct.conflicts.length > 0 || direct.treeSha !== currentTree) return null;
  } catch (error) {
    if (error instanceof MergeTooComplexError) return null;
    throw error;
  }
  return { kind: "rebased", sha: current, objects: [...created.values()] };
}

/** Lands a rebase merge in the Worker with a compare-and-swap push; null defers to the git container. */
export async function rebaseInWorker(input: {
  pull: PullRow;
  merger: { name: string; email: string; timezone?: string };
}): Promise<Landed | null> {
  const { pull } = input;
  const { baseRepo } = await loadPullRepos(pull);
  const before = await currentBaseSha(baseRepo, pull.base_ref);
  try {
    return await withPullStore(pull, "base", [before, pull.head_sha], async (store) => {
      const plan = await planRebase(store, before, pull.head_sha, input.merger, {
        ours: pull.base_ref,
        theirs: pull.head_ref,
      });
      if (!plan) return null;
      if (plan.kind === "up_to_date") throw new MergeError("The base branch already contains these changes");
      const after = plan.kind === "fast_forward" ? pull.head_sha : plan.sha;
      await pushToArtifacts(
        baseRepo,
        [{ ref: `refs/heads/${pull.base_ref}`, oldSha: before, newSha: after }],
        plan.kind === "rebased" ? plan.objects : [],
      );
      return { before, after };
    });
  } catch (error) {
    if (error instanceof PushRejectedError) throw new MergeError("The base branch moved during the merge; try again");
    throw error;
  }
}