#!/usr/bin/env python3
"""
Compare simple benchmark timings between commits using GitPython.

This script keeps repository objects local to `main()` so it fails fast when
not run from the llg3d repository.
"""

import argparse
import os
import re
import shlex
import subprocess
import sys
import logging
from pathlib import Path


from .utils import ChdirTemporaryDirectory
from git import Repo


def get_llg3d_repo() -> Repo:
    """Return the llg3d GitPython Repo object."""
    repo = Repo(Path(__file__).resolve(), search_parent_directories=True)
    if repo.working_tree_dir is None:
        raise RuntimeError("Could not determine repository working tree")
    root = Path(repo.working_tree_dir)

    if (root / "src" / "llg3d").exists():
        return repo

    py = root / "pyproject.toml"
    if py.exists():
        try:
            import tomllib

            with py.open("rb") as fh:
                data = tomllib.load(fh)
                if data.get("project", {}).get("name") == "llg3d":
                    return repo
        except Exception:
            pass

    try:
        for remote in repo.remotes:
            try:
                url = remote.url
            except Exception:
                url = repo.git.config(f"--get remote.{remote.name}.url")
            if url and "llg3d" in str(url):
                return repo
    except Exception:
        pass

    logging.error("This script must be run inside the llg3d repository.")
    logging.error(f"Detected repo root: {root}")
    sys.exit(2)


def get_current_commit(repo: Repo) -> str:
    """Return the current commit hexsha for `repo`."""
    return repo.head.commit.hexsha


def get_current_ref(repo: Repo) -> str:
    """Return the current branch name, or the commit if HEAD is detached."""
    if repo.head.is_detached:
        return get_current_commit(repo)
    return repo.active_branch.name


def ensure_clean_worktree(repo: Repo) -> None:
    """Exit if the worktree contains changes or untracked files."""
    if repo.is_dirty(untracked_files=True):
        logging.error("Working tree is not clean. Stash or commit changes first.")
        logging.error("Hint: git stash -u")
        sys.exit(1)


def checkout_commit(repo: Repo, commit: str) -> None:
    """Checkout the given `commit` in `repo`."""
    repo.git.checkout(commit)


def run_single_benchmark(
    solver: str,
    Jx: int,
    Jy: int,
    Jz: int,
    n_mean: int,
    N: int,
    launcher: str | None = None,
    np_procs: int = 1,
    precision: str = "single",
    repo_root: Path | None = None,
) -> dict:
    """Run a single benchmark and return a dict with timing results."""
    if repo_root is None:
        raise RuntimeError("repo_root must be provided")

    arg_new = (
        f"--element Cobalt --N {N} --dt 1.0e-14 --Jx {Jx} --Jy {Jy} --Jz {Jz} "
        f"--dx 1.0e-9 --T 1100.0 --H_ext 0.0 --solver {solver} --precision {precision} "
        f"--n_mean {n_mean} --n_profile 0 --seed 42"
    )
    arg_old = (
        f"-element Cobalt -N {N} -dt 1.0e-14 -Jx {Jx} -Jy {Jy} -Jz {Jz} "
        f"-dx 1.0e-9 -T 1100.0 -H_ext 0.0 -n_average {n_mean} -n_profile 0"
    )

    candidates = [
        ["python", "-m", "llg3d"],
        ["python", "-m", "llg3d.__main__"],
        ["python", "-m", "llg3d.main"],
        ["python", "-m", "llg3d.llg3d"],
    ]

    effective_launcher = launcher
    if solver == "mpi" and not effective_launcher:
        effective_launcher = f"mpirun -np {np_procs}"

    env = os.environ.copy()
    env["PYTHONPATH"] = f"{repo_root / 'src'}" + (":" + env.get("PYTHONPATH", ""))

    last_err = None
    output = ""

    with ChdirTemporaryDirectory() as tmpdir:
        for arg_str in (arg_new, arg_old):
            arg_list = shlex.split(arg_str)
            for cli in candidates:
                cmd = cli + arg_list
                if effective_launcher:
                    cmd = shlex.split(effective_launcher) + cmd
                proc = subprocess.run(
                    cmd,
                    cwd=tmpdir,
                    capture_output=True,
                    text=True,
                    env=env,
                    check=False,
                )
                if proc.returncode == 0:
                    output = proc.stdout
                    break
                last_err = (
                    "Cmd: "
                    + " ".join(cmd)
                    + "\n"
                    + "Stdout:\n"
                    + proc.stdout
                    + "\n"
                    + "Stderr:\n"
                    + proc.stderr
                )
                # log the last failing command at DEBUG level; it can be
                # unrelated to the solver implementation/version
                logging.debug(last_err)
            if output:
                break

    if not output:
        raise RuntimeError("Benchmark failed for all CLI variants\n" + (last_err or ""))

    total_time = None
    for line in output.splitlines():
        if "total_time [s]" in line:
            m = re.search(r"total_time \[s\]\s*=\s*([0-9.]+)", line)
            if m:
                total_time = float(m.group(1))
                break

    if total_time is None:
        raise RuntimeError("Failed to extract total_time from output.\n" + output)

    return {"total_time": total_time, "time_per_ite": total_time / N if N else 0.0}


