"""Email sampling for calibration."""
import logging
import random
from typing import List, Tuple
from src.email_providers.base import Email

logger = logging.getLogger(__name__)


class EmailSampler:
    """Sample emails for calibration phase."""

    def __init__(self, random_seed: int = 42):
        """Initialize sampler."""
        self.random_seed = random_seed
        random.seed(random_seed)

    def stratified_sample(
        self,
        emails: List[Email],
        sample_size: int,
        field_getter = None
    ) -> Tuple[List[Email], List[Email]]:
        """
        Stratified sampling based on sender domain type.

        Args:
            emails: All emails
            sample_size: Target sample size
            field_getter: Function to group by (default: sender_domain_type)

        Returns:
            (sample, remaining)
        """
        if sample_size >= len(emails):
            logger.warning(f"Sample size ({sample_size}) >= total ({len(emails)}), returning all")
            return emails, []

        if field_getter is None:
            # Default: group by sender domain type
            field_getter = lambda e: self._extract_domain_type(e)

        # Group emails by field
        groups = {}
        for email in emails:
            try:
                key = field_getter(email)
                if key not in groups:
                    groups[key] = []
                groups[key].append(email)
            except Exception:
                if 'unknown' not in groups:
                    groups['unknown'] = []
                groups['unknown'].append(email)

        logger.info(f"Stratified sampling: {len(groups)} groups, total {len(emails)} emails")

        # Proportional sampling from each group
        sample = []
        remaining = []

        for group_name, group_emails in groups.items():
            proportion = len(group_emails) / len(emails)
            group_sample_size = max(1, int(sample_size * proportion))

            group_sample = random.sample(group_emails, min(group_sample_size, len(group_emails)))
            group_remaining = [e for e in group_emails if e not in group_sample]

            sample.extend(group_sample)
            remaining.extend(group_remaining)

            logger.debug(f"Group '{group_name}': {len(group_sample)} sampled, {len(group_remaining)} remaining")

        # Adjust if we didn't get exactly the right size
        if len(sample) < sample_size and remaining:
            shortage = sample_size - len(sample)
            additional = random.sample(remaining, min(shortage, len(remaining)))
            sample.extend(additional)
            remaining = [e for e in remaining if e not in additional]

        logger.info(f"Stratified sample complete: {len(sample)} sampled")

        return sample[:sample_size], remaining

    def random_sample(
        self,
        emails: List[Email],
        sample_size: int
    ) -> Tuple[List[Email], List[Email]]:
        """Random sampling without stratification."""
        if sample_size >= len(emails):
            return emails, []

        sample = random.sample(emails, sample_size)
        remaining = [e for e in emails if e not in sample]

        logger.info(f"Random sample: {len(sample)} sampled")
        return sample, remaining

    def _extract_domain_type(self, email: Email) -> str:
        """Extract domain type for stratification."""
        sender = email.sender or ""
        if '@' in sender:
            domain = sender.split('@')[1].lower()
            freemail_domains = {'gmail.com', 'yahoo.com', 'hotmail.com', 'outlook.com'}

            if domain in freemail_domains:
                return 'freemail'
            elif 'noreply' in sender.lower():
                return 'noreply'
            else:
                return 'corporate'

        return 'unknown'
