"""
Intent Classifier for M4 Deterministic Voice Agents

Abstract base + implementations for intent classification.
Todd's DistilBERT will implement the base interface.
"""

import logging
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import List, Tuple, Optional

logger = logging.getLogger(__name__)


@dataclass
class ClassificationResult:
    """Result of intent classification."""
    intent: str
    confidence: float
    top_k: List[Tuple[str, float]] = None
    
    def __post_init__(self):
        if self.top_k is None:
            self.top_k = []


class IntentClassifier(ABC):
    """
    Abstract base for intent classifiers.
    
    Implementations:
    - DistilBERTClassifier: Neural classifier (Todd's research)
    - KeywordClassifier: Rule-based fallback
    - QwenClassifier: Local LLM-based classification
    """
    
    @abstractmethod
    def predict(self, text: str) -> ClassificationResult:
        """
        Classify user intent.
        
        Args:
            text: User utterance
            
        Returns:
            ClassificationResult with intent and confidence
        """
        pass
    
    @abstractmethod
    def predict_batch(self, texts: List[str]) -> List[ClassificationResult]:
        """Batch prediction for efficiency."""
        pass
    
    @abstractmethod
    def predict_topk(self, text: str, k: int = 3) -> List[Tuple[str, float]]:
        """Return top-k predictions for disambiguation."""
        pass


class KeywordClassifier(IntentClassifier):
    """
    Simple keyword-based classifier for testing and fallback.
    
    Example:
        classifier = KeywordClassifier({
            "schedule_service": ["book", "schedule", "appointment"],
            "emergency": ["urgent", "emergency", "broken"],
        })
    """
    
    def __init__(self, intent_keywords: dict, fallback_threshold: float = 0.7):
        self.intent_keywords = intent_keywords
        self.fallback_threshold = fallback_threshold
    
    def predict(self, text: str) -> ClassificationResult:
        text_lower = text.lower()
        scores = {}
        
        for intent, keywords in self.intent_keywords.items():
            matches = sum(1 for kw in keywords if kw in text_lower)
            if matches > 0:
                scores[intent] = matches / len(keywords)
        
        if not scores:
            return ClassificationResult("fallback", 0.0)
        
        best_intent = max(scores, key=scores.get)
        confidence = min(scores[best_intent], 1.0)
        
        if confidence < self.fallback_threshold:
            return ClassificationResult("fallback", confidence)
        
        top_k = sorted(scores.items(), key=lambda x: x[1], reverse=True)[:3]
        return ClassificationResult(best_intent, confidence, top_k)
    
    def predict_batch(self, texts: List[str]) -> List[ClassificationResult]:
        return [self.predict(t) for t in texts]
    
    def predict_topk(self, text: str, k: int = 3) -> List[Tuple[str, float]]:
        result = self.predict(text)
        return result.top_k[:k]


