#!/usr/bin/env python3

"""Run a TIG algorithm over a nonce range and save per-nonce results.

Run this script inside the challenge's TIG runtime Docker container after using
``download_algorithm``. It uses only the Python standard library.
"""

import argparse
import json
import os
import platform
import re
import shlex
import shutil
import subprocess
import sys
import tempfile
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone
from pathlib import Path


CHALLENGE_IDS = {
    "satisfiability": "c001",
    "vehicle_routing": "c002",
    "knapsack": "c003",
    "vector_search": "c004",
    "hypergraph": "c005",
    "neuralnet_optimizer": "c006",
    "job_scheduling": "c007",
    "energy_arbitrage": "c008",
}

RUNTIME_ERROR_TYPES = {
    82: "cuda_out_of_memory",
    83: "host_out_of_memory",
    84: "runtime_error",
    85: "no_solution",
    86: "invalid_solution",
    87: "out_of_fuel",
}

QUALITY_SCALE = 1_000_000
PRINT_LOCK = threading.Lock()


def parse_args():
    tig_runtime_path = shutil.which("tig-runtime")
    tig_verifier_path = shutil.which("tig-verifier")

    parser = argparse.ArgumentParser(
        description="Run a TIG algorithm and save per-nonce results as JSON.",
        epilog=(
            "Timing is most meaningful with the default --workers 1. "
            "Higher worker counts introduce resource contention."
        ),
    )
    parser.add_argument("algorithm", help="Algorithm name")
    parser.add_argument(
        "track_id",
        help="Track ID as comma-separated key-value pairs (for example, a=1,b=2)",
    )
    parser.add_argument(
        "hyperparameters",
        help="JSON hyperparameters, or null to use the algorithm defaults",
    )
    parser.add_argument(
        "--tig-runtime-path",
        default=tig_runtime_path,
        help=f"Path to tig-runtime (default: {tig_runtime_path})",
    )
    parser.add_argument(
        "--tig-verifier-path",
        default=tig_verifier_path,
        help=f"Path to tig-verifier (default: {tig_verifier_path})",
    )
    parser.add_argument(
        "--lib-dir",
        default="./tig-algorithms/lib",
        help="Algorithm library directory (default: ./tig-algorithms/lib)",
    )
    parser.add_argument(
        "--seed",
        default="rand_hash",
        help="Seed string used for instance generation (default: rand_hash)",
    )
    parser.add_argument("--start", type=int, default=0, help="Starting nonce (default: 0)")
    parser.add_argument(
        "--nonces",
        type=int,
        default=100,
        help="Number of nonces to process (default: 100)",
    )
    parser.add_argument(
        "--fuel",
        type=int,
        default=int(100e9),
        help="Maximum fuel per nonce (default: 100 billion)",
    )
    parser.add_argument(
        "--workers",
        type=int,
        default=1,
        help="Worker threads (default: 1; higher values make timing less accurate)",
    )
    parser.add_argument(
        "--output",
        default="test_algorithm_results.json",
        help="JSON report path (default: test_algorithm_results.json)",
    )
    parser.add_argument("--verbose", action="store_true", help="Print commands and subprocess output")
    return parser, parser.parse_args()


def detect_cpu_arch():
    machine = platform.machine().lower()
    if machine in ("x86_64", "amd64"):
        return "amd64"
    if machine in ("arm64", "aarch64"):
        return "arm64"
    raise ValueError(f"Unsupported CPU architecture: {machine}")


def detect_visible_gpus():
    if shutil.which("nvidia-smi") is None:
        return []

    configured_devices = os.environ.get("CUDA_VISIBLE_DEVICES")
    if configured_devices is not None:
        try:
            return [int(device) for device in configured_devices.split(",") if device]
        except ValueError as error:
            raise ValueError("CUDA_VISIBLE_DEVICES must be a comma-separated list of GPU numbers") from error

    output = subprocess.check_output(["nvidia-smi", "-L"], text=True)
    return [
        int(match.group(1))
        for line in output.splitlines()
        if (match := re.match(r"^GPU (\d+):", line)) is not None
    ]


def configure_visible_cpus():
    configured_cores = os.environ.get("CPU_VISIBLE_CORES")
    if configured_cores is None:
        return
    cores = [int(core) for core in configured_cores.split(",")]
    os.sched_setaffinity(0, cores)


def elapsed_ms(start_ns):
    return (time.perf_counter_ns() - start_ns) / 1_000_000


def run_command(command):
    start_ns = time.perf_counter_ns()
    completed = subprocess.run(command, capture_output=True, text=True, check=False)
    return completed, elapsed_ms(start_ns)


def error_record(stage, error_type, exit_code=None, message=None):
    record = {"stage": stage, "type": error_type}
    if exit_code is not None:
        record["exit_code"] = exit_code
    if message:
        record["message"] = message.strip()
    return record


