import argparse
import csv
import json
import sys


def parse_args():
    parser = argparse.ArgumentParser(
        description="Compare benchmark results between main and PR branches"
    )
    parser.add_argument("--main", required=True, help="Path to main branch CSV results")
    parser.add_argument("--pr", required=True, help="Path to PR branch CSV results")
    parser.add_argument(
        "--output", required=True, help="Path to write comparison JSON output"
    )
    parser.add_argument("--dataset", required=True, help="Dataset name for display")
    parser.add_argument(
        "--qps-threshold",
        type=float,
        default=0.025,
        help="QPS regression threshold (default: 0.025)",
    )
    parser.add_argument(
        "--recall-threshold",
        type=float,
        default=0.03,
        help="Recall regression threshold (default: 0.03)",
    )
    parser.add_argument(
        "--build-time-threshold",
        type=float,
        default=0.50,
        help="Build time regression threshold (default: 0.50)",
    )
    return parser.parse_args()


def parse_csv_results(csv_file):
    runs = []
    build_time = None
    index_size = None

    with open(csv_file) as f:
        reader = csv.DictReader(f)
        rows = list(reader)
        if not rows:
            return None

        for row in rows:
            runs.append(
                {
                    "recall": float(row["recall"]),
                    "qps": float(row["qps"]),
                    "n_probes": int(row["n_probes"]),
                }
            )
            if build_time is None:
                build_time = float(row["build_time_s"])
                index_size = float(row["index_size_kb"])

        runs.sort(key=lambda x: x["recall"])

    return {"runs": runs, "build_time_s": build_time, "index_size_kb": index_size}


def pct_change(new, old):
    return ((new - old) / old) * 100


def main():
    args = parse_args()

    try:
        main_data = parse_csv_results(args.main)
        pr_data = parse_csv_results(args.pr)
    except Exception as e:
        print(f"ERROR: Could not parse CSV files: {e}")
        sys.exit(1)

    if not main_data or not pr_data:
        print("ERROR: Could not extract metrics from CSV files")
        sys.exit(1)

    if len(main_data["runs"]) != len(pr_data["runs"]):
        print(
            f"ERROR: Mismatch in number of runs: {len(main_data['runs'])} (main) vs {len(pr_data['runs'])} (PR)"
        )
        sys.exit(1)

    if len(main_data["runs"]) == 0:
        print("ERROR: No runs found in results")
        sys.exit(1)

    main_runs = main_data["runs"]
    pr_runs = pr_data["runs"]
    main_build_time = main_data["build_time_s"]
    main_index_size = main_data["index_size_kb"]
    pr_build_time = pr_data["build_time_s"]
    pr_index_size = pr_data["index_size_kb"]

    qps_thresh = args.qps_threshold * 100
    recall_thresh = args.recall_threshold * 100
    build_time_thresh = args.build_time_threshold * 100

    changes = {}
    regressions = []

    for i, (main_run, pr_run) in enumerate(zip(main_runs, pr_runs)):
        recall_level = f"{main_run['recall'] * 100:.0f}%"
        qps_change = pct_change(pr_run["qps"], main_run["qps"])
        recall_change = pct_change(pr_run["recall"], main_run["recall"])
        changes[f"qps_{i}"] = {"recall": recall_level, "change": qps_change}
        changes[f"recall_{i}"] = {"recall": recall_level, "change": recall_change}

        if qps_change < -qps_thresh:
            regressions.append(
                f"QPS @ {recall_level} recall regressed by {abs(qps_change):.1f}%"
            )
        if recall_change < -recall_thresh:
            regressions.append(
                f"Recall @ {recall_level} regressed by {abs(recall_change):.1f}%"
            )

    build_time_change = pct_change(pr_build_time, main_build_time)
    index_size_change = pct_change(pr_index_size, main_index_size)
    changes["build_time"] = build_time_change
    changes["index_size"] = index_size_change

    if build_time_change > build_time_thresh:
        regressions.append(f"Build time regressed by {build_time_change:.1f}%")
    if index_size_change > qps_thresh:
        regressions.append(f"Index size regressed by {index_size_change:.1f}%")

    # Print results
    print("\n" + "=" * 70)
    print(f"Dataset: {args.dataset}")
    print("=" * 70)

    print("\nResults for Main Branch:")
    print("┌──────────────────────────────────────────────────┐")
    for run in main_runs:
        print(
            f"│ QPS @ {run['recall'] * 100:>5.1f}% recall:   {run['qps']:>8.2f} queries/second │"
        )
    print(f"│ Build time:             {main_build_time:>8.2f} seconds       │")
    print(f"│ Index size:             {main_index_size:>8.2f} KB            │")
    print("└──────────────────────────────────────────────────┘")

    print("\nResults for Feature Branch:")
    print(
        "┌──────────────────────────────────────────────────────────────────────────────┐"
    )
    for i, run in enumerate(pr_runs):
        qps_change = changes[f"qps_{i}"]["change"]
        recall_change = changes[f"recall_{i}"]["change"]
        qps_warn = "WARNING " if qps_change < -qps_thresh else "        "
        recall_warn = "WARNING " if recall_change < -recall_thresh else "        "
        print(
            f"│ QPS @ {run['recall'] * 100:>5.1f}% recall:   {run['qps']:>8.2f}  {qps_warn}{qps_change:>+8.1f}% {recall_warn}recall: {recall_change:>+8.1f}% │"
        )

    warn = "WARNING " if build_time_change > build_time_thresh else "        "
    print(
        f"│ Build time:             {pr_build_time:>8.2f}s {warn}{build_time_change:>+8.1f}%                        │"
    )

    warn = "WARNING " if index_size_change > qps_thresh else "        "
    print(
        f"│ Index size:             {pr_index_size:>8.2f} KB {warn}{index_size_change:>+8.1f}%                        │"
    )
    print(
        "└──────────────────────────────────────────────────────────────────────────────┘"
    )

    with open(args.output, "w") as f:
        json.dump({"changes": changes, "regressions": regressions}, f, indent=2)

    if regressions:
        print("\nFAILED:")
        for reg in regressions:
            print(f"   {reg}")
        sys.exit(1)
    else:
        print(f"\nPASSED - No regressions detected for {args.dataset}")
        sys.exit(0)


if __name__ == "__main__":
    main()
