"""
llm_guard/multilingual_sc_scorer.py — Multilingual SC Scorer
=============================================================
Drop-in scorer using a multilingual sentence-transformer backbone.
Supports 50+ languages via paraphrase-multilingual-MiniLM-L12-v2.

Stage A (backbone swap): Same 12 MFE Phase A features, but computed using
  multilingual embeddings. Trained on English HP chains, evaluated on
  Arabic/Spanish/Chinese XQuAD chains.

Usage
-----
    from llm_guard.features.multilingual_sc_scorer import MultilingualSCScorer

    scorer = MultilingualSCScorer(model_path="models/multilingual-minilm")
    scorer.fit(hp_chains, hp_labels)            # train on English HP
    risk = scorer.score_chain(arabic_chain)     # cross-lingual inference
Purpose:    Drop-in scorer supporting 50+ languages via paraphrase-multilingual-MiniLM-L12-v2. Trained on English HP chains and evaluated cross-lingually on Arabic/Spanish/Chinese XQuAD.
Exports:    MultilingualSCScorer
Added:      v0.59.0
Status:     Experimental
"""
from __future__ import annotations

from typing import Dict, List, Optional

import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler


class MultilingualSCScorer:
    """
    Multilingual risk scorer using multilingual sentence-transformer backbone.

    Parameters
    ----------
    model_path : str
        Path to multilingual sentence-transformer model (local or HF hub).
    C : float
        Logistic regression regularization. Default: 1.0.
    """

    def __init__(
        self,
        model_path: str,
        C: float = 1.0,
    ) -> None:
        self.model_path = model_path
        self.C = C
        self._extractor = None
        self._clf: Optional[LogisticRegression] = None
        self._scaler: Optional[StandardScaler] = None
        self._fitted = False

    @property
    def is_fitted(self) -> bool:
        return self._fitted

    def _get_extractor(self):
        """Lazy-initialize MechanismFeatureExtractor with multilingual model."""
        if self._extractor is None:
            from llm_guard.features.mechanism_layer import MechanismFeatureExtractor
            self._extractor = MechanismFeatureExtractor(model_name=self.model_path)
        return self._extractor

    def extract_features(self, chain: Dict) -> np.ndarray:
        """
        Extract 12 MFE Phase A features using multilingual backbone.

        Parameters
        ----------
        chain : dict  with keys: question, final_answer, steps

        Returns
        -------
        np.ndarray shape (12,)
        """
        extractor = self._get_extractor()
        ctx = ""
        for s in chain.get("steps", []):
            if s.get("action_type") == "Search" and s.get("observation"):
                ctx = s["observation"][:500]
                break
        return extractor.extract_features(
            chain.get("question", ""),
            chain.get("final_answer", ""),
            context=ctx,
        )

    def fit(
        self,
        chains: List[Dict],
        labels: np.ndarray,
    ) -> "MultilingualSCScorer":
        """
        Fit the scorer on labeled chains (typically English HP).

        Parameters
        ----------
        chains : List[Dict]  labeled agent chains
        labels : ndarray     binary (0=correct, 1=wrong)

        Returns
        -------
        self
        """
        X = np.array([self.extract_features(c) for c in chains])
        y = np.asarray(labels, dtype=np.int64)

        self._scaler = StandardScaler()
        X_scaled = self._scaler.fit_transform(X)

        self._clf = LogisticRegression(
            C=self.C, class_weight="balanced",
            max_iter=1000, solver="lbfgs",
        )
        self._clf.fit(X_scaled, y)
        self._fitted = True
        return self

    def score_chain(self, chain: Dict) -> float:
        """
        Score a chain, returning risk probability in [0, 1].

        Parameters
        ----------
        chain : dict

        Returns
        -------
        float — risk probability (1.0 = high risk of wrong answer)
        """
        if not self._fitted:
            raise RuntimeError("Call fit() before score_chain()")
        feat = self.extract_features(chain).reshape(1, -1)
        feat_scaled = self._scaler.transform(feat)
        return float(self._clf.predict_proba(feat_scaled)[0, 1])
