# -*- coding: utf-8 -*-
"""
Intent-Classifier Runtime.

Laedt das trainierte Intent-Modell (ONNX oder PyTorch)
und klassifiziert Eingabesaetze in Intents + Slots.

Inferenz auf Pi 4: <100ms
Modellgroesse: ~30-50MB

Beispiel:
    ```python
    classifier = IntentClassifier()
    await classifier.load("models/intent")

    result = await classifier.classify("Mach das Licht im Wohnzimmer an")
    # IntentPrediction(
    #     intent="lightswitch",
    #     confidence=0.97,
    #     slots={"room": "Wohnzimmer", "status": "an"},
    # )
    ```
"""

from __future__ import annotations

import json
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from trixy_core.utils.debug import pinfo, pdebug, perror


@dataclass
class IntentPrediction:
    """Ergebnis einer Intent-Klassifikation."""
    intent: str
    confidence: float
    slots: dict[str, str | list[str]] = field(default_factory=dict)
    tones: list[str] = field(default_factory=list)  # Emotionaler Ton des Satzes
    inference_ms: float = 0.0
    # Top-N Alternativen
    alternatives: list[tuple[str, float]] = field(default_factory=list)

    @property
    def is_confident(self) -> bool:
        """True wenn Confidence ueber Schwelle (0.5)."""
        return self.confidence >= 0.5

    @property
    def is_friendly(self) -> bool:
        """True wenn Ton freundlich ist."""
        return any(t in self.tones for t in ("friendly", "polite", "nice", "grateful"))

    @property
    def is_rude(self) -> bool:
        """True wenn Ton unhoeflich/aggressiv ist."""
        return any(t in self.tones for t in ("rude", "aggressive"))

    @property
    def is_urgent(self) -> bool:
        """True wenn dringend."""
        return "urgent" in self.tones

    def to_dict(self) -> dict[str, Any]:
        from trixy_core.nlp.tone import TAG_TO_CATEGORY
        return {
            "intent": self.intent,
            "confidence": self.confidence,
            "slots": self.slots,
            "tones": [
                {"tag": t, "category": TAG_TO_CATEGORY.get(t, "unknown")}
                for t in self.tones
            ],
            "inference_ms": self.inference_ms,
        }


