"""Command-line interface for email-sorter."""
import logging
import sys
from pathlib import Path
from typing import Optional

import click

from src.utils.config import load_config, load_categories
from src.utils.logging import setup_logging
from src.email_providers.base import MockProvider
from src.email_providers.gmail import GmailProvider
from src.email_providers.imap import IMAPProvider
from src.email_providers.enron import EnronProvider
from src.email_providers.outlook import OutlookProvider
from src.email_providers.local_file import LocalFileProvider
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.llm.ollama import OllamaProvider
from src.llm.openai_compat import OpenAIProvider


@click.group()
def cli():
    """Email Sorter - Hybrid ML/LLM Email Classification System."""
    pass


@cli.command()
@click.option('--source', type=click.Choice(['gmail', 'outlook', 'imap', 'mock', 'enron', 'local']), default='mock',
              help='Email provider')
@click.option('--credentials', type=click.Path(exists=False),
              help='Path to credentials file')
@click.option('--directory', type=click.Path(exists=True),
              help='Directory path for local file provider (.msg/.eml files)')
@click.option('--output', type=click.Path(), default='results/',
              help='Output directory')
@click.option('--config', type=click.Path(exists=False), default='config/default_config.yaml',
              help='Config file')
@click.option('--limit', type=int, default=None,
              help='Limit number of emails')
@click.option('--llm-provider', type=click.Choice(['ollama', 'openai']), default='ollama',
              help='LLM provider')
@click.option('--dry-run', is_flag=True,
              help='Do not sync results back')
@click.option('--verbose', is_flag=True,
              help='Verbose logging')
@click.option('--no-llm-fallback', is_flag=True,
              help='Disable LLM fallback - test pure ML performance')
@click.option('--verify-categories', is_flag=True,
              help='Verify model categories fit new mailbox (single LLM call)')
@click.option('--verify-sample', type=int, default=20,
              help='Number of emails to sample for category verification')
