"""Main orchestration for complete pipeline."""
import logging
import time
from typing import Dict, List, Any, Optional
from pathlib import Path

from src.utils.config import Config, load_config, load_categories
from src.email_providers.base import Email, ClassificationResult
from src.classification.feature_extractor import FeatureExtractor
from src.classification.ml_classifier import MLClassifier
from src.classification.llm_classifier import LLMClassifier
from src.classification.adaptive_classifier import AdaptiveClassifier
from src.processing.bulk_processor import BulkProcessor
from src.calibration.sampler import EmailSampler
from src.calibration.llm_analyzer import CalibrationAnalyzer
from src.llm.base import BaseLLMProvider
from src.export.exporter import ResultsExporter, ReportGenerator

logger = logging.getLogger(__name__)


class EmailSorterOrchestrator:
    """
    Main orchestrator for complete email sorting pipeline.

    Pipeline phases:
    1. Calibration: Sample emails, discover categories, train ML
    2. Bulk processing: Classify all emails with checkpointing
    3. LLM review: Process uncertain cases
    4. Export: Results, reports, and sync
    """

    def __init__(
        self,
        config: Config,
        llm_provider: Optional[BaseLLMProvider] = None
    ):
        """Initialize orchestrator."""
        self.config = config
        self.llm_provider = llm_provider
        self.categories = load_categories()

        # Components (lazy-initialized)
        self.feature_extractor = None
        self.ml_classifier = None
        self.llm_classifier = None
        self.adaptive_classifier = None
        self.bulk_processor = None
        self.exporter = None
        self.report_generator = None

        self.start_time = None
        self.stats = {}

    def initialize_components(self) -> None:
        """Initialize all classifier components."""
        logger.info("Initializing classifier components...")

        # Feature extraction
        self.feature_extractor = FeatureExtractor(self.config.features.dict())

        # ML Classifier
        self.ml_classifier = MLClassifier()

        # LLM Classifier
        if self.llm_provider:
            self.llm_classifier = LLMClassifier(
                self.llm_provider,
                self.categories,
                self.config.dict()
            )
        else:
            logger.warning("No LLM provider, LLM classification disabled")
            self.llm_classifier = None

        # Adaptive classifier
        self.adaptive_classifier = AdaptiveClassifier(
            self.feature_extractor,
            self.ml_classifier,
            self.llm_classifier,
            self.categories,
            self.config.dict()
        )

        # Bulk processor
        self.bulk_processor = BulkProcessor(
            self.adaptive_classifier,
            batch_size=self.config.processing.batch_size,
            checkpoint_dir=self.config.processing.checkpoint_dir,
            checkpoint_interval=self.config.processing.checkpoint_interval
        )

        # Export
        self.exporter = ResultsExporter(self.config.export.output_dir)
        self.report_generator = ReportGenerator(self.config.export.output_dir)

        logger.info("Initialization complete")

    def run_calibration(self, sample_emails: List[Email]) -> None:
        """Run calibration phase."""
        logger.info("=" * 80)
        logger.info("PHASE 1: CALIBRATION")
        logger.info("=" * 80)

        # Analyze with LLM if available
        if self.llm_classifier and self.llm_classifier.llm_available:
            logger.info("Analyzing sample with LLM...")
            analyzer = CalibrationAnalyzer(self.llm_provider, self.config.dict())
            categories, labels = analyzer.discover_categories(sample_emails)
            logger.info(f"Discovered {len(categories)} categories from LLM")

            # TODO: Use discovered categories to adjust thresholds

        logger.info("Calibration phase complete")

    def run_bulk_processing(
        self,
        emails: List[Email],
        resume: bool = True,
        progress_callback=None
    ) -> List[ClassificationResult]:
        """Run bulk processing phase."""
        logger.info("=" * 80)
        logger.info("PHASE 2: BULK PROCESSING")
        logger.info("=" * 80)

        results, llm_queue = self.bulk_processor.process(
            emails,
            resume=resume,
            progress_callback=progress_callback
        )

        self.stats['results'] = results
        self.stats['llm_queue'] = llm_queue

        return results

    def run_llm_review(self, emails_dict: Dict[str, Email]) -> None:
        """Run LLM review phase for uncertain classifications."""
        if not self.llm_classifier or not self.llm_classifier.llm_available:
            logger.warning("LLM unavailable, skipping LLM review phase")
            return

        logger.info("=" * 80)
        logger.info("PHASE 3: LLM REVIEW")
        logger.info("=" * 80)

        llm_queue = self.stats.get('llm_queue', [])
        if not llm_queue:
            logger.info("No emails need LLM review")
            return

        logger.info(f"Processing {len(llm_queue)} emails with LLM...")

        results = self.stats.get('results', [])

        # Build email dict for lookup
        email_lookup = {e.id: e for e in emails_dict.values()}

        # Update results with LLM review
        for result in results:
            if result.email_id in llm_queue:
                email = email_lookup.get(result.email_id)
                if email:
                    updated_result = self.adaptive_classifier.classify_with_llm(result, email)
                    # Replace in results
                    idx = results.index(result)
                    results[idx] = updated_result

        logger.info("LLM review phase complete")

    def run_export(self, results: List[ClassificationResult]) -> Dict[str, Path]:
        """Run export phase."""
        logger.info("=" * 80)
        logger.info("PHASE 4: EXPORT & REPORTING")
        logger.info("=" * 80)

        # Prepare metadata
        metadata = {
            'total_emails': len(results),
            'accuracy_estimate': self.adaptive_classifier.get_stats().accuracy_estimate(),
            'processing_time': int(time.time() - self.start_time) if self.start_time else 0,
            'classification_stats': {
                'rule_matched': self.adaptive_classifier.get_stats().rule_matched,
                'ml_classified': self.adaptive_classifier.get_stats().ml_classified,
                'llm_classified': self.adaptive_classifier.get_stats().llm_classified,
                'needs_review': self.adaptive_classifier.get_stats().needs_review,
            }
        }

        # Export results
        export_files = {}

        try:
            export_files['json'] = self.exporter.export_json(results, metadata)
            export_files['csv'] = self.exporter.export_csv(results)
            export_files['by_category'] = self.exporter.export_by_category(results)
        except Exception as e:
            logger.error(f"Error during export: {e}")

        # Generate report
        try:
            export_files['report'] = self.report_generator.generate_report(results, metadata)
        except Exception as e:
            logger.error(f"Error generating report: {e}")

        logger.info("Export phase complete")
        return export_files

    def run_full_pipeline(
        self,
        all_emails: List[Email],
        sample_size: int = 1500,
        resume: bool = True,
        progress_callback=None
    ) -> Dict[str, Any]:
        """
        Run complete pipeline from start to finish.

        Args:
            all_emails: All emails to process
            sample_size: Size of calibration sample
            resume: Resume from checkpoint if exists
            progress_callback: Function for progress updates

        Returns:
            Pipeline results
        """
        self.start_time = time.time()

        logger.info("=" * 80)
        logger.info("EMAIL SORTER v1.0 - FULL PIPELINE")
        logger.info(f"Total emails: {len(all_emails):,}")
        logger.info("=" * 80)

        # Initialize
        self.initialize_components()

        # Phase 1: Calibration
        sampler = EmailSampler()
        sample_emails, remaining_emails = sampler.stratified_sample(
            all_emails,
            sample_size
        )

        self.run_calibration(sample_emails)

        # Phase 2: Bulk processing
        results = self.run_bulk_processing(
            all_emails,  # Process all, not just remaining
            resume=resume,
            progress_callback=progress_callback
        )

        # Phase 3: LLM review
        emails_dict = {e.id: e for e in all_emails}
        self.run_llm_review(emails_dict)

        # Phase 4: Export
        export_files = self.run_export(results)

        # Summary
        elapsed = time.time() - self.start_time
        logger.info("=" * 80)
        logger.info("PIPELINE COMPLETE")
        logger.info(f"Time: {elapsed:.1f}s ({elapsed/60:.1f}m)")
        logger.info(f"Accuracy Estimate: {self.adaptive_classifier.get_stats().accuracy_estimate():.1%}")
        logger.info(f"Results: {export_files}")
        logger.info("=" * 80)

        return {
            'success': True,
            'elapsed_time': elapsed,
            'total_emails': len(all_emails),
            'results_processed': len(results),
            'export_files': export_files,
            'stats': self.stats
        }