class IntentClassifier:
    """
    Intent-Classifier fuer Trixy.

    Nutzt Sentence-Embeddings + trainiertes Klassifikationsmodell
    fuer schnelle Intent-Erkennung auf CPU.
    """

    def __init__(self) -> None:
        self._encoder = None           # SentenceTransformer oder ONNX-Encoder
        self._tokenizer = None         # HuggingFace Tokenizer (fuer ONNX-Encoder)
        self._classifier = None        # ONNX InferenceSession oder PyTorch
        self._slot_tagger = None       # ONNX InferenceSession fuer BIO-Tagger
        self._bio_tagset = None        # BIOTagSet fuer Tag-Decoding
        self._use_onnx = False
        self._use_onnx_encoder = False
        self._has_bio_tagger = False
        self._intent_names: list[str] = []
        self._slot_names: list[str] = []
        self._embedding_dim: int = 0
        self._loaded = False
        self._model_dir: Path | None = None

    @property
    def is_loaded(self) -> bool:
        """True wenn Modell geladen."""
        return self._loaded

    @property
    def intent_count(self) -> int:
        """Anzahl bekannter Intents."""
        return len(self._intent_names)

    @property
    def intent_names(self) -> list[str]:
        """Liste aller bekannten Intent-Namen."""
        return list(self._intent_names)

    async def load(self, model_dir: str | Path) -> bool:
        """
        Laedt das trainierte Modell.

        Args:
            model_dir: Verzeichnis mit metafile.json, classifier.onnx/pth

        Returns:
            True bei Erfolg
        """
        model_dir = Path(model_dir)
        self._model_dir = model_dir

        # Metadaten laden
        metafile = model_dir / "metafile.json"
        if not metafile.is_file():
            pdebug(f"IntentClassifier: Kein Modell in {model_dir}")
            return False

        try:
            with open(metafile) as f:
                meta = json.load(f)

            self._embedding_dim = meta["embedding_dim"]
            base_model = meta["base_model"]

            # Labels laden
            with open(model_dir / "intent_labels.json") as f:
                self._intent_names = json.load(f)
            self._slot_names = meta.get("slot_names", [])

        except (json.JSONDecodeError, KeyError, OSError) as e:
            perror(f"IntentClassifier: Metadaten-Fehler: {e}")
            return False

        # Encoder laden (ONNX bevorzugt, Fallback SentenceTransformer)
        encoder_dir = model_dir / "encoder_onnx"
        has_onnx_encoder = encoder_dir.is_dir() and any(
            f.suffix == ".onnx" for f in encoder_dir.iterdir() if f.is_file()
        )
        if has_onnx_encoder:
            # ONNX-Encoder (kein PyTorch noetig, ~112MB)
            try:
                from optimum.onnxruntime import ORTModelForFeatureExtraction
                from transformers import AutoTokenizer

                # Quantisiertes Modell bevorzugen
                onnx_files = [f.name for f in encoder_dir.iterdir() if f.suffix == ".onnx"]
                file_name = "model_quantized.onnx" if "model_quantized.onnx" in onnx_files else None
                self._encoder = ORTModelForFeatureExtraction.from_pretrained(
                    str(encoder_dir), file_name=file_name,
                )
                self._tokenizer = AutoTokenizer.from_pretrained(str(encoder_dir))
                self._use_onnx_encoder = True
                pinfo(f"IntentClassifier: ONNX-Encoder geladen ({encoder_dir})")
            except ImportError:
                pdebug("optimum nicht verfuegbar, versuche SentenceTransformer")
            except Exception as e:
                pdebug(f"ONNX-Encoder Fehler: {e}, versuche SentenceTransformer")

        if not self._use_onnx_encoder:
            # Fallback: SentenceTransformer (braucht PyTorch, ~400MB)
            try:
                from sentence_transformers import SentenceTransformer
                self._encoder = SentenceTransformer(base_model)
                pdebug(f"IntentClassifier: SentenceTransformer geladen ({base_model})")
            except ImportError:
                perror("Weder optimum noch sentence-transformers installiert")
                return False

        # Classifier laden (ONNX bevorzugt, Fallback PyTorch)
        onnx_path = model_dir / "intent_classifier.onnx"
        pth_path = model_dir / "classifier.pth"

        if onnx_path.is_file():
            try:
                import onnxruntime as ort
                self._classifier = ort.InferenceSession(
                    str(onnx_path),
                    providers=["CPUExecutionProvider"],
                )
                self._use_onnx = True
                pdebug(f"IntentClassifier: ONNX Modell geladen")
            except ImportError:
                pdebug("onnxruntime nicht verfuegbar, versuche PyTorch")

        if not self._use_onnx and pth_path.is_file():
            try:
                import torch
                num_intents = len(self._intent_names)
                self._classifier = torch.nn.Sequential(
                    torch.nn.Linear(self._embedding_dim, 256),
                    torch.nn.ReLU(),
                    torch.nn.Dropout(0.3),
                    torch.nn.Linear(256, num_intents),
                )
                self._classifier.load_state_dict(
                    torch.load(pth_path, map_location="cpu", weights_only=True)
                )
                self._classifier.eval()
                pdebug(f"IntentClassifier: PyTorch Modell geladen")
            except Exception as e:
                perror(f"IntentClassifier: PyTorch laden fehlgeschlagen: {e}")
                return False

        if self._classifier is None:
            perror(f"IntentClassifier: Kein Modell gefunden in {model_dir}")
            return False

        # BIO Slot-Tagger laden (optional)
        bio_onnx = model_dir / "slot_tagger.onnx"
        bio_labels = model_dir / "bio_labels.json"
        if bio_onnx.is_file() and bio_labels.is_file():
            try:
                import onnxruntime as ort
                self._slot_tagger = ort.InferenceSession(
                    str(bio_onnx), providers=["CPUExecutionProvider"],
                )
                with open(bio_labels) as f:
                    from trixy_core.trainer.core.intent.bio_tagger import BIOTagSet
                    self._bio_tagset = BIOTagSet.from_dict(json.load(f))
                self._has_bio_tagger = True
                pinfo(f"IntentClassifier: BIO Slot-Tagger geladen ({self._bio_tagset.num_tags} Tags)")
            except Exception as e:
                pdebug(f"BIO Slot-Tagger nicht geladen: {e}")

        self._loaded = True
        pinfo(
            f"IntentClassifier geladen: {len(self._intent_names)} Intents"
            f"{', ONNX-Encoder' if self._use_onnx_encoder else ', SentenceTransformer'}"
            f"{', BIO-Slots' if self._has_bio_tagger else ''}"
        )
        return True

    async def classify(
        self,
        text: str,
        min_confidence: float = 0.3,
        top_n: int = 3,
    ) -> IntentPrediction:
        """
        Klassifiziert einen Text in Intent + Slots.

        Args:
            text: Eingabetext
            min_confidence: Minimale Confidence fuer gueltigen Intent
            top_n: Anzahl Alternativen zurueckgeben

        Returns:
            IntentPrediction
        """
        if not self._loaded:
            return IntentPrediction(intent="unknown", confidence=0.0)

        start = time.monotonic()

        # 1. Embedding berechnen
        import numpy as np

        if self._use_onnx_encoder:
            # ONNX-Encoder: Tokenize → Encoder → Mean-Pooling
            inputs = self._tokenizer(
                text, return_tensors="np",
                padding=True, truncation=True, max_length=128,
            )
            outputs = self._encoder(**inputs)
            # Mean-Pooling ueber Token-Embeddings (ohne Padding)
            token_embeddings = outputs.last_hidden_state[0]  # (seq_len, dim)
            attention_mask = inputs["attention_mask"][0]      # (seq_len,)
            mask = attention_mask.astype(np.float32)
            masked = token_embeddings * mask[:, np.newaxis]
            embedding = masked.sum(axis=0) / mask.sum()
            embedding = embedding.reshape(1, -1).astype(np.float32)
        else:
            # SentenceTransformer: Direkte Embedding-Berechnung
            embedding = self._encoder.encode([text], show_progress_bar=False)
            embedding = np.array(embedding, dtype=np.float32)

        # 2. Classifier ausfuehren
        if self._use_onnx:
            outputs = self._classifier.run(None, {"embedding": embedding})
            logits = outputs[0][0]
        else:
            import torch
            with torch.no_grad():
                logits = self._classifier(torch.FloatTensor(embedding))
                logits = logits[0].numpy()

        # 3. Softmax + Top-N
        exp_logits = np.exp(logits - np.max(logits))
        probabilities = exp_logits / exp_logits.sum()

        top_indices = np.argsort(probabilities)[::-1][:top_n]
        best_idx = top_indices[0]
        best_confidence = float(probabilities[best_idx])
        best_intent = self._intent_names[best_idx]

        # Negativ-Klasse pruefen
        if best_intent == "__negative__" or best_confidence < min_confidence:
            best_intent = "unknown"
            best_confidence = 0.0

        # Alternativen
        alternatives = [
            (self._intent_names[idx], float(probabilities[idx]))
            for idx in top_indices[1:]
            if self._intent_names[idx] != "__negative__"
        ]

        # 4. Slot-Extraktion
        slots: dict[str, Any] = {}
        if best_intent != "unknown":
            # BIO-Tagger hat Vorrang (neuronale Slot-Extraktion)
            if self._has_bio_tagger and self._use_onnx_encoder:
                slots = self._extract_slots_bio(token_embeddings, inputs)
            # Fallback: Regelbasiert
            if not slots:
                slots = self._extract_slots(text)

        # 5. Tone-Analyse (Schlagwort + Satzstruktur)
        from trixy_core.nlp.tone import analyze_tone
        tone_results = analyze_tone(text)
        tones = [r.tag for r in tone_results]

        elapsed = (time.monotonic() - start) * 1000

        return IntentPrediction(
            intent=best_intent,
            confidence=best_confidence,
            slots=slots,
            tones=tones,
            inference_ms=elapsed,
            alternatives=alternatives,
        )

    def _extract_slots_bio(
        self, token_embeddings: "np.ndarray", inputs: dict,
    ) -> dict[str, Any]:
        """
        Extrahiert Slots via BIO-Tagger auf Token-Embeddings.

        Args:
            token_embeddings: (seq_len, 384) vom Encoder
            inputs: Tokenizer-Output mit input_ids

        Returns:
            Dict mit extrahierten Slots
        """
        import numpy as np

        if not self._slot_tagger or not self._bio_tagset:
            return {}

        try:
            # BIO-Logits berechnen: (1, seq_len, num_tags)
            token_emb_batch = token_embeddings.reshape(1, *token_embeddings.shape).astype(np.float32)
            bio_outputs = self._slot_tagger.run(
                None, {"token_embeddings": token_emb_batch},
            )
            bio_logits = bio_outputs[0][0]  # (seq_len, num_tags)

            # Argmax → Tag-Indizes
            tag_indices = bio_logits.argmax(axis=1).tolist()  # (seq_len,)

            # Token-IDs fuer Decoding
            token_ids = inputs["input_ids"][0].tolist() if hasattr(inputs["input_ids"], "tolist") else list(inputs["input_ids"][0])

            # BIO-Tags decodieren → Slots
            from trixy_core.trainer.core.intent.bio_tagger import decode_bio_tags
            slots = decode_bio_tags(
                tag_indices, token_ids, self._bio_tagset, self._tokenizer,
            )

            if slots:
                pdebug(f"IntentClassifier BIO-Slots: {slots}")

            return slots

        except Exception as e:
            pdebug(f"BIO Slot-Extraktion Fehler: {e}")
            return {}

    def _extract_slots(self, text: str) -> dict[str, str | list[str]]:
        """
        Extrahiert Slots aus dem Text via Pattern-Matching.

        Nutzt die Slot-Listen um bekannte Werte zu finden.
        """
        slots: dict[str, str | list[str]] = {}
        text_lower = text.lower()

        # Slot-Listen laden (gecacht aus Metafile)
        if not self._model_dir:
            return slots

        slot_file = self._model_dir / ".." / ".." / "config" / "slot_lists.json"
        if not slot_file.is_file():
            # Fallback: Default Slot-Registry
            from trixy_core.trainer.core.intent.slot_lists import SlotRegistry
            registry = SlotRegistry()
            for slot_name in self._slot_names:
                values = registry.get_values(slot_name)
                found = self._find_values_in_text(text_lower, values)
                if found:
                    slots[slot_name] = found if len(found) > 1 else found[0]
            return slots

        try:
            with open(slot_file) as f:
                slot_data = json.load(f)

            for slot_name in self._slot_names:
                entry = slot_data.get(slot_name, {})
                values = entry.get("values", []) if isinstance(entry, dict) else entry
                found = self._find_values_in_text(text_lower, values)
                if found:
                    slots[slot_name] = found if len(found) > 1 else found[0]

        except (json.JSONDecodeError, OSError):
            pass

        return slots

    @staticmethod
    def _find_values_in_text(text: str, values: list[str]) -> list[str]:
        """Findet Slot-Werte im Text (case-insensitive)."""
        found = []
        # Laengste Werte zuerst (damit "Wohnzimmer" vor "Wohn" matched)
        for value in sorted(values, key=len, reverse=True):
            if value.lower() in text:
                found.append(value)
                # Wert aus Text entfernen um Doppel-Matches zu vermeiden
                text = text.replace(value.lower(), "", 1)
        return found