"""Repeatable end-to-end translation benchmark against local LM Studio."""
import argparse
from contextlib import contextmanager
from datetime import datetime, timezone
import csv
import hashlib
import html
import json
import os
from pathlib import Path
import platform
import shutil
import statistics
import subprocess
import threading
import time
from urllib.parse import urlsplit
from uuid import uuid4

import httpx
from pydantic import Field
import uvicorn

from .api import create_app
from .config import ProviderConfig, Settings
from .models import Project, StrictModel, TextItem
from .storage import Storage
from .provider import Provider


class ModelSpec(StrictModel):
    model: str = Field(min_length=1)
    label: str | None = None
    prompt_style: str | None = None
    temperature: float | None = None
    top_p: float | None = None
    top_k: int | None = None
    repetition_penalty: float | None = None
    structured_output: bool = True


class BenchmarkConfig(StrictModel):
    game: str | None = None
    input: str | None = None
    project_file: str = "examples/tomie_project.json"
    models: list[ModelSpec] = Field(default_factory=list)
    batch_sizes: list[int] = Field(default=[1, 5, 10], min_length=1, max_length=10)
    repeats: int = Field(default=3, ge=1, le=20)
    limit: int = Field(default=30, ge=1, le=500)
    warmup_lines: int = Field(default=2, ge=1, le=20)
    context_length: int = Field(default=8192, ge=512, le=262144)
    gpu: float = Field(default=1, ge=0, le=1)
    reuse_loaded: bool = False
    port: int = Field(default=8788, ge=1024, le=65535)
    job_timeout: float = Field(default=600, gt=0, le=7200)
    retries: int = Field(default=2, ge=0, le=5)
    provider_url: str = "http://127.0.0.1:1234/v1"
    api_key_env: str = "LM_STUDIO_API_KEY"


def write_json(path, value):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(json.dumps(value, ensure_ascii=False, indent=2), encoding="utf-8")
    temporary.replace(path)


class LMS:
    def __init__(self, executable=None):
        self.executable = executable or shutil.which("lms")
        if not self.executable:
            raise RuntimeError("lms CLI not found in PATH. Install the LM Studio CLI first.")

    def run(self, *arguments):
        result = subprocess.run([self.executable, *arguments], capture_output=True, text=True,
                                encoding="utf-8", errors="replace", stdin=subprocess.DEVNULL, timeout=180,
                                creationflags=0x08000000 if os.name == "nt" else 0)
        if result.returncode:
            raise RuntimeError(f"lms {' '.join(arguments[:2])} failed: {(result.stderr or result.stdout)[-1500:]}")
        return result.stdout

    def json(self, command):
        value = json.loads(self.run(command, "--json"))
        if not isinstance(value, list):
            raise RuntimeError(f"Unexpected lms {command} --json response")
        return value


class ModelManager:
    """Switch only idle LLMs and restore the original model IDs on exit."""
    def __init__(self, lms, config):
        self.lms, self.config = lms, config
        self.originals = self.loaded()
        self.owned = set()
        self.restore_errors = []
        self.assert_idle(self.originals)

    def loaded(self):
        return [item for item in self.lms.json("ps") if item.get("type") == "llm"]

    @staticmethod
    def assert_idle(models):
        if any(item.get("status") != "idle" or item.get("queued", 0) for item in models):
            raise RuntimeError("An LM Studio model is busy. Stop other generation jobs before benchmarking.")

    def activate(self, metadata):
        loaded = self.loaded()
        self.assert_idle(loaded)
        selected = next((item for item in loaded if self.config.reuse_loaded and item.get("path") == metadata.get("path")
                         and item.get("contextLength") == self.config.context_length), None)
        for item in loaded:
            if selected and item["identifier"] == selected["identifier"]:
                continue
            print(f"Temporarily unloading idle model: {item['identifier']}", flush=True)
            self.lms.run("unload", item["identifier"])
        if selected:
            return selected, 0.0
        identifier = "lt-bench-" + uuid4().hex[:10]
        # Register before loading so a partial load is also cleaned up on error.
        self.owned.add(identifier)
        started = time.perf_counter()
        key = metadata["modelKey"]
        self.lms.run("load", key, "--identifier", identifier, "--context-length", str(self.config.context_length), "--gpu", str(self.config.gpu), "--yes")
        elapsed = time.perf_counter() - started
        selected = next((item for item in self.loaded() if item["identifier"] == identifier), None)
        if not selected:
            raise RuntimeError("Model was not registered after lms load")
        if selected.get("path") != metadata.get("path"):
            raise RuntimeError("LM Studio loaded a different model file; refusing a mislabeled benchmark")
        return selected, elapsed

    def restore(self):
        try:
            current = self.loaded()
            self.assert_idle(current)
            for item in current:
                if item["identifier"] in self.owned:
                    self.lms.run("unload", item["identifier"])
            identifiers = {item["identifier"] for item in self.loaded()}
            for item in self.originals:
                if item["identifier"] not in identifiers:
                    key = item["modelKey"]
                    self.lms.run("load", key, "--identifier", item["identifier"], "--context-length", str(item["contextLength"]), "--yes")
        except Exception as exc:
            self.restore_errors.append(str(exc))


