"""JevBench evaluation using the released runtime or its HTTP endpoint."""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
import platform
import sys
from collections import defaultdict
from dataclasses import dataclass
from importlib.metadata import version
from pathlib import Path
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.parse import urlsplit
from urllib.request import Request, urlopen

if __package__:
    from .serve import (
        LoadedRuntime,
        load_runtime,
        normalize_state,
        parse_question,
        score_question,
    )
else:
    from serve import (
        LoadedRuntime,
        load_runtime,
        normalize_state,
        parse_question,
        score_question,
    )

ECE_BINS = 15
PROBABILITY_FLOOR = 1e-12


@dataclass(frozen=True)
class HttpRuntime:
    """HTTP scoring client with server-advertised inference settings."""

    endpoint: str
    timeout: float
    token: str | None
    settings: dict[str, Any]

    @property
    def max_options(self) -> int:
        return self.settings["max_options"]

    def request(
        self, route: str, payload: dict[str, Any] | None = None
    ) -> dict[str, Any]:
        """Send a bounded JSON request, surfacing transport and protocol failures."""
        headers = {"Content-Type": "application/json"}
        if self.token:
            headers["Authorization"] = f"Bearer {self.token}"
        request = Request(
            self.endpoint + route,
            data=None
            if payload is None
            else json.dumps(payload, allow_nan=False).encode(),
            headers=headers,
        )
        try:
            with urlopen(request, timeout=self.timeout) as response:
                result = json.load(response)
        except HTTPError as exc:
            raise ValueError(f"{route} returned HTTP {exc.code}") from exc
        except (URLError, TimeoutError, OSError) as exc:
            raise ValueError(f"{route}: connection failed or timed out") from exc
        if not isinstance(result, dict):
            raise TypeError(f"{route}: response must be a JSON object")
        return result


def load_http_runtime(endpoint: str, timeout: float) -> HttpRuntime:
    """Connect to a server and validate the metadata needed for evaluation."""
    parsed = urlsplit(endpoint)
    if (
        parsed.scheme not in ("http", "https")
        or not parsed.hostname
        or parsed.username
        or parsed.password
        or parsed.query
        or parsed.fragment
    ):
        raise ValueError(
            "Endpoint must be an HTTP(S) URL without credentials, query or fragment"
        )
    client = HttpRuntime(
        endpoint.rstrip("/"), timeout, os.environ.get("CUEREFLEX_TOKEN"), {}
    )
    settings = client.request("/v1/version")
    if (
        type(settings.get("max_options")) is not int
        or not 2 <= settings["max_options"] <= 26
    ):
        raise ValueError("Server version response requires max_options in 2..26")
    if not isinstance(settings.get("model"), str) or not settings["model"]:
        raise ValueError("Server version response requires a model identity")
    return HttpRuntime(client.endpoint, timeout, client.token, settings)


def read_rows(paths: list[Path]) -> list[dict[str, Any]]:
    """Read JSONL files, rejecting malformed rows and duplicate IDs."""
    rows: list[dict[str, Any]] = []
    seen: set[str] = set()
    for path in paths:
        with path.open(encoding="utf-8") as handle:
            for line_number, line in enumerate(handle, 1):
                if not line.strip():
                    continue
                try:
                    row = json.loads(line)
                    if not isinstance(row, dict):
                        raise TypeError("Each row must be an object")
                    row_id = row.get("id")
                    if not isinstance(row_id, str) or not row_id:
                        raise ValueError("Each row requires a nonempty string ID")
                    if row_id in seen:
                        raise ValueError(f"Duplicate ID: {row_id}")
                except (ValueError, TypeError) as exc:
                    raise ValueError(f"{path}:{line_number}: {exc}") from exc
                seen.add(row_id)
                rows.append(row)
    if not rows:
        raise ValueError("No evaluation rows found")
    return rows


