"""Train LightGBM model on labeled emails."""
import logging
import numpy as np
import pickle
from typing import List, Dict, Any, Tuple, Optional
from pathlib import Path

try:
    import lightgbm as lgb
except ImportError:
    lgb = None

from src.email_providers.base import Email
from src.classification.feature_extractor import FeatureExtractor

logger = logging.getLogger(__name__)


class ModelTrainer:
    """
    Train LightGBM classifier on labeled emails.

    This trains a REAL model (not mock) on actual email data.
    Used during calibration after LLM labels emails.
    """

    def __init__(
        self,
        feature_extractor: FeatureExtractor,
        categories: List[str],
        config: Dict[str, Any] = None
    ):
        """Initialize trainer."""
        if not lgb:
            raise ImportError("lightgbm not installed")

        self.feature_extractor = feature_extractor
        self.categories = categories
        self.category_to_idx = {cat: idx for idx, cat in enumerate(categories)}
        self.idx_to_category = {idx: cat for cat, idx in self.category_to_idx.items()}
        self.config = config or {}
        self.model = None

    def train(
        self,
        labeled_emails: List[Tuple[Email, str]],
        validation_emails: Optional[List[Tuple[Email, str]]] = None,
        n_estimators: int = 200,
        learning_rate: float = 0.1,
        max_depth: int = 8
    ) -> Dict[str, Any]:
        """
        Train LightGBM model on labeled emails.

        Args:
            labeled_emails: List of (email, category) tuples
            validation_emails: Optional validation set
            n_estimators: Number of boosting rounds
            learning_rate: Learning rate
            max_depth: Maximum tree depth

        Returns:
            Training results with metrics
        """
        logger.info(f"Starting LightGBM training on {len(labeled_emails)} emails")

        # Extract features
        logger.info("Extracting features...")
        X_list = []
        y_list = []

        for email, category in labeled_emails:
            try:
                features = self.feature_extractor.extract(email)
                embedding = features.get('embedding', np.zeros(384))

                # Convert to feature vector
                if hasattr(embedding, 'shape'):
                    x_vector = embedding.flatten()
                else:
                    x_vector = np.array(embedding)

                X_list.append(x_vector)
                y_list.append(self.category_to_idx[category])

            except Exception as e:
                logger.warning(f"Error extracting features for {email.id}: {e}")

        if not X_list:
            raise ValueError("No features extracted")

        X = np.array(X_list)
        y = np.array(y_list)

        logger.info(f"Extracted features: X shape {X.shape}, y shape {y.shape}")

        # Train/validation split
        train_data = lgb.Dataset(
            X, label=y,
            params={'verbose': -1}
        )

        # Optional validation data
        eval_set = None
        val_names = None
        if validation_emails:
            logger.info(f"Preparing validation set with {len(validation_emails)} emails")
            X_val_list = []
            y_val_list = []

            for email, category in validation_emails:
                try:
                    features = self.feature_extractor.extract(email)
                    embedding = features.get('embedding', np.zeros(384))
                    x_vector = embedding.flatten() if hasattr(embedding, 'shape') else np.array(embedding)
                    X_val_list.append(x_vector)
                    y_val_list.append(self.category_to_idx[category])
                except Exception as e:
                    logger.warning(f"Error with validation email {email.id}: {e}")

            if X_val_list:
                X_val = np.array(X_val_list)
                y_val = np.array(y_val_list)
                eval_set = [lgb.Dataset(X_val, label=y_val, reference=train_data)]
                val_names = ['valid']

        # Train model
        logger.info("Training LightGBM classifier...")

        params = {
            'objective': 'multiclass',
            'num_class': len(self.categories),
            'metric': 'multi_logloss',
            'learning_rate': learning_rate,
            'num_leaves': 31,
            'max_depth': max_depth,
            'feature_fraction': 0.8,
            'bagging_fraction': 0.8,
            'bagging_freq': 5,
            'verbose': -1,
            'num_threads': 28
        }

        self.model = lgb.train(
            params,
            train_data,
            num_boost_round=n_estimators,
            valid_sets=eval_set,
            valid_names=val_names,
            callbacks=[
                lgb.log_evaluation(period=50)
            ] if eval_set else None
        )

        logger.info("Training complete")

        # Evaluate on training set
        train_pred = self.model.predict(X)
        train_pred_classes = np.argmax(train_pred, axis=1)
        train_acc = np.mean(train_pred_classes == y)

        results = {
            'training_accuracy': train_acc,
            'n_estimators': self.model.num_trees(),
            'feature_importance': dict(zip(
                [f'feature_{i}' for i in range(X.shape[1])],
                self.model.feature_importance()
            ))
        }

        # Validation accuracy
        if eval_set:
            val_pred = self.model.predict(X_val)
            val_pred_classes = np.argmax(val_pred, axis=1)
            val_acc = np.mean(val_pred_classes == y_val)
            results['validation_accuracy'] = val_acc
            logger.info(f"Training accuracy: {train_acc:.1%}, Validation accuracy: {val_acc:.1%}")
        else:
            logger.info(f"Training accuracy: {train_acc:.1%}")

        return results

    def predict(self, email: Email) -> Dict[str, Any]:
        """Predict category for email using trained model."""
        if not self.model:
            raise ValueError("Model not trained")

        try:
            features = self.feature_extractor.extract(email)
            embedding = features.get('embedding', np.zeros(384))
            x_vector = embedding.flatten() if hasattr(embedding, 'shape') else np.array(embedding)

            # Predict
            probs = self.model.predict([x_vector])[0]
            pred_class = np.argmax(probs)
            category = self.idx_to_category[pred_class]
            confidence = float(probs[pred_class])

            return {
                'category': category,
                'confidence': confidence,
                'probabilities': {
                    self.idx_to_category[i]: float(probs[i])
                    for i in range(len(self.categories))
                }
            }

        except Exception as e:
            logger.error(f"Prediction error: {e}")
            return {
                'category': 'unknown',
                'confidence': 0.0,
                'error': str(e)
            }

    def save_model(self, filepath: str) -> bool:
        """Save trained model to file."""
        if not self.model:
            logger.error("No model to save")
            return False

        try:
            Path(filepath).parent.mkdir(parents=True, exist_ok=True)

            model_data = {
                'model': self.model,
                'categories': self.categories,
                'category_to_idx': self.category_to_idx,
                'is_mock': False,  # This is a REAL model
                'model_type': 'LightGBM'
            }

            with open(filepath, 'wb') as f:
                pickle.dump(model_data, f)

            logger.info(f"Model saved to {filepath}")
            return True

        except Exception as e:
            logger.error(f"Error saving model: {e}")
            return False

    def load_model(self, filepath: str) -> bool:
        """Load pre-trained model from file."""
        try:
            with open(filepath, 'rb') as f:
                model_data = pickle.load(f)

            self.model = model_data['model']
            self.categories = model_data['categories']
            self.category_to_idx = model_data['category_to_idx']
            self.idx_to_category = {idx: cat for cat, idx in self.category_to_idx.items()}

            logger.info(f"Model loaded from {filepath}")
            return True

        except Exception as e:
            logger.error(f"Error loading model: {e}")
            return False

    def get_model_info(self) -> Dict[str, Any]:
        """Get model information."""
        if not self.model:
            return {'status': 'not_trained'}

        return {
            'status': 'trained',
            'model_type': 'LightGBM',
            'categories': self.categories,
            'n_classes': len(self.categories),
            'n_trees': self.model.num_trees(),
            'feature_count': self.model.num_feature()
        }