class CapturingProvider(Provider):
    def __init__(self, config, directory):
        super().__init__(config)
        self.directory = directory

    async def translate(self, *args, **kwargs):
        try:
            return await super().translate(*args, **kwargs)
        finally:
            if self.last_response is not None:
                write_json(self.directory / (uuid4().hex + ".json"), self.last_response)


@contextmanager
def isolated_api(settings, provider_factory=Provider):
    # A dedicated port and DB avoid old queues, cache and project mutations.
    import socket
    listener = socket.socket()
    try:
        listener.bind(("127.0.0.1", settings.port))
    except OSError as exc:
        raise RuntimeError(f"Benchmark port {settings.port} is already in use") from exc
    finally:
        listener.close()
    server = uvicorn.Server(uvicorn.Config(create_app(settings, provider_factory), host="127.0.0.1", port=settings.port,
                                           log_level="warning", access_log=False))
    thread = threading.Thread(target=server.run, daemon=True)
    thread.start()
    try:
        deadline = time.monotonic() + 20
        while not server.started:
            if not thread.is_alive() or time.monotonic() >= deadline:
                raise RuntimeError("Benchmark API failed to start")
            time.sleep(0.05)
        with httpx.Client(base_url=f"http://127.0.0.1:{settings.port}", timeout=30, trust_env=False) as client:
            yield client
    finally:
        server.should_exit = True
        thread.join(timeout=30)
        if thread.is_alive():
            raise RuntimeError("Benchmark API did not stop; do not start another benchmark on this port")


def request(client, method, path, payload=None):
    response = client.request(method, path, json=payload)
    if not response.is_success:
        raise RuntimeError(f"Translation API HTTP {response.status_code}: {response.text[:1000]}")
    return response.json()


def timed_job(client, project_id, items, batch_size, config, on_created=None):
    started = time.perf_counter()
    job = request(client, "POST", f"/projects/{project_id}/jobs", {
        "items": items, "batch_size": batch_size, "use_cache": False, "retries": config.retries,
    })
    if on_created:
        on_created(job["id"])
    deadline = time.monotonic() + config.job_timeout
    while job["status"] in {"queued", "running"}:
        if time.monotonic() >= deadline:
            # Never continue with another model while this job is still running.
            raise TimeoutError(f"Job {job['id']} timed out; aborting benchmark. Results remain in benchmark DB.")
        time.sleep(0.05)
        job = request(client, "GET", "/jobs/" + job["id"])
    return job, time.perf_counter() - started