def runtime_error_type(exit_code):
    if exit_code < 0:
        return "terminated_by_signal"
    return RUNTIME_ERROR_TYPES.get(exit_code, "unknown_runtime_error")


def verifier_error_type(stderr):
    if "verification error:" in stderr.lower():
        return "invalid_solution"
    return "verifier_error"


def parse_quality(stdout):
    matches = re.findall(r"^quality:\s*(-?\d+)\s*$", stdout, flags=re.MULTILINE)
    return int(matches[-1]) if matches else None


def empty_result(nonce):
    return {
        "nonce": nonce,
        "status": "error",
        "verified": False,
        "quality": None,
        "fuel_consumed": None,
        "runtime_ms": None,
        "verifier_ms": None,
        "total_ms": None,
        "runtime_exit_code": None,
        "verifier_exit_code": None,
        "errors": [],
    }


def print_verbose_result(nonce, runtime_command, runtime, verifier_command, verifier, result):
    lines = [
        f"[nonce {nonce}]",
        f"  runtime command: {shlex.join(runtime_command)}",
        f"  runtime elapsed: {result['runtime_ms']:.3f} ms",
        f"  runtime exit code: {result['runtime_exit_code']}",
    ]
    if runtime.stdout.strip():
        lines.append(f"  runtime stdout:\n{runtime.stdout.rstrip()}")
    if runtime.stderr.strip():
        lines.append(f"  runtime stderr:\n{runtime.stderr.rstrip()}")

    if verifier is not None:
        lines.extend(
            [
                f"  verifier command: {shlex.join(verifier_command)}",
                f"  verifier elapsed: {result['verifier_ms']:.3f} ms",
                f"  verifier exit code: {result['verifier_exit_code']}",
            ]
        )
        if verifier.stdout.strip():
            lines.append(f"  verifier stdout:\n{verifier.stdout.rstrip()}")
        if verifier.stderr.strip():
            lines.append(f"  verifier stderr:\n{verifier.stderr.rstrip()}")

    with PRINT_LOCK:
        print("\n".join(lines))


def run_nonce(nonce, args, settings_json, so_path, ptx_path, visible_gpus):
    result = empty_result(nonce)
    total_start_ns = time.perf_counter_ns()

    try:
        with tempfile.TemporaryDirectory() as temp_dir:
            output_file = Path(temp_dir) / f"{nonce}.json"
            runtime_command = [
                args.tig_runtime_path,
                settings_json,
                args.seed,
                str(nonce),
                str(so_path),
                "--fuel",
                str(args.fuel),
                "--output",
                temp_dir,
            ]

            gpu_device = None
            if ptx_path is not None:
                gpu_device = visible_gpus[nonce % len(visible_gpus)]
                runtime_command.extend(["--ptx", str(ptx_path), "--gpu", str(gpu_device)])
            if args.hyperparameters is not None:
                runtime_command.extend(["--hyperparameters", args.hyperparameters])

            runtime, runtime_duration = run_command(runtime_command)
            result["runtime_ms"] = runtime_duration
            result["runtime_exit_code"] = runtime.returncode

            if runtime.returncode != 0:
                result["errors"].append(
                    error_record(
                        "runtime",
                        runtime_error_type(runtime.returncode),
                        runtime.returncode,
                        runtime.stderr,
                    )
                )

            if output_file.exists():
                try:
                    with output_file.open(encoding="utf-8") as file:
                        output_data = json.load(file)
                    result["fuel_consumed"] = output_data.get("fuel_consumed")
                except (OSError, json.JSONDecodeError) as error:
                    result["errors"].append(
                        error_record("output", "invalid_output_json", message=str(error))
                    )
            else:
                result["errors"].append(error_record("output", "missing_output"))

            verifier = None
            verifier_command = None
            if output_file.exists():
                verifier_command = [
                    args.tig_verifier_path,
                    settings_json,
                    args.seed,
                    str(nonce),
                    str(output_file),
                ]
                if ptx_path is not None:
                    verifier_command.extend(["--ptx", str(ptx_path), "--gpu", str(gpu_device)])
                if args.verbose:
                    verifier_command.append("--verbose")

                verifier, verifier_duration = run_command(verifier_command)
                result["verifier_ms"] = verifier_duration
                result["verifier_exit_code"] = verifier.returncode

                quality = parse_quality(verifier.stdout)
                if verifier.returncode == 0 and quality is not None:
                    result["verified"] = True
                    result["quality"] = quality
                elif verifier.returncode != 0:
                    result["errors"].append(
                        error_record(
                            "verifier",
                            verifier_error_type(verifier.stderr),
                            verifier.returncode,
                            verifier.stderr,
                        )
                    )
                else:
                    result["errors"].append(
                        error_record(
                            "verifier",
                            "missing_quality",
                            verifier.returncode,
                            "Verifier succeeded but did not print a quality value",
                        )
                    )

            result["total_ms"] = elapsed_ms(total_start_ns)
            if not result["errors"]:
                result["status"] = "success"
            elif result["verified"]:
                result["status"] = "verified_with_errors"

            if args.verbose:
                print_verbose_result(
                    nonce,
                    runtime_command,
                    runtime,
                    verifier_command,
                    verifier,
                    result,
                )
    except Exception as error:  # Keep the remaining nonce set running and report harness failures.
        result["total_ms"] = elapsed_ms(total_start_ns)
        result["errors"].append(error_record("harness", "harness_error", message=str(error)))

    return result


