"""
Model Server
============

This module provides a model server for loading and serving predictions.
"""

import logging
from pathlib import Path
from typing import Dict, Any, List, Optional
import joblib
import json
import numpy as np
from sklearn.pipeline import Pipeline

logger = logging.getLogger(__name__)


class ModelServer:
    """
    Serves a trained sentiment analysis model for predictions.
    """
    
    def __init__(self, model_path: Path, label_map_path: Optional[Path] = None):
        """
        Initialize the model server.
        
        Args:
            model_path: Path to the trained model file
            label_map_path: Optional path to label map JSON file
        """
        self.model_path = model_path
        self.label_map_path = label_map_path
        self.model: Optional[Pipeline] = None
        self.label_map: Dict[int, str] = {0: "negative", 2: "neutral", 4: "positive"}
        self.model_info: Dict[str, Any] = {}
        
    def load_model(self):
        """Load the model and label map from disk."""
        logger.info(f"Loading model from {self.model_path}")
        
        try:
            self.model = joblib.load(self.model_path)
            logger.info("Model loaded successfully")
            
            # Load label map if provided
            if self.label_map_path and self.label_map_path.exists():
                with open(self.label_map_path, "r", encoding="utf-8") as f:
                    loaded_map = json.load(f)
                    # Convert string keys to int
                    self.label_map = {int(k): v for k, v in loaded_map.items()}
                logger.info(f"Label map loaded: {self.label_map}")
            
            # Load model info if available
            info_path = self.model_path.parent / "training_info.json"
            if info_path.exists():
                with open(info_path, "r", encoding="utf-8") as f:
                    self.model_info = json.load(f)
                logger.info("Model info loaded")
            
        except Exception as e:
            logger.error(f"Failed to load model: {e}")
            raise
    
    def predict(self, texts: List[str]) -> List[Dict[str, Any]]:
        """
        Make predictions on a list of texts.
        
        Args:
            texts: List of texts to classify
            
        Returns:
            List of prediction dictionaries
        """
        if self.model is None:
            raise ValueError("Model not loaded. Call load_model() first.")
        
        try:
            # Make predictions
            predictions = self.model.predict(texts)
            
            # Get decision scores if available
            try:
                decision_scores = self.model.decision_function(texts)
            except Exception:
                decision_scores = None
            
            # Format results
            results = []
            for i, (text, pred) in enumerate(zip(texts, predictions)):
                result = {
                    "text": text,
                    "prediction": int(pred),
                    "sentiment": self.label_map.get(int(pred), "unknown")
                }
                
                # Add confidence score if available
                if decision_scores is not None:
                    if len(decision_scores.shape) == 1:
                        # Binary classification
                        result["confidence"] = float(abs(decision_scores[i]))
                    else:
                        # Multi-class classification
                        result["confidence"] = float(np.max(np.abs(decision_scores[i])))
                
                results.append(result)
            
            return results
            
        except Exception as e:
            logger.error(f"Prediction failed: {e}")
            raise
    
    def predict_single(self, text: str) -> Dict[str, Any]:
        """
        Make a prediction on a single text.
        
        Args:
            text: Text to classify
            
        Returns:
            Prediction dictionary
        """
        results = self.predict([text])
        return results[0]
    
    def get_model_info(self) -> Dict[str, Any]:
        """
        Get information about the loaded model.
        
        Returns:
            Dictionary with model information
        """
        return {
            "model_loaded": self.model is not None,
            "model_path": str(self.model_path),
            "label_map": self.label_map,
            "training_info": self.model_info
        }
    
    def health_check(self) -> Dict[str, Any]:
        """
        Perform a health check on the model server.
        
        Returns:
            Health check status
        """
        is_healthy = self.model is not None
        
        status = {
            "status": "healthy" if is_healthy else "unhealthy",
            "model_loaded": is_healthy,
            "model_path": str(self.model_path)
        }
        
        if is_healthy:
            # Test prediction
            try:
                test_result = self.predict_single("This is a test")
                status["test_prediction"] = test_result
                status["can_predict"] = True
            except Exception as e:
                status["can_predict"] = False
                status["error"] = str(e)
        
        return status


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    # Example usage
    model_path = Path("artifacts_sentiment140/tfidf_svm_sentiment140.joblib")
    label_map_path = Path("artifacts_sentiment140/label_map.json")
    
    if model_path.exists():
        server = ModelServer(model_path, label_map_path)
        server.load_model()
        
        # Test predictions
        test_texts = [
            "I love this product!",
            "This is terrible",
            "It's okay, nothing special"
        ]
        
        results = server.predict(test_texts)
        
        for result in results:
            print(f"\nText: {result['text']}")
            print(f"Sentiment: {result['sentiment']}")
            print(f"Confidence: {result.get('confidence', 'N/A')}")
    else:
        print(f"Model not found at {model_path}")