def summarize_run(job, elapsed):
    rows = job["items"]
    successes = [row for row in rows if row["status"] == "completed"]
    inferred = [row for row in successes if row["attempts"] > 0 and not row["cached"]]
    metrics = job.get("metrics", {})
    provider_seconds = metrics.get("provider_seconds", 0)
    complete_usage = metrics.get("provider_calls", 0) > 0 and metrics.get("calls_with_completion_usage") == metrics.get("provider_calls")
    tokens = metrics.get("completion_tokens") if complete_usage else None
    return {
        "job_id": job["id"], "status": job["status"], "seconds": elapsed,
        "completed": len(successes), "failed": len(rows) - len(successes),
        "model_lines": len(inferred), "cached_lines": sum(bool(row["cached"]) for row in rows),
        "local_lines": sum(row["attempts"] == 0 for row in successes),
        "retry_attempts": sum(max(0, row["attempts"] - 1) for row in rows),
        "source_chars": sum(len(row["source"]) for row in inferred),
        "lines_per_second": len(inferred) / elapsed if elapsed > 0 else 0,
        "chars_per_second": sum(len(row["source"]) for row in inferred) / elapsed if elapsed > 0 else 0,
        "unchanged_ids": [row["id"] for row in inferred if row["source"].strip() == row["translation"].strip()],
        "completion_tokens_per_request_second": tokens / provider_seconds if tokens is not None and provider_seconds > 0 else None,
        "error": job.get("error"), **metrics,
    }


def aggregate(runs):
    result = []
    keys = dict.fromkeys((run["model"], run["batch_size"]) for run in runs)
    for model, batch in keys:
        group = [run for run in runs if run["model"] == model and run["batch_size"] == batch]
        successful = [run for run in group if run["status"] == "completed" and run["failed"] == 0]
        result.append({
            "model": model, "batch_size": batch, "runs": len(group), "successful_runs": len(successful),
            "median_seconds": statistics.median(run["seconds"] for run in successful) if successful else None,
            "median_lines_per_second": statistics.median(run["lines_per_second"] for run in successful) if successful else None,
            "median_chars_per_second": statistics.median(run["chars_per_second"] for run in successful) if successful else None,
            "min_seconds": min((run["seconds"] for run in successful), default=None),
            "max_seconds": max((run["seconds"] for run in successful), default=None),
            "failed_items": sum(run["failed"] for run in group),
            "retry_attempts": sum(run["retry_attempts"] for run in group),
            "unchanged_items": sum(len(run["unchanged_ids"]) for run in group),
        })
    return result


def save_report(directory, report):
    report["summary"] = aggregate(report["runs"])
    write_json(directory / "report.json", report)
    fields = ["model", "batch_size", "repeat", "status", "seconds", "completed", "failed", "model_lines", "cached_lines",
              "local_lines", "retry_attempts", "lines_per_second", "chars_per_second", "provider_calls", "provider_seconds",
              "prompt_tokens", "completion_tokens", "completion_tokens_per_request_second", "job_id"]
    with (directory / "runs.csv").open("w", encoding="utf-8-sig", newline="") as stream:
        writer = csv.DictWriter(stream, fieldnames=fields, extrasaction="ignore")
        writer.writeheader()
        writer.writerows(report["runs"])
    escape = lambda value: html.escape(str(value))
    number = lambda value: f"{value:.2f}" if value is not None else "—"
    rows = "".join(f"<tr><td>{escape(row['model'])}</td><td>{row['batch_size']}</td>"
                   f"<td>{row['successful_runs']}/{row['runs']}</td><td>{number(row['median_seconds'])}</td>"
                   f"<td>{number(row['median_lines_per_second'])}</td><td>{number(row['median_chars_per_second'])}</td>"
                   f"<td>{row['retry_attempts']}</td><td>{row['unchanged_items']}</td></tr>" for row in report["summary"])
    errors = "".join(f"<li>{escape(error)}</li>" for error in report["errors"] + report.get("restore_errors", []))
    source = escape(report["dataset"]["source"])
    document = f"""<!doctype html><html lang="ru"><meta charset="utf-8"><title>Translation benchmark</title>
<style>body{{font:16px system-ui;max-width:1200px;margin:40px auto;padding:0 24px;color:#20242a;background:#fafafa}}
table{{border-collapse:collapse;width:100%;background:white}}td,th{{padding:12px;text-align:left;border-bottom:1px solid #ddd}}
h1{{font-size:26px}}.note{{color:#555;line-height:1.6}}a{{color:#245b8c}}code{{overflow-wrap:anywhere}}</style>
<h1>Скорость перевода игры через Local Translator</h1>
<p class="note">Источник: {source} · {report['dataset']['items']} строк · {report['dataset']['source_chars']} символов · {escape(report['started_at'])}</p>
<p>Показаны медианы успешных повторов. Загрузка модели и прогрев исключены; HTTP, очередь, валидация, запись SQLite и повторы включены.
Кэш отключён. Полные совпадения со словарём удалены из тестового словаря. Неизменённые строки требуют ручной проверки и не считаются автоматически ошибками.</p>
<div style="overflow:auto"><table><thead><tr><th>Модель</th><th>Пачка</th><th>Успешные прогоны</th><th>Медиана, с</th><th>Строк/с</th><th>Символов/с</th><th>Повторы строк</th><th>Без изменений</th></tr></thead><tbody>{rows}</tbody></table></div>
<p class="note">Это проверка скорости и структуры, а не оценка литературного качества или отсутствия цензуры.
Токены/с в JSON/CSV делятся на полное время HTTP-запросов модели, включая обработку промпта; это не чистая скорость декодирования.
Нативный prompt cache LM Studio не сбрасывается. GPU, другие процессы и порядок прогонов влияют на результат.</p>
<ul>{errors}</ul><p><a href="report.json">Полный JSON</a> · <a href="runs.csv">CSV</a> · Переводы каждого прогона находятся в папке jobs.</p></html>"""
    (directory / "report.html").write_text(document, encoding="utf-8")