def run(
    source: str,
    credentials: Optional[str],
    directory: Optional[str],
    output: str,
    config: str,
    limit: Optional[int],
    llm_provider: str,
    dry_run: bool,
    verbose: bool,
    no_llm_fallback: bool,
    verify_categories: bool,
    verify_sample: int
):
    """Run email sorter pipeline."""

    # Setup logging
    log_level = 'DEBUG' if verbose else 'INFO'
    setup_logging(level=log_level)
    logger = logging.getLogger(__name__)

    logger.info("=" * 80)
    logger.info("EMAIL SORTER v1.0")
    logger.info("=" * 80)

    # Load config
    logger.info(f"Loading config from: {config}")
    cfg = load_config(config)
    categories = load_categories()

    # Setup email provider
    logger.info(f"Setting up email provider: {source}")
    if source == 'gmail':
        provider = GmailProvider()
        if not credentials:
            logger.error("Gmail provider requires --credentials")
            sys.exit(1)
    elif source == 'outlook':
        provider = OutlookProvider()
        if not credentials:
            logger.error("Outlook provider requires --credentials")
            sys.exit(1)
    elif source == 'imap':
        provider = IMAPProvider()
        if not credentials:
            logger.error("IMAP provider requires --credentials")
            sys.exit(1)
    elif source == 'enron':
        provider = EnronProvider(maildir_path=".")
        credentials = None
    elif source == 'local':
        if not directory:
            logger.error("Local file provider requires --directory")
            sys.exit(1)
        provider = LocalFileProvider(directory_path=directory)
        credentials = None
    else:  # mock
        logger.warning("Using MOCK provider for testing")
        provider = MockProvider()
        credentials = None

    # Connect provider
    creds_dict = {'credentials_path': credentials} if credentials else {}
    if not provider.connect(creds_dict):
        logger.error("Failed to connect to email provider")
        sys.exit(1)

    # Setup LLM provider
    logger.info(f"Setting up LLM provider: {llm_provider}")
    if llm_provider == 'ollama':
        llm = OllamaProvider(
            base_url=cfg.llm.ollama.base_url,
            model=cfg.llm.ollama.classification_model,
            temperature=cfg.llm.ollama.temperature,
            max_tokens=cfg.llm.ollama.max_tokens
        )
    else:  # openai
        llm = OpenAIProvider(
            base_url=cfg.llm.openai.base_url,
            model=cfg.llm.openai.classification_model,
            temperature=cfg.llm.openai.temperature,
            max_tokens=cfg.llm.openai.max_tokens
        )

    if not llm.is_available():
        logger.warning(f"LLM provider ({llm_provider}) not available, running in degraded mode")

    # Setup classifiers
    logger.info("Setting up classifiers")
    feature_extractor = FeatureExtractor(cfg.features.dict())
    ml_classifier = MLClassifier()
    llm_classifier = LLMClassifier(llm, categories, cfg.dict())
    adaptive_classifier = AdaptiveClassifier(
        feature_extractor,
        ml_classifier,
        llm_classifier,
        categories,
        cfg.dict(),
        disable_llm_fallback=no_llm_fallback
    )

    # Fetch emails
    logger.info(f"Fetching emails (limit: {limit})")
    emails = provider.fetch_emails(limit=limit)

    if not emails:
        logger.error("No emails fetched")
        sys.exit(1)

    logger.info(f"Fetched {len(emails)} emails")

    # Category verification (if requested and model exists)
    if verify_categories and not ml_classifier.is_mock and ml_classifier.model:
        logger.info("=" * 80)
        logger.info("VERIFYING MODEL CATEGORIES")
        logger.info("=" * 80)

        from src.calibration.category_verifier import verify_model_categories

        verification_result = verify_model_categories(
            emails=emails,
            model_categories=ml_classifier.categories,
            llm_provider=llm,
            sample_size=min(verify_sample, len(emails))
        )

        logger.info(f"Verification: {verification_result['verdict']}")
        logger.info(f"Confidence: {verification_result['confidence']:.0%}")

        if verification_result['verdict'] == 'POOR_MATCH':
            logger.warning("=" * 80)
            logger.warning("WARNING: Model categories may not fit this mailbox well")
            logger.warning(f"Suggested categories: {verification_result.get('suggested_categories', [])}")
            logger.warning("Consider running full calibration for better accuracy")
            logger.warning("Proceeding with existing model anyway...")
            logger.warning("=" * 80)
        elif verification_result['verdict'] == 'GOOD_MATCH':
            logger.info("Model categories look appropriate for this mailbox")

        logger.info("=" * 80)

    # Intelligent scaling: Decide if we need ML at all
    total_emails = len(emails)

    # Skip ML for small datasets (<1000 emails) - use LLM only
    if total_emails < 1000:
        logger.warning(f"Only {total_emails} emails - too few for ML training")
        logger.warning("Using LLM-only classification (no ML model)")
        ml_classifier.is_mock = True

    # Check if we need calibration (no good ML model)
    if ml_classifier.is_mock or not ml_classifier.model:
        if total_emails >= 1000:
            logger.info("=" * 80)
            logger.info("RUNNING CALIBRATION - Training ML model")
            logger.info("=" * 80)

            from src.calibration.workflow import CalibrationWorkflow, CalibrationConfig

            # Intelligent scaling for calibration and validation
            # Calibration: 3% of emails (min 250, max 1500)
            calibration_size = max(250, min(1500, int(total_emails * 0.03)))
            # Validation: 1% of emails (min 100, max 300)
            validation_size = max(100, min(300, int(total_emails * 0.01)))

            logger.info(f"Total emails: {total_emails:,}")
            logger.info(f"Calibration samples: {calibration_size} ({calibration_size/total_emails*100:.1f}%)")
            logger.info(f"Validation samples: {validation_size} ({validation_size/total_emails*100:.1f}%)")

            # Create calibration LLM provider
            calibration_llm = OllamaProvider(
                base_url=cfg.llm.ollama.base_url,
                model=cfg.llm.ollama.calibration_model,
                temperature=cfg.llm.ollama.temperature,
                max_tokens=cfg.llm.ollama.max_tokens
            )
            logger.info(f"Calibration model: {cfg.llm.ollama.calibration_model}")

            # Create consolidation LLM provider
            consolidation_model = getattr(cfg.llm.ollama, 'consolidation_model', cfg.llm.ollama.calibration_model)
            consolidation_llm = OllamaProvider(
                base_url=cfg.llm.ollama.base_url,
                model=consolidation_model,
                temperature=cfg.llm.ollama.temperature,
                max_tokens=cfg.llm.ollama.max_tokens
            )
            logger.info(f"Consolidation model: {consolidation_model}")

            calibration_config = CalibrationConfig(
                sample_size=calibration_size,
                validation_size=validation_size,
                llm_batch_size=50
            )

            calibration = CalibrationWorkflow(
                llm_provider=calibration_llm,
                consolidation_llm_provider=consolidation_llm,
                feature_extractor=feature_extractor,
                categories={},  # Don't pass hardcoded - let LLM discover
                config=calibration_config
            )

            # Run calibration to train ML model
            cal_results = calibration.run(emails, model_output_path="src/models/calibrated/classifier.pkl")

            # Reload the ML classifier with the new model
            ml_classifier = MLClassifier(model_path="src/models/calibrated/classifier.pkl")
            adaptive_classifier.ml_classifier = ml_classifier

            logger.info(f"Calibration complete! Accuracy: {cal_results.get('validation_accuracy', 0):.1%}")
            logger.info("=" * 80)

    # Classify emails
    logger.info("Starting classification")

    # Batch size for embedding extraction (larger = fewer API calls but more memory)
    batch_size = 512
    logger.info(f"Extracting features in batches (batch_size={batch_size})...")

    # Extract all features in batches (MUCH faster than one-at-a-time)
    all_features = feature_extractor.extract_batch(emails, batch_size=batch_size)

    logger.info(f"Feature extraction complete, classifying {len(emails)} emails...")
    results = []

    for i, (email, features) in enumerate(zip(emails, all_features)):
        if (i + 1) % 1000 == 0:
            logger.info(f"Progress: {i+1}/{len(emails)}")

        result = adaptive_classifier.classify_with_features(email, features)

        # If low confidence and LLM available: Use LLM
        if result.needs_review and llm.is_available():
            result = adaptive_classifier.classify_with_llm(result, email)

        results.append(result)

    # Export results
    logger.info("Exporting results")
    Path(output).mkdir(parents=True, exist_ok=True)

    import json
    results_data = {
        'metadata': {
            'total_emails': len(emails),
            'accuracy_estimate': adaptive_classifier.get_stats().accuracy_estimate(),
            'classification_stats': {
                'rule_matched': adaptive_classifier.get_stats().rule_matched,
                'ml_classified': adaptive_classifier.get_stats().ml_classified,
                'llm_classified': adaptive_classifier.get_stats().llm_classified,
                'needs_review': adaptive_classifier.get_stats().needs_review,
            }
        },
        'classifications': [
            {
                'email_id': r.email_id,
                'category': r.category,
                'confidence': r.confidence,
                'method': r.method
            }
            for r in results
        ]
    }

    results_file = Path(output) / 'results.json'
    with open(results_file, 'w') as f:
        json.dump(results_data, f, indent=2)

    logger.info(f"Results saved to: {results_file}")
    logger.info(f"\n{adaptive_classifier.get_stats()}")

    # Print summary
    logger.info("=" * 80)
    logger.info(f"Classification complete!")
    logger.info(f"Accuracy estimate: {adaptive_classifier.get_stats().accuracy_estimate():.1%}")
    logger.info("=" * 80)

    provider.disconnect()


