"""Managed, loopback-only llama.cpp server used by the validation pipeline."""
import contextlib
import socket
import subprocess
import time
import unicodedata

import requests

from test_equivalence import ROOT


class Server:
    def __init__(self, model, gpu_layers=99, context=4096, mmproj=None):
        self.model, self.gpu_layers, self.context, self.mmproj = model, gpu_layers, context, mmproj

    def __enter__(self):
        (ROOT / "logs").mkdir(exist_ok=True)
        with socket.socket() as sock:
            sock.bind(("127.0.0.1", 0))
            self.port = sock.getsockname()[1]
        self.url = f"http://127.0.0.1:{self.port}"
        self.log = (ROOT / "logs" / f"server-{self.model.stem}.log").open("a")
        args = [str(ROOT / "llama.cpp/build/bin/llama-server"), "-m", str(self.model),
            "--alias", "agnes", "--host", "127.0.0.1", "--port", str(self.port), "-c", str(self.context),
            "-ngl", str(self.gpu_layers), "--parallel", "1", "--no-warmup", "--jinja",
            "--flash-attn", "on", "--split-mode", "layer", "--tensor-split", "1,1",
            "--threads", "6", "--batch-size", "256", "--ubatch-size", "128"]
        if self.mmproj:
            args += ["--mmproj", str(self.mmproj)]
        self.process = subprocess.Popen(args, stdout=self.log, stderr=subprocess.STDOUT)
        try:
            deadline = time.monotonic() + 900
            while time.monotonic() < deadline:
                if self.process.poll() is not None:
                    raise RuntimeError(f"llama-server exited; see {self.log.name}")
                try:
                    if requests.get(self.url + "/health", timeout=3).status_code == 200:
                        return self
                except requests.RequestException:
                    pass
                time.sleep(1)
            raise TimeoutError("llama-server startup timed out")
        except BaseException:
            self.__exit__(None, None, None)
            raise

    def post(self, endpoint, payload):
        response = requests.post(self.url + endpoint, json=payload, timeout=1800)
        response.raise_for_status()
        return response.json()

    def complete(self, prompt, tokens=96, **extra):
        if isinstance(prompt, str):
            prompt = unicodedata.normalize("NFC", prompt)
        return self.post("/completion", {"prompt": prompt, "n_predict": tokens,
            "temperature": 0, "seed": 20260912, "cache_prompt": False, **extra})

    def __exit__(self, *args):
        if hasattr(self, "process") and self.process.poll() is None:
            self.process.terminate()
            try:
                self.process.wait(timeout=30)
            except subprocess.TimeoutExpired:
                self.process.kill()
                self.process.wait(timeout=15)
        if hasattr(self, "log"):
            self.log.close()


def chat(tokenizer, text):
    return unicodedata.normalize("NFC", tokenizer.apply_chat_template([{"role": "user", "content": text}],
        tokenize=False, add_generation_prompt=True, enable_thinking=False))