from __future__ import annotations

import argparse
import datetime as dt
import json
import time
import urllib.error
import urllib.request
from pathlib import Path

from common import ROOT, read_json, write_json


DEFAULT_MODEL_ID = "qwen3-coder-30b-a3b-instruct-q6_k"
DEFAULT_PROFILES = ["gpu-fast", "cpu-test"]
DEFAULT_REPORT = ROOT / "reports" / "benchmarks" / "runtime-profiles-qwen3-coder-q6.json"
RUNTIME_PROFILES = ROOT / "config" / "runtime_profiles.json"
DEFAULT_PROMPT = (
    "Ты эксперт 1С. Кратко, но предметно опиши безопасный план анализа ошибки "
    "проведения документа РеализацияТоваровУслуг, если нет metadata snapshot. Дай 8 пунктов."
)


def chat_completion(
    *,
    base_url: str,
    model: str,
    prompt: str,
    temperature: float,
    max_tokens: int,
    timeout: int,
) -> dict:
    payload = {
        "model": model,
        "messages": [
            {"role": "system", "content": "Отвечай по-русски, кратко и по делу."},
            {"role": "user", "content": prompt},
        ],
        "temperature": temperature,
        "max_tokens": max_tokens,
        "stream": False,
    }
    request = urllib.request.Request(
        f"{base_url.rstrip('/')}/v1/chat/completions",
        data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    started_at = time.perf_counter()
    with urllib.request.urlopen(request, timeout=timeout) as response:
        data = json.loads(response.read().decode("utf-8"))
    elapsed_sec = time.perf_counter() - started_at
    choices = data.get("choices") or []
    answer = ""
    if choices and isinstance(choices[0], dict):
        answer = str((choices[0].get("message") or {}).get("content") or "")
    usage = data.get("usage") or {}
    completion_tokens = int(usage.get("completion_tokens") or 0)
    total_tokens = int(usage.get("total_tokens") or 0)
    prompt_tokens = int(usage.get("prompt_tokens") or 0)
    return {
        "status": "ok",
        "elapsed_sec": round(elapsed_sec, 3),
        "prompt_tokens": prompt_tokens,
        "completion_tokens": completion_tokens,
        "total_tokens": total_tokens,
        "output_tokens_per_sec": round(completion_tokens / elapsed_sec, 3) if completion_tokens else None,
        "total_tokens_per_sec": round(total_tokens / elapsed_sec, 3) if total_tokens else None,
        "answer_preview": answer[:800],
    }


def profile_target(profile: dict, model_id: str, plugin: str) -> dict:
    overrides = profile.get("model_overrides") or {}
    override = overrides.get(model_id) or {}
    base_url = override.get("base_url") or (profile.get("endpoints") or {}).get(plugin)
    served_model_name = override.get("served_model_name")
    if not base_url:
        raise ValueError(f"profile `{profile.get('id')}` has no endpoint for model `{model_id}`")
    if not served_model_name:
        raise ValueError(f"profile `{profile.get('id')}` has no served_model_name override for `{model_id}`")
    return {
        "base_url": str(base_url),
        "served_model_name": str(served_model_name),
        "container_name": override.get("container_name"),
        "host": profile.get("host"),
        "docker_endpoint": profile.get("docker_endpoint"),
        "role": profile.get("role"),
    }


def benchmark_profile(profile_id: str, profile: dict, *, model_id: str, plugin: str, args: argparse.Namespace) -> dict:
    try:
        target = profile_target(profile, model_id, plugin)
        result = chat_completion(
            base_url=target["base_url"],
            model=target["served_model_name"],
            prompt=args.prompt,
            temperature=args.temperature,
            max_tokens=args.max_tokens,
            timeout=args.timeout,
        )
        return {
            "profile_id": profile_id,
            "label": profile.get("label") or profile_id,
            "model_id": model_id,
            "target": target,
            **result,
        }
    except (ValueError, TimeoutError, urllib.error.URLError, json.JSONDecodeError) as exc:
        return {
            "profile_id": profile_id,
            "label": profile.get("label") or profile_id,
            "model_id": model_id,
            "status": "error",
            "error": str(exc),
        }


def speedup_summary(results: list[dict]) -> dict:
    speeds = {
        str(item.get("profile_id")): float(item.get("output_tokens_per_sec") or 0)
        for item in results
        if item.get("status") == "ok" and item.get("output_tokens_per_sec")
    }
    gpu_speed = speeds.get("gpu-fast")
    cpu_speed = speeds.get("cpu-test")
    summary = {"output_tokens_per_sec": speeds}
    if gpu_speed and cpu_speed:
        summary["gpu_vs_cpu_ratio"] = round(gpu_speed / cpu_speed, 3)
    return summary


def main() -> int:
    parser = argparse.ArgumentParser(description="Benchmark the same model across runtime profiles.")
    parser.add_argument("--model-id", default=DEFAULT_MODEL_ID)
    parser.add_argument("--plugin", default="1c")
    parser.add_argument("--profiles", nargs="+", default=DEFAULT_PROFILES)
    parser.add_argument("--prompt", default=DEFAULT_PROMPT)
    parser.add_argument("--temperature", type=float, default=0.1)
    parser.add_argument("--max-tokens", type=int, default=384)
    parser.add_argument("--timeout", type=int, default=600)
    parser.add_argument("--report", type=Path, default=DEFAULT_REPORT)
    parser.add_argument("--print", action="store_true")
    args = parser.parse_args()

    config = read_json(RUNTIME_PROFILES)
    profiles = config.get("profiles") or {}
    results = []
    for profile_id in args.profiles:
        profile = profiles.get(profile_id)
        if not isinstance(profile, dict):
            results.append(
                {
                    "profile_id": profile_id,
                    "model_id": args.model_id,
                    "status": "error",
                    "error": f"unknown runtime profile: {profile_id}",
                }
            )
            continue
        profile = {**profile, "id": profile_id}
        results.append(benchmark_profile(profile_id, profile, model_id=args.model_id, plugin=args.plugin, args=args))

    ok_results = [item for item in results if item.get("status") == "ok"]
    fastest = None
    if ok_results:
        fastest = max(ok_results, key=lambda item: float(item.get("output_tokens_per_sec") or 0)).get("profile_id")
    report = {
        "created_at": dt.datetime.now(dt.UTC).isoformat(),
        "model_id": args.model_id,
        "plugin": args.plugin,
        "prompt": args.prompt,
        "temperature": args.temperature,
        "max_tokens": args.max_tokens,
        "status": "ok" if len(ok_results) == len(results) else "partial" if ok_results else "failed",
        "fastest_profile_id": fastest,
        "speedup": speedup_summary(results),
        "results": results,
    }
    write_json(args.report, report)
    if args.print:
        print(json.dumps(report, ensure_ascii=False, indent=2))
    else:
        print(f"Benchmark status: {report['status']}")
        print(f"Wrote report to {args.report}")
    return 0 if ok_results else 1


if __name__ == "__main__":
    raise SystemExit(main())
