from __future__ import annotations

import pickle
from dataclasses import dataclass
from typing import Iterable, Tuple

import numpy as np

try:
    from sklearn.linear_model import LogisticRegression
    from sklearn.preprocessing import StandardScaler
except Exception:  # pragma: no cover
    LogisticRegression = None
    StandardScaler = None

try:
    import joblib
except Exception:  # pragma: no cover
    joblib = None


@dataclass
class MonitorOutput:
    is_safe: bool
    spectral_hazard_score: float


class SpectralGuardDetector:
    """Detector with a stable public contract for spectral safety monitoring."""

    def __init__(self, threshold: float = 0.5):
        self.threshold = float(threshold)
        self.scaler = None
        self.model = None

    @staticmethod
    def _feature_vector(hidden_states) -> np.ndarray:
        if isinstance(hidden_states, dict):
            if "rho_layers" in hidden_states:
                rho = np.asarray(hidden_states["rho_layers"], dtype=float)
            elif "rho" in hidden_states:
                rho = np.asarray(hidden_states["rho"], dtype=float)
            else:
                raise ValueError("hidden_states dict must provide 'rho_layers' or 'rho'.")
        else:
            rho = np.asarray(hidden_states, dtype=float)

        if rho.ndim != 1:
            raise ValueError("hidden_states must map to a 1D layer-wise rho vector.")
        if rho.shape[0] == 0:
            raise ValueError("hidden_states must contain at least one layer value.")
        if np.isnan(rho).any():
            raise ValueError("hidden_states contains NaN values.")

        mean = float(np.mean(rho))
        std = float(np.std(rho))
        gap = float(np.max(rho) - np.min(rho))
        return np.concatenate([rho, [mean, std, gap]])

    def fit(self, hidden_state_vectors: Iterable, labels: Iterable[int]) -> "SpectralGuardDetector":
        if LogisticRegression is None or StandardScaler is None:
            raise ImportError("scikit-learn is required for fit(). Install dependencies from requirements.txt.")

        X = np.array([self._feature_vector(x) for x in hidden_state_vectors], dtype=float)
        y = np.asarray(list(labels), dtype=int)
        if X.shape[0] == 0:
            raise ValueError("fit() requires non-empty training data.")

        self.scaler = StandardScaler().fit(X)
        Xs = self.scaler.transform(X)
        self.model = LogisticRegression(class_weight="balanced", max_iter=1000)
        self.model.fit(Xs, y)
        return self

    def monitor(self, prompt: str, hidden_states) -> Tuple[bool, float]:
        if not isinstance(prompt, str):
            raise ValueError("prompt must be a string.")
        feat = self._feature_vector(hidden_states)

        if self.model is None or self.scaler is None:
            # Heuristic fallback for out-of-the-box deployment.
            mean = feat[-3]
            std = feat[-2]
            gap = feat[-1]
            score = float(np.clip((0.95 - mean) + 0.5 * std + 0.25 * gap, 0.0, 1.0))
        else:
            xs = self.scaler.transform(feat.reshape(1, -1))
            score = float(self.model.predict_proba(xs)[0, 1])

        is_safe = bool(score < self.threshold)
        return is_safe, score

    def save(self, path: str) -> None:
        payload = {"threshold": self.threshold, "scaler": self.scaler, "model": self.model}
        if joblib is not None:
            joblib.dump(payload, path)
            return
        with open(path, "wb") as handle:
            pickle.dump(payload, handle)

    @classmethod
    def load(cls, path: str) -> "SpectralGuardDetector":
        if joblib is not None:
            blob = joblib.load(path)
        else:
            with open(path, "rb") as handle:
                blob = pickle.load(handle)
        inst = cls(threshold=blob["threshold"])
        inst.scaler = blob["scaler"]
        inst.model = blob["model"]
        return inst


_DEFAULT_DETECTOR = SpectralGuardDetector()


def monitor(prompt: str, hidden_states) -> Tuple[bool, float]:
    """Official package API: monitor(prompt, hidden_states)."""
    return _DEFAULT_DETECTOR.monitor(prompt, hidden_states)


def set_default_detector(detector: SpectralGuardDetector) -> None:
    if not isinstance(detector, SpectralGuardDetector):
        raise ValueError("detector must be an instance of SpectralGuardDetector.")
    global _DEFAULT_DETECTOR
    _DEFAULT_DETECTOR = detector