import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";

async function git(pi: ExtensionAPI, cwd: string, args: string[]) {
  return pi.exec("git", args, { cwd });
}

function firstNonEmpty(...values: Array<string | undefined>): string | undefined {
  return values.map((value) => value?.trim()).find(Boolean);
}

async function findDefaultBranch(pi: ExtensionAPI, cwd: string): Promise<string | undefined> {
  // Prefer the remote's configured default branch when available.
  const originHead = await git(pi, cwd, ["symbolic-ref", "--quiet", "--short", "refs/remotes/origin/HEAD"]);
  if (originHead.code === 0) {
    const branch = originHead.stdout.trim().replace(/^origin\//, "");
    if (branch) return branch;
  }

  // Fall back to common local default branch names.
  const branches = await git(pi, cwd, ["branch", "--list", "main", "master"]);
  if (branches.code !== 0) return undefined;

  const branchList = branches.stdout
    .split("\n")
    .map((branch) => branch.replace("*", "").trim())
    .filter(Boolean);

  if (branchList.includes("main")) return "main";
  if (branchList.includes("master")) return "master";
  return undefined;
}

export default function doneExtension(pi: ExtensionAPI) {
  pi.registerCommand("done", {
    description: "Return to the default git branch, pull latest changes, and start a new pi session",
    handler: async (_args, ctx) => {
      const status = await git(pi, ctx.cwd, ["status", "--porcelain"]);
      if (status.code !== 0) {
        ctx.ui.notify("/done must be run inside a git repository.", "error");
        return;
      }

      const dirty = status.stdout.trim();
      if (dirty) {
        ctx.ui.notify(
          `Uncommitted changes detected:\n${dirty}\n\nCommit or stash before running /done.`,
          "error",
        );
        return;
      }

      const defaultBranch = await findDefaultBranch(pi, ctx.cwd);
      if (!defaultBranch) {
        ctx.ui.notify("Could not find a default branch. Expected origin/HEAD, main, or master.", "error");
        return;
      }

      const checkout = await git(pi, ctx.cwd, ["checkout", defaultBranch]);
      if (checkout.code !== 0) {
        ctx.ui.notify(
          `Failed to checkout ${defaultBranch}: ${firstNonEmpty(checkout.stderr, checkout.stdout) ?? "unknown error"}`,
          "error",
        );
        return;
      }

      const pull = await git(pi, ctx.cwd, ["pull", "--ff-only"]);
      if (pull.code !== 0) {
        ctx.ui.notify(
          `Checked out ${defaultBranch}, but failed to pull latest changes: ${
            firstNonEmpty(pull.stderr, pull.stdout) ?? "unknown error"
          }`,
          "error",
        );
        return;
      }

      ctx.ui.notify(`Ready on ${defaultBranch}. Starting a fresh pi session...`, "info");
      await ctx.newSession();
    },
  });
}