def score_row(
    runtime: LoadedRuntime | HttpRuntime, row: dict[str, Any]
) -> dict[str, Any]:
    """Score only inference fields, then attach the held-out expected answer."""
    labels = row["labels"]
    if (
        not isinstance(labels, list)
        or any(not isinstance(label, str) or not label for label in labels)
        or len(set(labels)) != len(labels)
        or not 2 <= len(labels) <= runtime.max_options
    ):
        raise ValueError("Labels must be distinct strings within the option limit")
    question = row["question"]
    if not isinstance(question, dict):
        raise TypeError("Question must be an object")
    kind = question.get("type")
    if kind == "noul" and labels != ["no", "yes"]:
        raise ValueError("Noul labels must be ['no', 'yes']")
    if kind == "score" and labels != [str(i) for i in range(len(labels))]:
        raise ValueError("Score labels must be consecutive indices starting at '0'")
    expected = row["expected"]
    if kind == "score":
        if type(expected) is not int:
            raise ValueError("Score expected answer must be an integer")
        expected_label = str(expected)
    else:
        if not isinstance(expected, str):
            raise ValueError("Expected answer must be a string label")
        expected_label = expected
    if expected_label not in labels:
        raise ValueError("Expected answer must belong to labels")

    spec = {
        key: question[key]
        for key in ("type", "instructions", "criteria")
        if key in question
    }
    if kind == "choice":
        spec = {**spec, "labels": labels}
    prepared = parse_question(normalize_state(row["state"]), spec, runtime.max_options)
    if kind == "score" and len(prepared.options) != len(labels):
        raise ValueError("Score criteria must align with labels")
    if isinstance(runtime, HttpRuntime):
        response = runtime.request(
            "/v1/systemone",
            {"state": row["state"], "questions": {"decision": spec}},
        )
        if response.get("model") != runtime.settings["model"]:
            raise ValueError("Server model identity changed during evaluation")
        answer = response["answers"]["decision"]
        input_tokens = response["usage"]["input_tokens"]
        if not isinstance(answer, dict) or answer.get("type") != kind:
            raise ValueError("Server answer type does not match question")
        if type(input_tokens) is not int or input_tokens < 0:
            raise ValueError("Server input_tokens must be a nonnegative integer")
    else:
        answer, input_tokens = score_question(runtime, prepared)
    probabilities = answer["probabilities"]
    if not isinstance(probabilities, dict):
        raise TypeError("Response probabilities must be an object")
    if set(probabilities) != set(labels):
        raise ValueError("Response probability labels do not match dataset labels")
    if any(
        isinstance(p, bool)
        or not isinstance(p, (int, float))
        or not math.isfinite(p)
        or not 0 <= p <= 1
        for p in probabilities.values()
    ) or not math.isclose(sum(probabilities.values()), 1.0, abs_tol=1e-3):
        raise ValueError("Response must contain a normalized finite distribution")
    predicted = min(labels, key=lambda label: (-probabilities[label], label))
    return {
        "id": row["id"],
        "kind": kind,
        "expected": expected,
        "predicted": int(predicted) if kind == "score" else predicted,
        "correct": predicted == expected_label,
        "confidence": probabilities[predicted],
        "loss": -math.log(max(probabilities[expected_label], PROBABILITY_FLOOR)),
        "brier": sum(
            (probabilities[label] - float(label == expected_label)) ** 2
            for label in labels
        ),
        "probabilities": probabilities,
        "input_tokens": input_tokens,
    }


def summarize(items: list[dict[str, Any]]) -> dict[str, Any]:
    """Calculate pooled hard-label metrics and equal-width confidence-bin ECE."""
    if not items:
        raise ValueError("Cannot summarize zero rows")
    buckets: dict[int, list[dict[str, Any]]] = defaultdict(list)
    for item in items:
        buckets[min(int(item["confidence"] * ECE_BINS), ECE_BINS - 1)].append(item)
    correct = sum(item["correct"] for item in items)
    return {
        "rows": len(items),
        "correct": correct,
        "accuracy": correct / len(items),
        "loss": sum(item["loss"] for item in items) / len(items),
        "brier": sum(item["brier"] for item in items) / len(items),
        "ece": sum(
            abs(sum(item["correct"] - item["confidence"] for item in bucket))
            for bucket in buckets.values()
        )
        / len(items),
    }