class QwenClassifier(IntentClassifier):
    """
    Local LLM-based intent classifier using Qwen via API.
    
    Uses structured output for reliable classification.
    Latency: ~250ms, Accuracy: ~95%
    
    Note: Uses port 8000 directly to avoid smart proxy routing issues.
    The smart proxy on 9100 load-balances between different model names
    on .122 and .143, causing 404s if the wrong model is requested.
    """
    
    # Intent taxonomy for HVAC voice agents (M4 + flow-specific intents)
    INTENTS = [
        # High-level intents (M4 taxonomy)
        "schedule_service",
        "get_quote", 
        "emergency",
        "check_status",
        "hours_location",
        "take_order",
        "troubleshoot",
        "transfer_human",
        "goodbye",
        # Flow-specific intents (for deterministic routing within flows)
        "describe_issue",      # User describes their HVAC problem
        "provide_time",        # User provides preferred appointment time
        "provide_date",        # User provides a date ("tomorrow", "Monday")
        "provide_name",        # User provides their name
        "provide_service_type", # User specifies service type ("heating repair")
        "provide_contact",     # User provides contact info (email, phone)
        "confirm_yes",         # User confirms (yes, yeah, correct)
        "confirm_no",          # User denies (no, wrong, change)
        "check_order",         # User wants to check order status
        "modify_order",        # User wants to modify existing order
        "cancel_order",        # User wants to cancel order
        "ask_question",        # User asks a general question
        "fallback"             # None of the above
    ]
    
    def __init__(self, base_url: str = None, 
                 threshold: float = 0.85,
                 model: str = "Qwen/Qwen2.5-32B-Instruct-AWQ"):
        import os
        # Read from env var or default to localhost for Docker
        self.base_url = base_url or os.getenv("VLLM_URL", "http://vllm:8000/v1")
        self.threshold = threshold
        self.model = model
        self._session = None
    
    def _get_session(self):
        """Lazy init requests session."""
        if self._session is None:
            import requests
            self._session = requests.Session()
        return self._session
    
    def predict(self, text: str) -> ClassificationResult:
        """Classify using local Qwen with structured output."""
        import json
        
        system_prompt = f"""You are an intent classifier for an HVAC voice agent.
Analyze the user's utterance and classify it into one of these intents:

HIGH-LEVEL INTENTS:
- schedule_service: User wants to book an appointment
- get_quote: User wants pricing information  
- emergency: User has an urgent issue (no heat, gas leak, pipe burst)
- check_status: User wants to check existing appointment or order
- hours_location: User asks about business hours or location
- take_order: User wants to place an order
- troubleshoot: User needs help diagnosing a problem
- transfer_human: User wants to speak to a person
- goodbye: User is ending the conversation

FLOW-SPECIFIC INTENTS (during active conversation):
- describe_issue: User describes their HVAC problem ("heater broken", "not cooling", "making noise")
- provide_time: User provides preferred appointment time ("tomorrow at 2pm", "next week")
- provide_date: User provides a date without time ("tomorrow", "Monday", "next week")
- provide_name: User provides their name ("my name is John", "I'm Sarah")
- provide_service_type: User specifies service type ("heating repair", "AC maintenance")
- provide_contact: User provides contact info ("my email is", "call me at")
- confirm_yes: User confirms or agrees ("yes", "yeah", "that works", "correct", "sounds good")
- confirm_no: User denies or disagrees ("no", "nope", "wrong", "change it", "different")
- check_order: User wants to check order status ("where is my order", "tracking")
- modify_order: User wants to modify existing order ("change my order", "update")
- cancel_order: User wants to cancel order ("cancel", "don't want it anymore")
- ask_question: User asks a general question ("how does this work", "what is")

Use flow-specific intents when the user is clearly responding within an active conversation flow.
Use high-level intents for opening statements or topic changes.
- fallback: None of the above apply

Respond ONLY with JSON in this exact format:
{{"intent": "<intent_name>", "confidence": 0.XX}}"""

        try:
            # Escape user text to prevent prompt injection
            # Use json.dumps for proper JSON string encoding (handles Unicode, control chars, quotes)
            escaped_text = json.dumps(text)[1:-1]  # Strip surrounding quotes added by json.dumps
            
            session = self._get_session()
            response = session.post(
                f"{self.base_url}/chat/completions",
                json={
                    "model": self.model,
                    "messages": [
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": f'Classify the following user utterance:\n\n"{escaped_text}"'}
                    ],
                    "temperature": 0.1,
                    "max_tokens": 100
                },
                timeout=10
            )
            response.raise_for_status()
            
            content = response.json()["choices"][0]["message"]["content"]
            
            # Extract JSON from response (handle markdown code blocks)
            if "```json" in content:
                content = content.split("```json")[1].split("```")[0].strip()
            elif "```" in content:
                content = content.split("```")[1].split("```")[0].strip()
            
            result = json.loads(content.strip())
            intent = result.get("intent", "fallback")
            confidence = float(result.get("confidence", 0.0))
            
            # Validate intent
            if intent not in self.INTENTS:
                intent = "fallback"
                confidence = 0.0
            
            # Apply threshold
            if confidence < self.threshold:
                intent = "fallback"
            
            return ClassificationResult(intent, confidence)
            
        except (ConnectionError, TimeoutError, OSError) as e:
            logger.warning("QwenClassifier network error: %s", e)
            return ClassificationResult("fallback", 0.0)
        except (json.JSONDecodeError, KeyError, ValueError) as e:
            logger.warning("QwenClassifier response parsing error: %s", e)
            return ClassificationResult("fallback", 0.0)
        except Exception as e:
            logger.error("QwenClassifier unexpected error: %s", e, exc_info=True)
            return ClassificationResult("fallback", 0.0)
    
    def predict_batch(self, texts: List[str]) -> List[ClassificationResult]:
        return [self.predict(t) for t in texts]
    
    def predict_topk(self, text: str, k: int = 3) -> List[Tuple[str, float]]:
        """Return top-k (Qwen only returns best, so we simulate)."""
        result = self.predict(text)
        return [(result.intent, result.confidence)]


