"""CLI for the Piggie benchmark suite."""

from __future__ import annotations

import argparse
import subprocess
import sys
from pathlib import Path
from typing import Any

from benchmarks.config import (
    ALGO_FAMILIES,
    ALL_DBS,
    ALL_WORKLOADS,
    DEFAULT_SCALE,
    SCALES,
)

BENCHMARKS_DIR = Path(__file__).parent


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="benchmarks",
        description="Piggie benchmark suite — graph database comparison",
    )
    sub = parser.add_subparsers(dest="command", help="Command to run")

    # -- run (default) ---------------------------------------------------
    run_p = sub.add_parser("run", help="Run benchmarks (default)")
    run_p.add_argument(
        "--db",
        nargs="+",
        choices=ALL_DBS,
        default=list(ALL_DBS),
        help="Databases to benchmark",
    )
    run_p.add_argument(
        "--workload",
        nargs="+",
        default=list(ALL_WORKLOADS),
        help="Workload IDs to run",
    )
    run_p.add_argument(
        "--scale",
        choices=list(SCALES.keys()),
        default=DEFAULT_SCALE,
        help="Scale preset (default: 10k)",
    )
    run_p.add_argument(
        "--results-dir",
        default=str(BENCHMARKS_DIR / "results"),
        help="Directory for raw JSON results",
    )

    # -- setup -----------------------------------------------------------
    sub.add_parser("setup", help="Start Docker containers for Neo4j and NebulaGraph")

    # -- teardown --------------------------------------------------------
    sub.add_parser("teardown", help="Stop and remove Docker containers")

    # -- report ----------------------------------------------------------
    report_p = sub.add_parser("report", help="Generate markdown report from results")
    report_p.add_argument(
        "results_dir",
        nargs="?",
        default=str(BENCHMARKS_DIR / "results"),
        help="Path to results directory",
    )
    report_p.add_argument(
        "--output-dir",
        default=str(BENCHMARKS_DIR / "reports"),
        help="Directory for generated reports",
    )

    # -- algo ------------------------------------------------------------
    algo_p = sub.add_parser("algo", help="Run algorithm performance benchmarks")
    algo_p.add_argument(
        "--scale",
        choices=list(SCALES.keys()),
        default=DEFAULT_SCALE,
        help="Scale preset (default: 10k)",
    )
    algo_p.add_argument(
        "--db",
        choices=["age", "neo4j"],
        default="age",
        help="Database to benchmark (default: age)",
    )
    algo_p.add_argument(
        "--family",
        nargs="+",
        choices=list(ALGO_FAMILIES),
        default=None,
        help="Algorithm families to benchmark",
    )
    algo_p.add_argument(
        "--algo",
        nargs="+",
        default=None,
        help="Specific algorithm IDs to benchmark",
    )
    algo_p.add_argument(
        "--round",
        type=int,
        default=None,
        dest="bench_round",
        help="Optimization round number for tracking progress",
    )
    algo_p.add_argument(
        "--results-dir",
        default=str(BENCHMARKS_DIR / "results"),
        help="Directory for raw JSON results",
    )

    # -- algo-report -----------------------------------------------------
    algo_report_p = sub.add_parser(
        "algo-report", help="Generate algorithm benchmark report"
    )
    algo_report_p.add_argument(
        "results_dir",
        nargs="?",
        default=str(BENCHMARKS_DIR / "results"),
        help="Path to results directory",
    )
    algo_report_p.add_argument(
        "--output-dir",
        default=str(BENCHMARKS_DIR / "reports"),
        help="Directory for generated reports",
    )

    # -- scorecard -------------------------------------------------------
    scorecard_p = sub.add_parser(
        "scorecard", help="Generate competitor scorecard (Piggie vs Neo4j GDS)"
    )
    scorecard_p.add_argument(
        "results_dir",
        nargs="?",
        default=str(BENCHMARKS_DIR / "results"),
        help="Path to results directory",
    )
    scorecard_p.add_argument(
        "--output-dir",
        default=str(BENCHMARKS_DIR / "reports"),
        help="Directory for generated reports",
    )

    return parser


def cmd_setup() -> None:
    """Start Docker containers."""
    compose_file = BENCHMARKS_DIR / "docker-compose.yml"
    if not compose_file.exists():
        print("ERROR: docker-compose.yml not found", file=sys.stderr)
        sys.exit(1)
    print("Starting benchmark containers...")
    subprocess.run(
        ["docker", "compose", "-f", str(compose_file), "up", "-d"],
        check=True,
    )
    print("Containers started. Waiting for services to be ready...")
    _wait_for_neo4j()
    _wait_for_nebula()
    print("All services ready.")


