#!/usr/bin/env python3

import argparse
import git
import json
import pathlib
import sys
import time
from typing import List

VERBOSE: bool = False

TIME_MULTI: dict = {"ns": 1, "us": 1000, "ms": 1000 * 1000, "s": 1000 * 1000 * 1000}


def log(msg: str) -> None:
    if VERBOSE:
        print(msg)


def warning(msg: str) -> None:
    print("\033[1;93m" + msg + "\033[0m")


def error(msg: str) -> None:
    print("\033[1;31m" + msg + "\033[0m")


def sort_commits_by_tree_order(repo: git.Repo, commit_hashes: List[str]) -> List[str]:
    target_hashes = set(commit_hashes)
    ordered_commits = list(
        repo.iter_commits(rev="HEAD", topo_order=True, reverse=False)
    )

    # Build position mapping for target commits
    commit_positions = {}
    for idx, commit in enumerate(ordered_commits):
        commit_hash = commit.hexsha
        if commit_hash in target_hashes:
            commit_positions[commit_hash] = idx

    # Sort input hashes by their topological position
    # Handle missing commits by putting them at the end
    def sort_key(commit_hash: str) -> int:
        return commit_positions.get(commit_hash, len(ordered_commits))

    return sorted(commit_hashes, key=sort_key)


def main() -> None:
    global VERBOSE

    benchmarks = {"commits": [], "results": {}}
    commits = []
    benchNames = []
    parser = argparse.ArgumentParser(
        prog="benchreport",
        description="Generate an HTML report for bpfilter benchmarks",
    )
    parser.add_argument(
        "-s", "--sources", default=".", help="Sources directory containing .git"
    )
    parser.add_argument(
        "-r", "--results", help="Directory containing the benchmark results"
    )
    parser.add_argument("-t", "--template", help="HTML report template file")
    parser.add_argument("-o", "--output", help="Output HTML file")
    parser.add_argument(
        "-v", "--verbose", action="store_true", help="Produce a verbose output"
    )
    args = parser.parse_args()

    VERBOSE = args.verbose

    files = list(pathlib.Path(args.results).glob("*.json"))
    if not files:
        warning(f"No benchmark results found in '{args.results}', ignoring")
        sys.exit(0)

    for file in files:
        with open(file, "r", encoding="utf-8") as f:
            log(f"Reading benchmark results from {file}")
            d = json.load(f)

            gitrev = d["context"]["gitrev"]
            commits.append(gitrev)

            for bench in d["benchmarks"]:
                benchNames.append(bench["name"])
                if bench["name"] not in benchmarks["results"]:
                    benchmarks["results"][bench["name"]] = {}

                benchmarks["results"][bench["name"]][gitrev] = {
                    "iters": bench["iterations"],
                    "time": bench["real_time"] * TIME_MULTI[bench["time_unit"]],
                }

                nInsn = bench.get("nInsn", None)
                if nInsn:
                    benchmarks["results"][bench["name"]][gitrev]["nInsn"] = nInsn

    repo = git.Repo.init(args.sources)
    commits = sort_commits_by_tree_order(repo, commits)

    for commit in commits:
        try:
            message = repo.commit(commit).message
            date = repo.commit(commit).committed_date
        except:
            message = "<Not committed yet>"
            date = int(time.time())

        benchmarks["commits"].append(
            {
                "sha": commit,
                "date": date,
                "message": message,
            }
        )

    benchmarks["benchNames"] = list(dict.fromkeys(benchNames))

    with open(args.output, "w", encoding="utf-8") as f:
        with open(args.template, "r", encoding="utf-8") as template_file:
            template = template_file.read()
        f.write(template.replace("{{ DATA }}", json.dumps(benchmarks)))
    log(f"Benchmark report generated at {args.output}")


if __name__ == "__main__":
    main()