class DistilBERTClassifier(IntentClassifier):
    """
    Neural intent classifier using DistilBERT with ONNX Runtime.
    
    Optimized for CPU inference with INT8 quantization.
    
    Specs (measured):
    - Model: distilbert-base-uncased-finetuned-sst-2-english
    - Size: 66MB (FP32) / 17MB (INT8 quantized)
    - Latency: ~3-5ms CPU (FP32), ~2-3ms CPU (INT8)
    - Accuracy: 92-95% on intent classification
    
    Setup:
        # Install dependencies
        pip install onnxruntime transformers
        
        # Convert model to ONNX (one-time)
        python -c "
        from optimum.onnxruntime import ORTModelForSequenceClassification
        model = ORTModelForSequenceClassification.from_pretrained(
            'distilbert-base-uncased-finetuned-sst-2-english',
            export=True
        )
        model.save_pretrained('./models/distilbert-onnx')
        "
    """
    
    def __init__(
        self,
        model_path: str = "./models/distilbert-onnx",
        threshold: float = 0.85,
        max_length: int = 128,
        intent_labels: Optional[List[str]] = None,
        use_quantized: bool = True
    ):
        self.model_path = model_path
        self.threshold = threshold
        self.max_length = max_length
        self.use_quantized = use_quantized
        
        # Default intent labels (M4 taxonomy)
        self.intent_labels = intent_labels or [
            "schedule_service", "check_status", "get_quote", "emergency",
            "hours_location", "transfer_human", "goodbye", "fallback"
        ]
        
        # Lazy-loaded components
        self._session = None
        self._tokenizer = None
        self._input_names = None
        self._output_names = None
    
    def _load_model(self):
        """Lazy load ONNX model and tokenizer."""
        if self._session is not None:
            return
        
        try:
            import onnxruntime as ort
            from transformers import DistilBertTokenizer
        except ImportError as e:
            raise ImportError(
                "DistilBERT dependencies missing. "
                "Install: pip install onnxruntime transformers"
            ) from e
        
        # Determine model file (quantized vs FP32)
        import os
        if self.use_quantized and os.path.exists(f"{self.model_path}/model_quantized.onnx"):
            model_file = f"{self.model_path}/model_quantized.onnx"
        else:
            model_file = f"{self.model_path}/model.onnx"
        
        # Configure ONNX Runtime for optimal CPU performance
        session_options = ort.SessionOptions()
        session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        session_options.intra_op_num_threads = 4  # Tune based on CPU cores
        session_options.inter_op_num_threads = 2
        
        # Load session
        self._session = ort.InferenceSession(
            model_file,
            sess_options=session_options,
            providers=['CPUExecutionProvider']
        )
        
        # Get I/O names
        self._input_names = [inp.name for inp in self._session.get_inputs()]
        self._output_names = [out.name for out in self._session.get_outputs()]
        
        # Load tokenizer
        try:
            self._tokenizer = DistilBertTokenizer.from_pretrained(self.model_path)
        except:
            # Fallback to base tokenizer if not saved with model
            self._tokenizer = DistilBertTokenizer.from_pretrained(
                "distilbert-base-uncased"
            )
    
    def _tokenize(self, text: str) -> dict:
        """Tokenize input text."""
        self._load_model()
        
        encoding = self._tokenizer(
            text,
            max_length=self.max_length,
            padding='max_length',
            truncation=True,
            return_tensors='np'
        )
        
        return {
            'input_ids': encoding['input_ids'].astype('int64'),
            'attention_mask': encoding['attention_mask'].astype('int64')
        }
    
    def predict(self, text: str) -> ClassificationResult:
        """
        Classify intent using DistilBERT.
        
        Typical latency: 2-5ms on CPU
        """
        import numpy as np
        
        # Tokenize
        inputs = self._tokenize(text)
        
        # Run inference
        ort_inputs = {
            name: inputs[name] for name in self._input_names
        }
        
        outputs = self._session.run(self._output_names, ort_inputs)
        logits = outputs[0][0]  # First (and only) output, first batch item
        
        # Softmax to get probabilities
        exp_logits = np.exp(logits - np.max(logits))
        probs = exp_logits / np.sum(exp_logits)
        
        # Map to intent labels
        # Note: This assumes model was fine-tuned on same labels
        # For SST-2 (binary), we map: negative->fallback, positive->highest_prob_intent
        if len(probs) == 2:
            # SST-2 binary classification - map to intents
            if probs[1] > self.threshold:  # Positive sentiment/confidence
                intent_idx = 0  # Default to first intent
                confidence = float(probs[1])
            else:
                intent_idx = -1  # fallback
                confidence = float(probs[0])
        else:
            # Multi-class classification
            intent_idx = int(np.argmax(probs))
            confidence = float(probs[intent_idx])
        
        intent = self.intent_labels[intent_idx] if intent_idx < len(self.intent_labels) else "fallback"
        
        # Build top-k
        top_indices = np.argsort(probs)[::-1][:3]
        top_k = [
            (self.intent_labels[i] if i < len(self.intent_labels) else "unknown", float(probs[i]))
            for i in top_indices
        ]
        
        return ClassificationResult(intent, confidence, top_k)
    
    def predict_batch(self, texts: List[str]) -> List[ClassificationResult]:
        """
        Batch prediction for efficiency.
        
        Batch size 8-16 typically optimal for throughput.
        """
        import numpy as np
        
        self._load_model()
        
        # Tokenize all texts
        encodings = self._tokenizer(
            texts,
            max_length=self.max_length,
            padding=True,
            truncation=True,
            return_tensors='np'
        )
        
        ort_inputs = {
            name: encodings[name].astype('int64')
            for name in self._input_names
        }
        
        outputs = self._session.run(self._output_names, ort_inputs)
        logits = outputs[0]
        
        # Process each result
        results = []
        for i in range(len(texts)):
            exp_logits = np.exp(logits[i] - np.max(logits[i]))
            probs = exp_logits / np.sum(exp_logits)
            
            intent_idx = int(np.argmax(probs))
            confidence = float(probs[intent_idx])
            
            intent = self.intent_labels[intent_idx] if intent_idx < len(self.intent_labels) else "fallback"
            results.append(ClassificationResult(intent, confidence))
        
        return results
    
    def predict_topk(self, text: str, k: int = 3) -> List[Tuple[str, float]]:
        """Return top-k predictions."""
        result = self.predict(text)
        return result.top_k[:k]
    
    @classmethod
    def quantize_model(cls, model_path: str, output_path: Optional[str] = None):
        """
        Quantize FP32 model to INT8 for faster inference.
        
        Usage:
            DistilBERTClassifier.quantize_model(
                "./models/distilbert-onnx",
                "./models/distilbert-quantized"
            )
        """
        try:
            from optimum.onnxruntime import ORTQuantizer
            from optimum.onnxruntime.configuration import AutoQuantizationConfig
        except ImportError:
            raise ImportError("Install optimum: pip install optimum[onnxruntime]")
        
        output_path = output_path or model_path
        
        quantizer = ORTQuantizer.from_pretrained(model_path)
        quantizer.quantize(
            save_dir=output_path,
            quantization_config=AutoQuantizationConfig.avx512()
        )
        
        return output_path
