import argparse
import base64
import json
import statistics
import subprocess
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Any, Dict, List

import requests


DEFAULT_OLD_ENDPOINT = "http://127.0.0.1:5002/v1/api/ai/suya_ocr/"
DEFAULT_NEW_ENDPOINT = "http://127.0.0.1:5002/v1/api/ai/suya_ocr_vllm/"


def _encode_image(path: Path) -> str:
    return base64.b64encode(path.read_bytes()).decode("utf-8")


def _preflight() -> Dict[str, Any]:
    checks: Dict[str, Any] = {}
    commands = {
        "docker_runtimes": ["docker", "info", "--format", "{{json .Runtimes}}"],
        "nvidia_smi": ["nvidia-smi", "-L"],
        "nvidia_container_runtime": ["which", "nvidia-container-runtime"],
        "nvidia_ctk": ["which", "nvidia-ctk"],
        "docker_gpus": [
            "docker",
            "run",
            "--rm",
            "--gpus",
            "all",
            "--entrypoint",
            "nvidia-smi",
            "vllm/vllm-openai:v0.20.1",
            "-L",
        ],
    }
    for name, cmd in commands.items():
        try:
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=15)
            checks[name] = {
                "ok": result.returncode == 0,
                "stdout": result.stdout.strip(),
                "stderr": result.stderr.strip(),
            }
        except Exception as exc:
            checks[name] = {"ok": False, "stdout": "", "stderr": str(exc)}
    runtimes = checks.get("docker_runtimes", {}).get("stdout", "")
    checks["docker_has_nvidia_runtime"] = '"nvidia"' in runtimes
    if not checks["docker_has_nvidia_runtime"] and checks.get("docker_gpus", {}).get("ok"):
        checks["recommended_fix"] = (
            "Docker GPU passthrough works with --gpus. Use the local "
            "scripts/docker_nvidia_runtime_compat.sh wrapper to strip the legacy "
            "--runtime nvidia flag from Surya's vLLM spawn command."
        )
    else:
        checks["recommended_fix"] = (
            "Run: sudo nvidia-ctk runtime configure --runtime=docker "
            "--config=/etc/docker/daemon.json && sudo systemctl restart docker"
        )
    return checks


def _percentile(values: List[float], percentile: float) -> float:
    if not values:
        return 0.0
    ordered = sorted(values)
    index = min(len(ordered) - 1, int(round((percentile / 100) * (len(ordered) - 1))))
    return ordered[index]


def _post_once(url: str, payload: Dict[str, Any], timeout: float) -> Dict[str, Any]:
    start = time.perf_counter()
    try:
        response = requests.post(url, json=payload, timeout=timeout)
        elapsed = time.perf_counter() - start
        body = response.json() if response.content else {}
        return {
            "ok": response.status_code == 200 and body.get("code") == 200,
            "status_code": response.status_code,
            "api_code": body.get("code"),
            "elapsed_seconds": elapsed,
            "error": body.get("message") if body.get("code") != 200 else None,
        }
    except Exception as exc:
        return {
            "ok": False,
            "status_code": None,
            "api_code": None,
            "elapsed_seconds": time.perf_counter() - start,
            "error": str(exc),
        }


def run_endpoint(
    name: str,
    url: str,
    payload: Dict[str, Any],
    requests_count: int,
    concurrency: int,
    timeout: float,
) -> Dict[str, Any]:
    start = time.perf_counter()
    samples: List[Dict[str, Any]] = []
    with ThreadPoolExecutor(max_workers=concurrency) as pool:
        futures = [
            pool.submit(_post_once, url, payload, timeout)
            for _ in range(requests_count)
        ]
        for future in as_completed(futures):
            samples.append(future.result())
    wall_seconds = time.perf_counter() - start
    latencies = [sample["elapsed_seconds"] for sample in samples]
    success_count = sum(1 for sample in samples if sample["ok"])
    return {
        "name": name,
        "url": url,
        "requests": requests_count,
        "concurrency": concurrency,
        "success": success_count,
        "failed": requests_count - success_count,
        "wall_seconds": wall_seconds,
        "throughput_rps": requests_count / wall_seconds if wall_seconds else 0.0,
        "latency_seconds": {
            "min": min(latencies) if latencies else 0.0,
            "mean": statistics.mean(latencies) if latencies else 0.0,
            "p50": _percentile(latencies, 50),
            "p95": _percentile(latencies, 95),
            "max": max(latencies) if latencies else 0.0,
        },
        "errors": [sample for sample in samples if not sample["ok"]][:10],
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="Compare legacy and vLLM OCR endpoint concurrency.")
    parser.add_argument("--image", required=True, help="Path to an input png/jpg image.")
    parser.add_argument("--old-url", default=DEFAULT_OLD_ENDPOINT)
    parser.add_argument("--new-url", default=DEFAULT_NEW_ENDPOINT)
    parser.add_argument("--requests", type=int, default=20)
    parser.add_argument("--concurrency", type=int, default=8)
    parser.add_argument("--timeout", type=float, default=900)
    parser.add_argument("--output", default="concurrency_test_results.json")
    parser.add_argument("--skip-old", action="store_true")
    parser.add_argument("--skip-new", action="store_true")
    args = parser.parse_args()

    image_path = Path(args.image)
    image_type = image_path.suffix.lstrip(".").lower() or "png"
    payload = {
        "file": _encode_image(image_path),
        "type": "jpg" if image_type == "jpeg" else image_type,
        "skip_text_detection": False,
        "skip_table_detection": False,
        "recognize_math": False,
        "ocr_with_boxes": True,
    }

    results = {
        "started_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
        "image": str(image_path),
        "requests": args.requests,
        "concurrency": args.concurrency,
        "preflight": _preflight(),
        "endpoints": [],
    }
    if not args.skip_old:
        results["endpoints"].append(
            run_endpoint("legacy", args.old_url, payload, args.requests, args.concurrency, args.timeout)
        )
    if not args.skip_new:
        results["endpoints"].append(
            run_endpoint("vllm", args.new_url, payload, args.requests, args.concurrency, args.timeout)
        )

    output_path = Path(args.output)
    output_path.write_text(json.dumps(results, indent=2), encoding="utf-8")
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()