def profile(spec, config):
    hy = "hy-mt" in spec.model.lower() or "hy_mt" in spec.model.lower()
    values = {"base_url": config.provider_url, "api_key_env": config.api_key_env, "timeout": min(config.job_timeout, 180),
              "model": spec.model, "prompt_style": "hy_mt" if hy else "generic", "temperature": 0.7 if hy else 0.2,
              "structured_output": spec.structured_output}
    if hy:
        values.update(top_p=0.6, top_k=20, repetition_penalty=1.05)
    values.update({key: value for key, value in spec.model_dump().items() if key not in {"model", "label"} and value is not None})
    return ProviderConfig.model_validate(values)


def resolve_metadata(spec, downloaded):
    matches = [item for item in downloaded if spec.model in {item.get("modelKey"), item.get("path"), item.get("indexedModelIdentifier")}]
    if len(matches) != 1:
        raise ValueError(f"Expected one downloaded model for {spec.model}, found {len(matches)}. Use its exact indexedModelIdentifier from lms ls --json.")
    return matches[0]


def load_dataset(config, root, directory):
    if config.input:
        source = (root / config.input).resolve()
    elif config.game:
        # The optional extractor runs independently and never writes into the game.
        source = directory / "input.json"
        import sys
        result = subprocess.run([sys.executable, "-X", "utf8", str(Path(__file__).resolve().parents[1] / "examples/extract_tomie_ui.py"),
                                 str((root / config.game).resolve()), "--output", str(source), "--limit", str(config.limit)],
                                capture_output=True, text=True, encoding="utf-8", errors="replace", timeout=120,
                                creationflags=0x08000000 if os.name == "nt" else 0)
        if result.returncode:
            raise RuntimeError("Game extraction failed. Install optional dependency with pip install -e .[unity].\n" + result.stderr[-2000:])
    else:
        raise ValueError("Set input or game in benchmark configuration")
    raw = json.loads(source.read_text(encoding="utf-8-sig"))
    if not isinstance(raw, list):
        raise ValueError("Dataset must be an array of {id, text} items")
    items = [TextItem.model_validate(item).model_dump() for item in raw][:config.limit]
    items = [item for item in items if item["text"].strip()]
    if not items or len({item["id"] for item in items}) != len(items):
        raise ValueError("Dataset must contain nonempty text and unique IDs")
    if sum(len(item["text"]) for item in items) > 1000000:
        raise ValueError("Dataset exceeds API job limit of 1,000,000 characters")
    write_json(directory / "dataset.json", items)
    return items, {"source": str(source), "items": len(items), "source_chars": sum(len(item["text"]) for item in items),
                   "sha256": hashlib.sha256(json.dumps(items, ensure_ascii=False, sort_keys=True).encode()).hexdigest()}


