"""
Hosted calibration service: fit a DeepLocalVerifier on labeled chains and
return the serialized model as base64.

POST /v2/calibrate/fit  (added to server.py)

Input:
    labeled_runs: list of {"question": str, "steps": [...], "correct": bool}
    ptrue_values: optional list[float] for P(True) scores per chain

Output:
    {
        "model_b64":   "<base64-encoded pickle of fitted DeepLocalVerifier>",
        "auroc":       0.80,          # 5-fold CV AUROC (or null if n < 20)
        "n_runs":      200,
        "n_correct":   91,
        "n_wrong":     109,
    }

Security: model_b64 is a pickle — only share with trusted clients.
"""

from __future__ import annotations

import base64
import io
import pickle
from typing import Any, Optional


def fit_verifier(
    labeled_runs: list[dict],
    ptrue_values: Optional[list[float]] = None,
) -> dict[str, Any]:
    """
    Fit a DeepLocalVerifier on labeled chains.

    Parameters
    ----------
    labeled_runs : list[dict]
        Each dict: {"question": str, "steps": list[dict], "correct": bool}
    ptrue_values : list[float] | None
        Optional P(True) score per chain (same order as labeled_runs).

    Returns
    -------
    dict with keys: model_b64, auroc, n_runs, n_correct, n_wrong
    """
    try:
        from llm_guard.deep_verifier import DeepLocalVerifier
        from sklearn.model_selection import StratifiedKFold, cross_val_score
        import numpy as np
    except ImportError as exc:
        raise RuntimeError(
            "DeepLocalVerifier requires: pip install llm-guard-kit[qara]"
        ) from exc

    if not labeled_runs:
        raise ValueError("labeled_runs must be non-empty")

    n_runs = len(labeled_runs)
    n_correct = sum(1 for r in labeled_runs if r.get("correct", False))
    n_wrong = n_runs - n_correct

    # Need at least 10 examples of each class for CV
    if n_correct < 5 or n_wrong < 5:
        raise ValueError(
            f"Need at least 5 correct AND 5 wrong chains. "
            f"Got {n_correct} correct, {n_wrong} wrong."
        )

    verifier = DeepLocalVerifier()
    verifier.fit(labeled_runs)

    # 5-fold CV AUROC if n >= 20 and both classes have >= 4 examples
    auroc = None
    if n_runs >= 20 and n_correct >= 4 and n_wrong >= 4:
        from llm_guard.local_verifier import extract_features
        import numpy as np

        try:
            X = np.array([extract_features(r) for r in labeled_runs])
            y = np.array([1 if r.get("correct", False) else 0 for r in labeled_runs])

            # Use the MLP from DeepLocalVerifier for CV
            from sklearn.neural_network import MLPClassifier
            cv_model = MLPClassifier(
                hidden_layer_sizes=(64, 32),
                max_iter=500,
                random_state=42,
            )
            skf = StratifiedKFold(n_splits=min(5, n_correct, n_wrong), shuffle=True, random_state=42)
            scores = cross_val_score(cv_model, X, y, cv=skf, scoring="roc_auc")
            auroc = float(np.mean(scores))
        except Exception:
            auroc = None  # CV is best-effort

    # Serialize model
    buf = io.BytesIO()
    pickle.dump(verifier, buf)
    model_b64 = base64.b64encode(buf.getvalue()).decode("ascii")

    return {
        "model_b64": model_b64,
        "auroc": auroc,
        "n_runs": n_runs,
        "n_correct": n_correct,
        "n_wrong": n_wrong,
    }


def load_verifier_from_b64(model_b64: str):
    """Deserialize a DeepLocalVerifier from base64-encoded pickle."""
    data = base64.b64decode(model_b64.encode("ascii"))
    return pickle.loads(data)  # noqa: S301 — trusted internal use only