@cli.command()
def test_config():
    """Test configuration loading."""
    logger = logging.getLogger(__name__)
    setup_logging()

    logger.info("Testing configuration...")

    cfg = load_config()
    logger.info(f"Config loaded: {cfg.version}")

    categories = load_categories()
    logger.info(f"Categories loaded: {len(categories)} categories")

    logger.info("Configuration test passed!")


@cli.command()
def test_ollama():
    """Test Ollama connection."""
    logger = logging.getLogger(__name__)
    setup_logging()

    logger.info("Testing Ollama connection...")

    from src.utils.config import load_config
    cfg = load_config()

    llm = OllamaProvider(
        base_url=cfg.llm.ollama.base_url,
        model=cfg.llm.ollama.classification_model
    )

    if llm.is_available():
        logger.info("Ollama connection successful!")
        logger.info(f"Model: {llm.model}")

        # Try a simple prompt
        logger.info("Testing inference...")
        response = llm.complete("Classify this email: 'Your verification code is 123456'")
        logger.info(f"Response: {response[:100]}...")
    else:
        logger.error("Ollama not available")
        sys.exit(1)


@cli.command()
def test_gmail():
    """Test Gmail connection."""
    logger = logging.getLogger(__name__)
    setup_logging()

    logger.info("Testing Gmail connection...")

    provider = GmailProvider()

    if provider.connect({'credentials_path': 'credentials.json'}):
        logger.info("Gmail connection successful!")
        emails = provider.fetch_emails(limit=5)
        logger.info(f"Fetched {len(emails)} test emails")
        provider.disconnect()
    else:
        logger.error("Gmail connection failed")
        sys.exit(1)


if __name__ == '__main__':
    cli()