def evaluate(
    runtime: LoadedRuntime | HttpRuntime, rows: list[dict[str, Any]]
) -> dict[str, Any]:
    """Evaluate every row, failing explicitly rather than dropping failed items."""
    items: list[dict[str, Any]] = []
    for index, row in enumerate(rows, 1):
        try:
            item = score_row(runtime, row)
        except (ValueError, KeyError, TypeError, RuntimeError) as exc:
            raise ValueError(f"Evaluation failed for row {row['id']}: {exc}") from exc
        items.append(item)
        print(
            f"[{index}/{len(rows)}] {row['id']}: {item['predicted']}", file=sys.stderr
        )
    kinds = sorted({item["kind"] for item in items})
    return {
        "overall": summarize(items),
        "by_kind": {
            kind: summarize([item for item in items if item["kind"] == kind])
            for kind in kinds
        },
        "items": items,
    }


def main() -> None:
    """Evaluate locally or over HTTP and save a reproducible diagnostic report."""
    parser = argparse.ArgumentParser(description=__doc__)
    source = parser.add_mutually_exclusive_group(required=True)
    source.add_argument("--model-dir", type=Path)
    source.add_argument(
        "--endpoint", help="Server base URL, e.g. http://127.0.0.1:8000"
    )
    parser.add_argument(
        "--timeout", type=float, default=120.0, help="HTTP timeout in seconds"
    )
    parser.add_argument("--base-model", default=None)
    parser.add_argument("--tasks", type=Path, nargs="+", required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    if not math.isfinite(args.timeout) or args.timeout <= 0:
        parser.error("--timeout must be a finite positive number")
    if args.endpoint and args.base_model:
        parser.error("--base-model is only valid with --model-dir")
    try:
        if args.output.exists():
            raise FileExistsError(f"Refusing to overwrite {args.output}")
        rows = read_rows(args.tasks)
        datasets = [
            {
                "path": str(path.resolve()),
                "sha256": hashlib.sha256(path.read_bytes()).hexdigest(),
            }
            for path in args.tasks
        ]
        runtime = (
            load_http_runtime(args.endpoint, args.timeout)
            if args.endpoint
            else load_runtime(args.model_dir.resolve(), args.base_model)
        )
        report = evaluate(runtime, rows)
        if isinstance(runtime, HttpRuntime):
            source_provenance = {
                "mode": "http",
                "endpoint": runtime.endpoint,
                "timeout_seconds": runtime.timeout,
                "server": runtime.settings,
            }
        else:
            source_provenance = {
                "mode": "offline",
                "model_dir": str(args.model_dir.resolve()),
                "model": runtime.model_id,
                "max_seq_length": runtime.max_seq_length,
                "max_options": runtime.max_options,
                "enable_thinking": runtime.enable_thinking,
                "temperatures": runtime.temperatures,
                "packages": {
                    name: version(name) for name in ("torch", "transformers", "peft")
                },
            }
        report["provenance"] = {
            **source_provenance,
            "datasets": datasets,
            "python": platform.python_version(),
            "ece_bins": ECE_BINS,
            "probability_floor_for_loss": PROBABILITY_FLOOR,
            "protocol": "explicit canonical choice labels; not stock typesafe adapter",
        }
        args.output.parent.mkdir(parents=True, exist_ok=True)
        with args.output.open("x", encoding="utf-8") as handle:
            json.dump(report, handle, indent=2, allow_nan=False)
            handle.write("\n")
    except (OSError, ValueError, TypeError, RuntimeError) as exc:
        parser.exit(1, f"Evaluation error: {exc}\n")
    print(json.dumps(report["overall"], indent=2))
    print(f"Report saved to {args.output}")


if __name__ == "__main__":
    main()