"""
guard_state.py — Serialization for AgentGuard learned state.

File format: single JSON file with two sections:
  - "plain": JSON-serializable scalars (threshold, buffer, version)
  - "sklearn_b64": Base64-encoded pickle of sklearn objects (IsotonicRegression)

Extension: .llmg (LLM Guard state)
"""
from __future__ import annotations

import base64
import json
import logging
import pickle
from pathlib import Path
from typing import Any, Dict, Optional

logger = logging.getLogger(__name__)

_FORMAT_VERSION = 1


def dump_state(state: Dict[str, Any], path: str) -> None:
    """
    Serialize AgentGuard state to *path*.

    Parameters
    ----------
    state : dict with keys:
        alert_threshold : float
        iso_buffer      : list of (float, int) tuples
        iso_calibrator  : IsotonicRegression | None
        ptrue_bandit    : dict | None  (counts, rewards, total)
        version         : str  (llm-guard-kit version string)
    path : str
        File path (use .llmg extension by convention).
    """
    Path(path).parent.mkdir(parents=True, exist_ok=True)

    iso_b64: Optional[str] = None
    iso_obj = state.get("iso_calibrator")
    if iso_obj is not None:
        try:
            iso_b64 = base64.b64encode(pickle.dumps(iso_obj)).decode("ascii")
        except Exception as exc:
            logger.warning("guard_state: could not pickle iso_calibrator: %s", exc)

    payload = {
        "_format": _FORMAT_VERSION,
        "plain": {
            "alert_threshold": state.get("alert_threshold", 0.70),
            "iso_buffer": state.get("iso_buffer", []),
            "ptrue_bandit": state.get("ptrue_bandit"),
            "version": state.get("version", "unknown"),
        },
        "sklearn_b64": iso_b64,
    }

    with open(path, "w", encoding="utf-8") as fh:
        json.dump(payload, fh, indent=2)


def load_state(path: str) -> Optional[Dict[str, Any]]:
    """
    Deserialize AgentGuard state from *path*.

    Returns None on any failure (missing file, corrupt JSON, pickle mismatch).
    Callers should treat None as "no saved state — use defaults".
    """
    try:
        with open(path, "r", encoding="utf-8") as fh:
            payload = json.load(fh)
    except (FileNotFoundError, json.JSONDecodeError) as exc:
        logger.debug("guard_state: load failed (%s): %s", path, exc)
        return None

    try:
        plain = payload["plain"]
        result: Dict[str, Any] = {
            "alert_threshold": float(plain.get("alert_threshold", 0.70)),
            "iso_buffer": [tuple(x) for x in plain.get("iso_buffer", [])],
            "ptrue_bandit": plain.get("ptrue_bandit"),
            "version": plain.get("version", "unknown"),
            "iso_calibrator": None,
        }

        iso_b64 = payload.get("sklearn_b64")
        if iso_b64:
            try:
                result["iso_calibrator"] = pickle.loads(base64.b64decode(iso_b64))
            except Exception as exc:
                logger.warning("guard_state: could not unpickle iso_calibrator: %s", exc)

        return result

    except (KeyError, TypeError, ValueError) as exc:
        logger.warning("guard_state: malformed state file (%s): %s", path, exc)
        return None