def cmd_teardown() -> None:
    """Stop Docker containers."""
    compose_file = BENCHMARKS_DIR / "docker-compose.yml"
    if not compose_file.exists():
        print("ERROR: docker-compose.yml not found", file=sys.stderr)
        sys.exit(1)
    print("Stopping benchmark containers...")
    subprocess.run(
        ["docker", "compose", "-f", str(compose_file), "down", "-v"],
        check=True,
    )
    print("Containers stopped.")


def cmd_report(results_dir: str, output_dir: str) -> None:
    """Generate markdown report."""
    from benchmarks.report import write_report

    try:
        path = write_report(results_dir, output_dir)
        print(f"Report written to {path}")
    except FileNotFoundError as e:
        print(f"ERROR: {e}", file=sys.stderr)
        sys.exit(1)


def cmd_algo(
    scale: str,
    db: str,
    families: list[str] | None,
    algo_ids: list[str] | None,
    bench_round: int | None,
    results_dir: str,
) -> None:
    """Run algorithm benchmarks."""
    from benchmarks.algo_bench import run_algo_benchmark

    try:
        run_algo_benchmark(
            scale,
            db=db,
            families=families,
            algo_ids=algo_ids,
            results_dir=results_dir,
            bench_round=bench_round,
        )
    except Exception as exc:
        print(f"FATAL: {type(exc).__name__}: {exc}", file=sys.stderr)
        sys.exit(1)


def cmd_algo_report(results_dir: str, output_dir: str) -> None:
    """Generate algorithm benchmark report."""
    from benchmarks.algo_report import write_algo_report

    try:
        path = write_algo_report(results_dir, output_dir)
        print(f"Algorithm report written to {path}")
    except FileNotFoundError as e:
        print(f"ERROR: {e}", file=sys.stderr)
        sys.exit(1)


def cmd_scorecard(results_dir: str, output_dir: str) -> None:
    """Generate competitor scorecard."""
    from benchmarks.algo_report import write_scorecard

    try:
        path = write_scorecard(results_dir, output_dir)
        print(f"Scorecard written to {path}")
    except FileNotFoundError as e:
        print(f"ERROR: {e}", file=sys.stderr)
        sys.exit(1)


def cmd_run(
    dbs: list[str],
    workloads: list[str],
    scale_name: str,
    results_dir: str,
) -> None:
    """Run benchmarks."""
    from benchmarks.datagen.synthetic import generate, generate_query_vector
    from benchmarks.harness import BenchmarkRun, collect_system_info, execute_benchmark

    scale = SCALES[scale_name]
    print(f"=== Piggie Benchmark Suite ===")
    print(f"Scale: {scale.name} ({scale.num_nodes:,} nodes, {scale.num_edges:,} edges)")
    print(f"Databases: {', '.join(dbs)}")
    print(f"Workloads: {len(workloads)}")
    print()

    # Generate data
    print("Generating synthetic data...", flush=True)
    data = generate(scale)
    query_vec = generate_query_vector(scale)
    print(f"  {len(data.nodes):,} nodes, {len(data.edges):,} edges, "
          f"{data.vectors.shape if data.vectors is not None else 'no'} vectors")
    print()

    # Build adapter instances
    adapters = _build_adapters(dbs)

    for adapter in adapters:
        print(f"--- {adapter.name} ---")
        try:
            adapter.connect()
            adapter.setup_schema(scale.vec_dims)

            # Load baseline data
            print(f"  Loading data...", flush=True)
            adapter.load_nodes(data.nodes)
            adapter.load_edges(data.edges)

            # Load vectors if adapter supports it
            if hasattr(adapter, "_load_vectors") and data.vectors is not None:
                print(f"  Loading vectors...", flush=True)
                adapter._load_vectors(data.nodes, data.vectors)

            # Build workload functions
            wl_fns = _build_workload_fns(adapter, data, query_vec, workloads, scale)

            # Execute
            run = execute_benchmark(adapter, wl_fns, scale)
            path = run.save(results_dir)
            print(f"  Results saved to {path}")

        except Exception as exc:
            print(f"  FATAL: {type(exc).__name__}: {exc}")
        finally:
            try:
                adapter.teardown()
            except Exception:
                pass
            try:
                adapter.close()
            except Exception:
                pass
        print()

    print("Done. Run `python -m benchmarks report` to generate the markdown report.")


def _build_adapters(dbs: list[str]) -> list[Any]:
    """Instantiate adapter objects for requested databases."""
    adapters = []
    for db in dbs:
        if db == "age":
            from benchmarks.adapters.age_adapter import AgeAdapter
            adapters.append(AgeAdapter())
        elif db == "neo4j":
            from benchmarks.adapters.neo4j_adapter import Neo4jAdapter
            adapters.append(Neo4jAdapter())
        elif db == "kuzu":
            from benchmarks.adapters.kuzu_adapter import KuzuAdapter
            adapters.append(KuzuAdapter())
        elif db == "nebula":
            from benchmarks.adapters.nebula_adapter import NebulaAdapter
            adapters.append(NebulaAdapter())
    return adapters


