"""Open board plumbing: generic chat prompt / stop ids / context limits, failure kinds,
the local OpenAI-compatible endpoint system, and the board's no-reweighting rows."""

import json
from pathlib import Path

from deg.bench.suite_open import failure, usage
from deg.bench.systems.constrained import chat_prompt, native_context, stop_ids


class FakeTok:
    unk_token_id = 0
    eos_token_id = 2

    def __init__(self, system_ok: bool = True, vocab: dict | None = None) -> None:
        self.system_ok = system_ok
        self.vocab = {"<|im_end|>": 7} if vocab is None else vocab

    def convert_tokens_to_ids(self, t):
        return self.vocab.get(t, self.unk_token_id)

    def apply_chat_template(self, messages, tokenize, add_generation_prompt, enable_thinking):
        if not self.system_ok and messages[0]["role"] == "system":
            raise ValueError("System role not supported")
        return "|".join(f"{m['role']}:{m['content'][:12]}" for m in messages)


REC = {"uid": "u", "context": "Order AB-1.", "schema": {"type": "object"}}


def test_chat_prompt_system_and_fallback():
    assert chat_prompt(FakeTok(), REC).startswith("system:Fill the JSO|user:SCHEMA:")
    # a template without a system role gets the instruction inside the user turn
    assert chat_prompt(FakeTok(system_ok=False), REC).startswith("user:Fill the JSO")


def test_stop_ids_and_native_context(tmp_path: Path):
    (tmp_path / "generation_config.json").write_text(json.dumps({"eos_token_id": [2, 106]}))
    (tmp_path / "config.json").write_text(json.dumps(
        {"text_config": {"max_position_embeddings": 32768}, "max_position_embeddings": 99}))
    assert stop_ids(FakeTok(), tmp_path) == [7, 2, 106]
    assert stop_ids(FakeTok(vocab={}), tmp_path) == [2, 106]  # no <|im_end|> in vocab
    assert native_context(tmp_path) == 32768


def test_failure_kinds():
    assert failure("context_exceeded: native 300000 prompt tokens >= 262144") == "ctx_native"
    assert failure("context_exceeded: server n_ctx 40000 prompt tokens >= 32768") == "ctx_server"
    assert failure("context_exceeded: HTTP 400: maximum context length is 12288") == "ctx_native"
    assert failure("OutOfMemoryError: CUDA out of memory") == "oom"
    assert failure("json: Expecting value") == "other"


def test_usage_counts_latest_prediction_per_uid(tmp_path: Path):
    rows = [{"uid": "a", "error": "OutOfMemoryError: CUDA out of memory"},
            {"uid": "a", "value": {}, "latency_s": 1.0, "generated_tokens": 10},
            {"uid": "b", "error": "context_exceeded: native 9 >= 8"}]
    (tmp_path / "extractbench.preds.jsonl").write_text(
        "\n".join(json.dumps(r) for r in rows) + "\n")
    u = usage(tmp_path)
    assert u["records"] == 2 and u["failures"] == {"ctx_native": 1}
    assert u["extractbench_failures"] == {"ctx_native": 1} and u["gen_tok_mean"] == 10


def test_open_endpoint_request(monkeypatch):
    from deg.bench.systems import frontier

    sys_ = frontier.OpenEndpointSystem("127.0.0.1:8780", "M", workers=2)
    assert sys_.url == "http://127.0.0.1:8780/v1/chat/completions" and sys_.name == "open-M"
    sent = {}

    class Resp:
        def __enter__(self):
            return self

        def __exit__(self, *a):
            return False

        def read(self):
            return json.dumps({"choices": [{"message": {"content": '{"id": "AB-1"}'}}],
                               "usage": {"completion_tokens": 5, "prompt_tokens": 50}}).encode()

    def fake_open(req, timeout):
        sent.update(json.loads(req.data))
        return Resp()

    monkeypatch.setattr(frontier.urllib.request, "urlopen", fake_open)
    rec = {"uid": "u", "context": "Order AB-1.",
           "schema": {"type": "object", "properties": {"id": {"type": "string"}}}}
    p = sys_.predict(rec)
    assert p.value == {"id": "AB-1"} and p.generated_tokens == 5
    assert sent["temperature"] == 0.0
    assert sent["chat_template_kwargs"] == {"enable_thinking": False}
    assert sent["response_format"]["json_schema"]["strict"] is True