def print_separator(char: str = "=", length: int = 80) -> None:
    """Print a simple separator line to stdout."""
    logging.info(char * length)


def main() -> None:
    """Command-line entry point: compare benchmark timings between commits."""
    # configure logging for the CLI; keep message format minimal
    logging.basicConfig(level=logging.INFO, format="%(message)s")

    repo = get_llg3d_repo()
    assert repo.working_tree_dir is not None
    repo_root = Path(repo.working_tree_dir)

    parser = argparse.ArgumentParser(
        description="Compare simple benchmark timings between two commits",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("--reference-commit", default="HEAD~1")
    parser.add_argument(
        "--solver",
        default="numpy",
        choices=["numpy", "mpi", "opencl"],
    )
    parser.add_argument(
        "--debug",
        action="store_true",
        help="Enable debug logging output",
    )
    parser.add_argument(
        "--np", type=int, default=1, help="Number of processes for MPI solver"
    )
    parser.add_argument(
        "--launcher", default="", help="Launcher command for MPI solver"
    )
    parser.add_argument(
        "--Jx", type=int, default=300, help="Number of grid points in x"
    )
    parser.add_argument("--Jy", type=int, default=21, help="Number of grid points in y")
    parser.add_argument("--Jz", type=int, default=21, help="Number of grid points in z")
    parser.add_argument(
        "--n_mean",
        type=int,
        default=0,
        help="Spatial average frequency (number of iterations)",
    )
    parser.add_argument(
        "--precision",
        choices=["single", "double"],
        default="single",
        help="Precision of the simulation",
    )
    parser.add_argument("--N", type=int, default=100, help="Number of iterations")
    parser.add_argument("--skip-reference", action="store_true")

    args = parser.parse_args()

    if args.debug:
        logging.getLogger().setLevel(logging.DEBUG)

    current_ref = get_current_ref(repo)
    current_commit = get_current_commit(repo)
    reference_commit = args.reference_commit

    if current_ref == current_commit:
        logging.error("detached HEAD — checkout a branch first")
        sys.exit(1)

    print_separator()
    logging.info("SIMPLE COMPARISON BENCHMARK")
    print_separator()
    logging.info(f"Current ref  : {current_ref}")
    logging.info(f"Current hash : {current_commit[:8]}")
    logging.info(f"Reference    : {reference_commit}")
    logging.info(f"Solver       : {args.solver}")
    print_separator()

    reference_result = None
    if not args.skip_reference:
        stash = False
        if repo.is_dirty(untracked_files=True):
            try:
                repo.git.stash("save", "-u")
                stash = True
            except Exception:
                stash = False

        try:
            checkout_commit(repo, reference_commit)
            reference_result = run_single_benchmark(
                args.solver,
                args.Jx,
                args.Jy,
                args.Jz,
                args.n_mean,
                args.N,
                args.launcher or None,
                args.np,
                args.precision,
                repo_root=repo_root,
            )
        finally:
            checkout_commit(repo, current_ref)
            if stash:
                try:
                    repo.git.stash("pop")
                except Exception:
                    pass

    current_result = run_single_benchmark(
        args.solver,
        args.Jx,
        args.Jy,
        args.Jz,
        args.n_mean,
        args.N,
        args.launcher or None,
        args.np,
        args.precision,
        repo_root=repo_root,
    )

    if reference_result:
        speedup = reference_result["total_time"] / current_result["total_time"]
        diff_pct = (
            100
            * (current_result["total_time"] - reference_result["total_time"])
            / reference_result["total_time"]
        )
        logging.info(f"Reference : {reference_result['total_time']:.3f} s")
        logging.info(f"Current   : {current_result['total_time']:.3f} s")
        logging.info(f"Speedup   : {speedup:.3f}x")
        logging.info(f"Diff      : {diff_pct:+.1f}%")
    else:
        logging.info(f"Current   : {current_result['total_time']:.3f} s")

    print_separator()
    logging.info("✓ Benchmark completed!")


if __name__ == "__main__":
    main()