def _build_workload_fns(
    adapter: Any,
    data: Any,
    query_vec: Any,
    workload_ids: list[str],
    scale: Any,
) -> dict[str, Any]:
    """Map workload IDs to callables for this adapter."""
    from benchmarks.datagen.synthetic import (
        generate_workload_edges,
        generate_workload_nodes,
    )

    fns: dict[str, Any] = {}
    write_nodes = generate_workload_nodes(100, start_id=scale.num_nodes + 1000)
    write_edges = generate_workload_edges(100, max_node_id=scale.num_nodes)

    # Pick Person nodes for traversal workloads (Person has the most edge types)
    person_ids = data.nodes_by_category.get("Person", [0, 100])
    start_node = person_ids[0]
    end_node = person_ids[min(100, len(person_ids) - 1)]

    mapping: dict[str, Any] = {
        "W01_node_creation": lambda: adapter.w01_node_creation(write_nodes),
        "W02_edge_creation": lambda: adapter.w02_edge_creation(write_edges),
        "W03_point_lookup": lambda: adapter.w03_point_lookup(start_node),
        "W04_pattern_match": lambda: adapter.w04_pattern_match(2, start_node),
        "W05_aggregation": lambda: adapter.w05_aggregation(),
        "W06_vle": lambda: adapter.w06_vle(1, 4, start_node),
        "W07_shortest_path": lambda: adapter.w07_shortest_path(start_node, end_node),
        "W08_pagerank": lambda: adapter.w08_pagerank(20),
        "W09_vector_search": lambda: adapter.w09_vector_search(query_vec, 10),
        "W10_hybrid_search": lambda: adapter.w10_hybrid_search(query_vec, 10, "Person"),
        "W11_bulk_load": lambda: adapter.w11_bulk_load(
            generate_workload_nodes(1000, start_id=scale.num_nodes + 100_000),
            generate_workload_edges(5000, max_node_id=scale.num_nodes),
        ),
        "W12_concurrent": lambda: adapter.w12_concurrent(
            ["MATCH (n {bench_id: 1}) RETURN n"] * 50,
            4,
        ),
    }

    for wid in workload_ids:
        fns[wid] = mapping.get(wid)

    return fns


def _wait_for_neo4j(timeout: int = 60) -> None:
    """Wait for Neo4j to accept connections."""
    import time

    from benchmarks.config import NEO4J_URI

    print("  Waiting for Neo4j...", end=" ", flush=True)
    start = time.time()
    while time.time() - start < timeout:
        try:
            import neo4j as neo4j_lib
            driver = neo4j_lib.GraphDatabase.driver(NEO4J_URI, auth=("neo4j", "benchpass123"))
            driver.verify_connectivity()
            driver.close()
            print("ready.")
            return
        except Exception:
            time.sleep(2)
    print("TIMEOUT (proceeding anyway)")


def _wait_for_nebula(timeout: int = 60) -> None:
    """Wait for NebulaGraph to accept connections."""
    import time

    from benchmarks.config import NEBULA_HOST, NEBULA_PORT

    print("  Waiting for NebulaGraph...", end=" ", flush=True)
    start = time.time()
    while time.time() - start < timeout:
        try:
            import socket
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            s.settimeout(2)
            s.connect((NEBULA_HOST, NEBULA_PORT))
            s.close()
            print("ready.")
            return
        except Exception:
            time.sleep(2)
    print("TIMEOUT (proceeding anyway)")


def main(argv: list[str] | None = None) -> None:
    parser = build_parser()
    args = parser.parse_args(argv)

    if args.command is None or args.command == "run":
        # Default to 'run' when no subcommand given
        if args.command is None:
            # Re-parse with run defaults
            args = build_parser().parse_args(["run"] + (argv or []))
        cmd_run(args.db, args.workload, args.scale, args.results_dir)
    elif args.command == "setup":
        cmd_setup()
    elif args.command == "teardown":
        cmd_teardown()
    elif args.command == "report":
        cmd_report(args.results_dir, args.output_dir)
    elif args.command == "algo":
        cmd_algo(args.scale, args.db, args.family, args.algo, args.bench_round, args.results_dir)
    elif args.command == "algo-report":
        cmd_algo_report(args.results_dir, args.output_dir)
    elif args.command == "scorecard":
        cmd_scorecard(args.results_dir, args.output_dir)