def run(config, root, directory):
    directory.mkdir(parents=True, exist_ok=False)
    items, dataset = load_dataset(config, root, directory)
    base_project = Project.model_validate(json.loads((root / config.project_file).read_text(encoding="utf-8-sig")))
    source_strings = {item["text"] for item in items}
    removed_keys = [key for key in base_project.glossary if key in source_strings]
    base_project.glossary = {key: value for key, value in base_project.glossary.items() if key not in source_strings}
    url = urlsplit(config.provider_url)
    if url.scheme != "http" or url.hostname not in {"localhost", "127.0.0.1"} or url.username or url.password or url.query or url.fragment:
        raise ValueError("Benchmark model switching requires a local HTTP LM Studio endpoint")
    lms = LMS()
    downloaded = [item for item in lms.json("ls") if item.get("type") == "llm"]
    specs = config.models or [ModelSpec(model=item.get("indexedModelIdentifier") or item["modelKey"]) for item in downloaded]
    if not specs:
        raise RuntimeError("No local language models found. Download a model in LM Studio first.")
    if len({spec.model for spec in specs}) != len(specs):
        raise ValueError("Model list contains duplicates")
    for batch_size in config.batch_sizes:
        if not 1 <= batch_size <= 50:
            raise ValueError("Batch sizes must be between 1 and 50")
    if len(set(config.batch_sizes)) != len(config.batch_sizes):
        raise ValueError("Batch sizes must be unique")
    # Validate every profile before changing loaded models.
    for spec in specs:
        profile(spec, config)
    source_root = Path(__file__).parent
    implementation_hash = hashlib.sha256(b"".join(file.read_bytes() for file in sorted(source_root.glob("*.py")))).hexdigest()
    gpu_info = None
    if shutil.which("nvidia-smi"):
        try:
            gpu_info = subprocess.run(["nvidia-smi", "--query-gpu=name,memory.total,driver_version", "--format=csv,noheader"],
                capture_output=True, text=True, timeout=5, creationflags=0x08000000 if os.name == "nt" else 0).stdout.strip() or None
        except (OSError, subprocess.SubprocessError):
            pass
    report = {"started_at": datetime.now(timezone.utc).isoformat(), "state": "running", "config": config.model_dump(),
              "implementation_sha256": implementation_hash,
              "dataset": dataset, "removed_exact_glossary_keys": removed_keys, "host": {"platform": platform.platform(), "python": platform.python_version(), "nvidia_gpu_info": gpu_info},
              "models": [], "runs": [], "errors": [], "restore_errors": []}
    save_report(directory, report)
    headers = {"Authorization": "Bearer " + os.environ[config.api_key_env]} if os.environ.get(config.api_key_env) else {}
    with httpx.Client(timeout=3, trust_env=False, headers=headers) as probe:
        try:
            response = probe.get(config.provider_url.rstrip("/") + "/models")
            response.raise_for_status()
        except httpx.HTTPError:
            lms.run("daemon", "up")
            lms.run("server", "start", "--port", str(url.port or 1234), "--bind", "127.0.0.1")
    manager = ModelManager(lms, config)
    write_json(directory / "original-models.json", manager.originals)
    try:
        for index, spec in enumerate(specs):
            try:
                metadata = resolve_metadata(spec, downloaded)
            except ValueError as exc:
                report["errors"].append(str(exc))
                save_report(directory, report)
                continue
            label = spec.label or f"{metadata.get('displayName', spec.model)} {metadata.get('quantization', {}).get('name', '')}".strip()
            if any(model["label"] == label for model in report["models"]):
                label += " (" + spec.model + ")"
            model_record = {"label": label, "key": spec.model, "metadata": metadata}
            report["models"].append(model_record)
            print(f"\nModel {index + 1}/{len(specs)}: {label}", flush=True)
            try:
                loaded, load_seconds = manager.activate(metadata)
                provider = profile(spec, config).model_copy(update={"model": loaded["identifier"]})
                model_record.update(loaded=loaded, load_seconds=load_seconds, provider=provider.model_dump())
                settings = Settings(database=str(directory / "translations.sqlite3"), port=config.port, providers={"lmstudio": provider})
                factory = lambda provider_config: CapturingProvider(provider_config, directory / "raw-responses" / f"model-{index}")
                with isolated_api(settings, factory) as client:
                    project = base_project.model_copy(update={"name": "Benchmark: " + label, "provider": "lmstudio", "model": loaded["identifier"]})
                    project_data = request(client, "POST", "/projects", project.model_dump())
                    warmup, warm_seconds = timed_job(client, project_data["id"], items[:config.warmup_lines], min(10, config.warmup_lines), config,
                        lambda identity: write_json(directory / "active-job.json", {"job_id": identity, "phase": "warmup", "model": label}))
                    write_json(directory / "jobs" / f"model-{index}-warmup.json", warmup)
                    model_record["warmup_seconds"] = warm_seconds
                    if warmup["status"] != "completed":
                        raise RuntimeError(f"Warmup failed: {warmup.get('error') or warmup['status']}; see jobs/model-{index}-warmup.json")
                    # Rotate batch order between repeats to reduce fixed-order bias.
                    for repeat in range(config.repeats):
                        order = config.batch_sizes[repeat % len(config.batch_sizes):] + config.batch_sizes[:repeat % len(config.batch_sizes)]
                        for batch_size in order:
                            print(f"  repeat {repeat + 1}/{config.repeats}, batch {batch_size} ...", flush=True)
                            job, elapsed = timed_job(client, project_data["id"], items, batch_size, config,
                                lambda identity: write_json(directory / "active-job.json", {"job_id": identity, "phase": "measured", "model": label}))
                            filename = f"model-{index}-batch-{batch_size}-repeat-{repeat + 1}.json"
                            write_json(directory / "jobs" / filename, job)
                            row = {"model": label, "model_index": index, "batch_size": batch_size, "repeat": repeat + 1,
                                   "result_file": "jobs/" + filename, **summarize_run(job, elapsed)}
                            report["runs"].append(row)
                            save_report(directory, report)
                            print(f"    {row['seconds']:.2f}s, {row['lines_per_second']:.2f} lines/s, failed={row['failed']}, retries={row['retry_attempts']}", flush=True)
            except (TimeoutError, KeyboardInterrupt):
                raise
            except Exception as exc:
                report["errors"].append(f"{label}: {exc}")
                print(f"  FAILED: {exc}", flush=True)
                save_report(directory, report)
    except BaseException as exc:
        report["errors"].append(f"Benchmark interrupted: {type(exc).__name__}: {exc}")
        report["state"] = "interrupted"
        raise
    finally:
        manager.restore()
        report["restore_errors"] = manager.restore_errors
        if report["state"] != "interrupted":
            report["state"] = "completed_with_errors" if report["errors"] or manager.restore_errors or any(row["status"] != "completed" for row in report["runs"]) else "completed"
        report["finished_at"] = datetime.now(timezone.utc).isoformat()
        save_report(directory, report)
    return report


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", type=Path, default=Path("benchmark.json"))
    parser.add_argument("--models", nargs="+", help="Downloaded model keys or paths from lms ls --json")
    parser.add_argument("--input", type=Path)
    parser.add_argument("--limit", type=int)
    parser.add_argument("--repeats", type=int)
    parser.add_argument("--batches", type=int, nargs="+")
    parser.add_argument("--output", type=Path, help="New output directory")
    args = parser.parse_args()
    config_path = args.config.resolve()
    values = json.loads(config_path.read_text(encoding="utf-8-sig"))
    root = config_path.parent
    if args.models:
        values["models"] = [{"model": model} for model in args.models]
    for name in ("limit", "repeats"):
        if getattr(args, name) is not None:
            values[name] = getattr(args, name)
    if args.batches:
        values["batch_sizes"] = args.batches
    if args.input:
        values["input"] = str(args.input.resolve())
    config = BenchmarkConfig.model_validate(values)
    directory = args.output.resolve() if args.output else root / "data" / "benchmarks" / (datetime.now().strftime("%Y%m%d-%H%M%S") + "-" + uuid4().hex[:6])
    lock = Storage(str(root / "data" / "benchmarks" / "runner.sqlite3"))
    with lock.exclusive_worker():
        report = run(config, root, directory)
    print(f"\nReport: {directory / 'report.html'}")
    print(f"State: {report['state']}")
    if report["state"] != "completed":
        raise SystemExit(1)


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        raise SystemExit("Benchmark interrupted; partial results preserved.")
    except (RuntimeError, ValueError, OSError, httpx.HTTPError, subprocess.SubprocessError) as exc:
        raise SystemExit(str(exc)) from exc