"""Complete calibration workflow."""
import logging
from typing import List, Dict, Any, Tuple
from dataclasses import dataclass

from src.email_providers.base import Email
from src.calibration.sampler import EmailSampler
from src.calibration.llm_analyzer import CalibrationAnalyzer
from src.calibration.trainer import ModelTrainer
from src.classification.feature_extractor import FeatureExtractor
from src.llm.base import BaseLLMProvider

logger = logging.getLogger(__name__)


@dataclass
class CalibrationConfig:
    """Calibration workflow configuration."""
    sample_size: int = 1500
    validation_size: int = 300
    llm_batch_size: int = 50
    model_n_estimators: int = 200
    model_learning_rate: float = 0.1
    model_max_depth: int = 8


class CalibrationWorkflow:
    """
    Complete calibration workflow.

    Steps:
    1. Sample emails stratified by sender type
    2. LLM analyzes sample to discover categories
    3. Trainer trains LightGBM on labeled data
    4. Validate on held-out test set
    5. Save trained model
    """

    def __init__(
        self,
        llm_provider: BaseLLMProvider,
        feature_extractor: FeatureExtractor,
        categories: Dict[str, Dict],
        config: CalibrationConfig = None,
        consolidation_llm_provider: BaseLLMProvider = None
    ):
        """Initialize calibration workflow."""
        self.llm_provider = llm_provider
        self.consolidation_llm_provider = consolidation_llm_provider or llm_provider
        self.feature_extractor = feature_extractor
        self.categories = list(categories.keys())
        self.config = config or CalibrationConfig()

        self.sampler = EmailSampler()
        self.analyzer = CalibrationAnalyzer(
            llm_provider,
            {'consolidation_llm': self.consolidation_llm_provider},
            embedding_model=feature_extractor.embedder
        )
        self.trainer = ModelTrainer(feature_extractor, self.categories)

        self.results = {}

    def run(
        self,
        all_emails: List[Email],
        model_output_path: str = None
    ) -> Dict[str, Any]:
        """
        Run complete calibration workflow.

        Returns:
            Workflow results with metrics
        """
        logger.info("=" * 80)
        logger.info("CALIBRATION WORKFLOW")
        logger.info("=" * 80)

        # Step 1: Sample
        logger.info("\nStep 1: Sampling emails...")
        sample_emails, remaining_emails = self.sampler.stratified_sample(
            all_emails,
            self.config.sample_size
        )

        validation_emails = remaining_emails[:self.config.validation_size]
        logger.info(f"Sample: {len(sample_emails)}, Validation: {len(validation_emails)}")

        # Step 2: LLM Analysis
        logger.info("\nStep 2: LLM category discovery...")
        discovered_categories, sample_labels = self.analyzer.discover_categories(sample_emails)

        logger.info(f"ANALYZER RETURNED: {len(discovered_categories)} categories, {len(sample_labels)} labels")
        logger.info(f"Discovered {len(discovered_categories)} categories:")
        for cat, desc in discovered_categories.items():
            logger.info(f"  - {cat}: {desc}")

        if len(sample_labels) > 0:
            logger.info(f"Sample labels (first 3): {sample_labels[:3]}")

        # Step 3: Label emails
        logger.info("\nStep 3: Labeling emails...")

        # Create lookup for LLM labels
        label_map = {email_id: category for email_id, category in sample_labels}

        # Use ONLY LLM-discovered categories for training
        # DO NOT merge self.categories (hardcoded) - those are for rule-based matching only
        label_categories = set(category for _, category in sample_labels)
        all_categories = list(set(discovered_categories.keys()) | label_categories)
        logger.info(f"Using categories (LLM-discovered): {all_categories}")
        logger.info(f"Categories count: {len(all_categories)}")

        # Update trainer with discovered categories
        self.trainer.categories = all_categories
        self.trainer.category_to_idx = {cat: idx for idx, cat in enumerate(all_categories)}
        self.trainer.idx_to_category = {idx: cat for cat, idx in self.trainer.category_to_idx.items()}

        # Build training set
        training_data = []
        matched = 0
        for email in sample_emails:
            category = label_map.get(email.id)
            if category:
                training_data.append((email, category))
                matched += 1

        logger.info(f"Training data: {len(training_data)} labeled emails (matched {matched}/{len(sample_emails)} emails)")

        if not training_data and len(label_map) > 0:
            logger.error(f"CRITICAL: Label ID mismatch! LLM returned {len(label_map)} labels but NONE match email IDs")
            logger.error(f"First 3 email IDs from sample: {[repr(e.id) for e in sample_emails[:3]]}")
            logger.error(f"First 3 label IDs from LLM: {[repr(k) for k in list(label_map.keys())[:3]]}")

            # Check for pattern differences
            if len(label_map) > 0 and len(sample_emails) > 0:
                sample_email_id = sample_emails[0].id
                sample_label_id = list(label_map.keys())[0]
                logger.error(f"Length: email_id={len(sample_email_id)}, label_id={len(sample_label_id)}")
                logger.error(f"Email ID bytes: {sample_email_id.encode()}")
                logger.error(f"Label ID bytes: {sample_label_id.encode()}")

        if not training_data:
            logger.error("No labeled training data!")
            return {'success': False, 'error': 'No labeled data'}

        # Step 4: Train model
        logger.info("\nStep 4: Training LightGBM model...")

        # Prepare validation data
        validation_data = []
        # Use first discovered category as default for validation
        default_category = all_categories[0] if all_categories else 'unknown'
        for email in validation_emails:
            validation_data.append((email, default_category))

        try:
            train_results = self.trainer.train(
                training_data,
                validation_emails=validation_data,
                n_estimators=self.config.model_n_estimators,
                learning_rate=self.config.model_learning_rate,
                max_depth=self.config.model_max_depth
            )

            logger.info(f"Training accuracy: {train_results.get('training_accuracy', 0):.1%}")
            if 'validation_accuracy' in train_results:
                logger.info(f"Validation accuracy: {train_results['validation_accuracy']:.1%}")

            self.results['training_results'] = train_results

        except Exception as e:
            logger.error(f"Training failed: {e}")
            return {'success': False, 'error': str(e)}

        # Step 5: Save model
        if model_output_path:
            logger.info(f"\nStep 5: Saving model to {model_output_path}...")

            try:
                self.trainer.save_model(model_output_path)
                logger.info("Model saved successfully")
                self.results['model_path'] = model_output_path

            except Exception as e:
                logger.error(f"Error saving model: {e}")
                return {'success': False, 'error': f'Save failed: {e}'}

        # Summary
        logger.info("\n" + "=" * 80)
        logger.info("CALIBRATION COMPLETE")
        logger.info(f"Sample size: {len(sample_emails)}")
        logger.info(f"Categories discovered: {len(discovered_categories)}")
        logger.info(f"Training accuracy: {train_results.get('training_accuracy', 0):.1%}")
        logger.info("=" * 80)

        return {
            'success': True,
            'sample_size': len(sample_emails),
            'categories_discovered': discovered_categories,
            'training_results': train_results,
            'model_path': model_output_path
        }

    def get_results(self) -> Dict[str, Any]:
        """Get calibration results."""
        return self.results