def validate_args(parser, args, challenge, cpu_arch, visible_gpus):
    if challenge not in CHALLENGE_IDS:
        parser.error(f"Unsupported CHALLENGE value: {challenge}")
    if args.tig_runtime_path is None:
        parser.error("tig-runtime was not found on PATH; run this script inside a TIG runtime image")
    if args.tig_verifier_path is None:
        parser.error("tig-verifier was not found on PATH; run this script inside a TIG runtime image")
    if args.start < 0:
        parser.error("--start must be zero or greater")
    if args.nonces < 1:
        parser.error("--nonces must be at least 1")
    if args.fuel < 1:
        parser.error("--fuel must be at least 1")
    if args.workers < 1:
        parser.error("--workers must be at least 1")

    so_path = Path(args.lib_dir) / challenge / cpu_arch / f"{args.algorithm}.so"
    ptx_candidate = Path(args.lib_dir) / challenge / "ptx" / f"{args.algorithm}.ptx"
    ptx_path = ptx_candidate if ptx_candidate.exists() else None

    if not so_path.exists():
        parser.error(
            f"Algorithm library not found at {so_path}. Run download_algorithm {args.algorithm} first."
        )
    if ptx_path is not None and not visible_gpus:
        parser.error(f"PTX found at {ptx_path}, but no NVIDIA GPU was detected")

    return so_path, ptx_path


def main():
    parser, args = parse_args()
    challenge = os.environ.get("CHALLENGE")
    if challenge is None:
        parser.error("CHALLENGE is not set; run this script inside a TIG runtime image")

    try:
        cpu_arch = detect_cpu_arch()
        visible_gpus = detect_visible_gpus()
        configure_visible_cpus()
    except (OSError, ValueError) as error:
        parser.error(str(error))

    if args.hyperparameters.lower() == "null":
        args.hyperparameters = None
    else:
        try:
            json.loads(args.hyperparameters)
        except json.JSONDecodeError as error:
            parser.error(f"Failed to parse hyperparameters as JSON: {error}")

    so_path, ptx_path = validate_args(parser, args, challenge, cpu_arch, visible_gpus)

    if args.workers > 1:
        print(
            "Warning: --workers is greater than 1. CPU/GPU contention and scheduling "
            "will make per-nonce timing less accurate."
        )

    settings = {
        "algorithm_id": "",
        "challenge_id": CHALLENGE_IDS[challenge],
        "track_id": args.track_id,
        "block_id": "",
        "player_id": "",
    }
    settings_json = json.dumps(settings, separators=(",", ":"))
    nonces = range(args.start, args.start + args.nonces)
    results = []

    print(f"Processing {args.nonces} nonces with {args.workers} worker(s)...")
    with ThreadPoolExecutor(max_workers=args.workers) as executor:
        futures = {
            executor.submit(
                run_nonce,
                nonce,
                args,
                settings_json,
                so_path,
                ptx_path,
                visible_gpus,
            ): nonce
            for nonce in nonces
        }
        for completed, future in enumerate(as_completed(futures), start=1):
            result = future.result()
            results.append(result)
            if not args.verbose:
                print(
                    f"\r#completed: {completed}/{args.nonces}",
                    end="",
                    flush=True,
                )

    if not args.verbose:
        print()

    results.sort(key=lambda result: result["nonce"])
    report = {
        "schema_version": 1,
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "configuration": {
            "challenge": challenge,
            "algorithm": args.algorithm,
            "track_id": args.track_id,
            "hyperparameters": args.hyperparameters,
            "seed": args.seed,
            "start_nonce": args.start,
            "num_nonces": args.nonces,
            "fuel_limit": args.fuel,
            "workers": args.workers,
            "timing_warning": (
                "Parallel workers introduce resource contention; use workers=1 for meaningful timing."
                if args.workers > 1
                else None
            ),
            "quality_scale": QUALITY_SCALE,
            "timing_definitions": {
                "runtime_ms": (
                    "Complete tig-runtime process: startup, instance generation, "
                    "algorithm execution, and output writing"
                ),
                "verifier_ms": (
                    "Complete tig-verifier process: instance regeneration and "
                    "solution evaluation"
                ),
                "total_ms": "Elapsed time spanning both processes and harness overhead",
            },
        },
        "results": results,
    }

    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with output_path.open("w", encoding="utf-8") as file:
        json.dump(report, file, indent=2)
        file.write("\n")

    print(f"Wrote per-nonce results to {output_path}")


if __name__ == "__main__":
    main()