import argparse
import re
import subprocess
import sys
from pathlib import Path

PULL_REQUEST = re.compile(r"Merge pull request #(\d+)")
BASES = ("upstream/main", "origin/main", "main")


def git(repo, *args, check=True):
    result = subprocess.run(["git", "-C", str(repo), *args], capture_output=True, text=True, encoding="utf-8", errors="replace")
    if check and result.returncode:
        raise RuntimeError(f"git {' '.join(args[:2])} failed: {result.stderr.strip()}")
    return result


def lines(repo, *args):
    return [line for line in git(repo, *args).stdout.splitlines() if line]


def names(repo, *args):
    return [name for name in git(repo, *args).stdout.split("\0") if name]


def resolve_base(repo, name):
    for candidate in (name,) if name else BASES:
        if git(repo, "rev-parse", "--verify", "--quiet", f"{candidate}^{{commit}}", check=False).returncode == 0:
            return candidate
    raise RuntimeError(f"no base ref found ({name or ', '.join(BASES)}); fetch the upstream remote first")


def conflicts(repo, base, head):
    result = git(repo, "merge-tree", "-z", "--write-tree", "--name-only", "--no-messages", base, head, check=False)
    if result.returncode not in (0, 1):
        raise RuntimeError(f"git merge-tree failed: {result.stderr.strip()}")
    if result.returncode == 0:
        return []
    fields = result.stdout.split("\0")[1:]
    return fields[:fields.index("")] if "" in fields else fields


def merged_pull_requests(repo, mergebase, base, paths):
    found = {}
    if not paths:
        return found
    for line in lines(repo, "-c", "core.quotepath=off", "log", "--first-parent", "--format=%h %s", f"{mergebase}..{base}", "--", *paths):
        match = PULL_REQUEST.search(line)
        if match:
            found[int(match.group(1))] = line.split(" ", 1)[0]
    return found


def report(repo, base, head):
    found = git(repo, "merge-base", base, head, check=False)
    if found.returncode:
        raise RuntimeError(f"{head} and {base} have no common ancestor")
    mergebase = found.stdout.strip()
    behind = int(git(repo, "rev-list", "--count", f"{head}..{base}").stdout)
    ahead = int(git(repo, "rev-list", "--count", f"{base}..{head}").stdout)
    mine = set(names(repo, "diff", "-z", "--name-only", mergebase, head))
    theirs = set(names(repo, "diff", "-z", "--name-only", mergebase, base))
    shared = sorted(mine & theirs)
    return {
        "base": base,
        "head": head,
        "mergebase": mergebase,
        "behind": behind,
        "ahead": ahead,
        "conflicts": conflicts(repo, base, head),
        "shared": shared,
        "pulls": merged_pull_requests(repo, mergebase, base, shared),
    }


def render(data):
    out = [f"{data['head']} is {data['ahead']} ahead and {data['behind']} behind {data['base']} (merge base {data['mergebase'][:9]})"]
    if data["conflicts"]:
        out.append(f"merge: CONFLICT in {len(data['conflicts'])} file(s)")
        out += [f"  {name}" for name in data["conflicts"]]
    else:
        out.append("merge: clean")
    if data["shared"]:
        out.append(f"{len(data['shared'])} file(s) changed on both sides since the merge base:")
        out += [f"  {name}" for name in data["shared"]]
        if data["pulls"]:
            out.append("merged pull requests that touched them: " + ", ".join(f"#{number}" for number in sorted(data["pulls"])))
    return "\n".join(out)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="report whether a branch is in step with the upstream main before you push or test")
    parser.add_argument("--base", help="the main ref to compare against (default: upstream/main, then origin/main, then main)")
    parser.add_argument("--head", default="HEAD", help="the branch or commit to check (default: HEAD)")
    parser.add_argument("--repo", type=Path, default=Path("."), help="the repository to check (default: the current directory)")
    parser.add_argument("--fetch", action="store_true", help="fetch the base's remote first")
    parser.add_argument("--strict", action="store_true", help="also fail when the branch is behind the base")
    args = parser.parse_args()
    try:
        base = resolve_base(args.repo, args.base)
        if args.fetch:
            remote, _, branch = base.partition("/")
            if not branch:
                parser.error("--fetch needs a remote-tracking base such as upstream/main")
            git(args.repo, "fetch", "--quiet", remote, branch)
        data = report(args.repo, base, args.head)
    except RuntimeError as error:
        parser.error(str(error))
    sys.stdout.reconfigure(errors="backslashreplace")
    print(render(data))
    if data["conflicts"]:
        raise SystemExit(1)
    if args.strict and data["behind"]:
        raise SystemExit(1)