"""Feature extraction from emails for classification."""
import re
import logging
from typing import Dict, List, Any, Optional, Tuple
from datetime import datetime
import numpy as np

try:
    import pandas as pd
except ImportError:
    pd = None

try:
    from sklearn.feature_extraction.text import TfidfVectorizer
except ImportError:
    TfidfVectorizer = None

try:
    from sentence_transformers import SentenceTransformer
except ImportError:
    SentenceTransformer = None

from src.email_providers.base import Email

logger = logging.getLogger(__name__)


class FeatureExtractor:
    """
    Extract features from emails for classification.

    Combines three feature types:
    1. Semantic embeddings (384 dimensions)
    2. Hard pattern detection (20+ boolean features)
    3. Structural metadata (20+ categorical/numerical)
    """

    def __init__(self, config: Optional[Dict] = None):
        """Initialize feature extractor."""
        self.config = config or self._default_config()
        self.embedder = None
        self.text_vectorizer = None
        self._initialize_embedder()
        self._initialize_vectorizer()

    def _default_config(self) -> Dict:
        """Default configuration."""
        return {
            'embedding_model': 'all-MiniLM-L6-v2',
            'embedding_batch_size': 32,
            'text_features': {
                'max_features': 10000,
                'ngram_range': [1, 2],
                'min_df': 2,
                'max_df': 0.95,
            }
        }

    def _initialize_embedder(self) -> None:
        """
        Initialize embedding model via Ollama.

        NOTE: We use Ollama's all-minilm:l6-v2 model instead of downloading sentence-transformers.
        This is MUCH faster (2-3 seconds vs 90 seconds) since Ollama caches the model.

        TODO: The original design used sentence-transformers which downloads the model each time.
        We bypassed it to use Ollama for speed. If sentence-transformers had proper caching,
        it would also be 2-3 seconds. Keep this Ollama approach for now.
        """
        try:
            import ollama
            self.embedder = ollama.Client(host="http://localhost:11434")
            logger.info("Embedder initialized: using Ollama (all-minilm:l6-v2)")
            logger.info("Embedding dimension: 384 dims")
        except ImportError:
            logger.error("ollama package not installed: pip install ollama")
            self.embedder = None
        except Exception as e:
            logger.error(f"Failed to initialize Ollama embedder: {e}")
            self.embedder = None

    def _initialize_vectorizer(self) -> None:
        """Initialize TF-IDF vectorizer."""
        if TfidfVectorizer is None:
            logger.warning("scikit-learn not installed, text vectorization unavailable")
            self.text_vectorizer = None
            return

        try:
            text_config = self.config.get('text_features', {})
            self.text_vectorizer = TfidfVectorizer(
                max_features=text_config.get('max_features', 10000),
                ngram_range=tuple(text_config.get('ngram_range', [1, 2])),
                min_df=text_config.get('min_df', 2),
                max_df=text_config.get('max_df', 0.95),
                sublinear_tf=True
            )
            logger.info("TF-IDF vectorizer initialized")
        except Exception as e:
            logger.error(f"Failed to initialize vectorizer: {e}")
            self.text_vectorizer = None

    def extract(self, email: Email) -> Dict[str, Any]:
        """Extract all features from a single email."""
        features = {}

        # Basic email content
        features['subject'] = email.subject
        features['body_snippet'] = email.body_snippet
        features['full_body'] = email.body

        # Structural features
        features.update(self._extract_structural(email))

        # Sender features
        features.update(self._extract_sender(email))

        # Pattern features
        features.update(self._extract_patterns(email))

        # Embedding (if available)
        if self.embedder:
            features['embedding'] = self._extract_embedding(email)
        else:
            logger.debug("Embedding model not available, using zeros")
            features['embedding'] = np.zeros(384)  # Default MiniLM dimension

        return features

    def _extract_structural(self, email: Email) -> Dict[str, Any]:
        """Extract structural/metadata features."""
        features = {}

        # Attachments
        features['has_attachments'] = email.has_attachments
        features['attachment_count'] = email.attachment_count
        features['attachment_types'] = email.attachment_types

        # Links and images
        body = email.body or email.body_snippet or ""
        subject = email.subject or ""
        combined = f"{subject} {body}"

        features['link_count'] = len(re.findall(r'https?://', combined))
        features['image_count'] = len(re.findall(r'<img', combined, re.IGNORECASE))

        # Lengths
        features['body_length'] = len(body)
        features['subject_length'] = len(subject)

        # Reply/Forward
        features['has_reply_prefix'] = bool(re.match(r'^(Re:|Fwd:)', subject, re.IGNORECASE))

        # Time features
        if email.date:
            hour = email.date.hour
            if 0 <= hour < 6:
                features['time_of_day'] = 'night'
            elif 6 <= hour < 12:
                features['time_of_day'] = 'morning'
            elif 12 <= hour < 18:
                features['time_of_day'] = 'afternoon'
            else:
                features['time_of_day'] = 'evening'
            features['day_of_week'] = email.date.strftime('%A').lower()
        else:
            features['time_of_day'] = 'unknown'
            features['day_of_week'] = 'unknown'

        return features

    def _extract_sender(self, email: Email) -> Dict[str, Any]:
        """Extract sender-based features."""
        features = {}

        sender = email.sender or ""
        if '@' in sender:
            domain = sender.split('@')[1].lower()
            features['sender_domain'] = domain

            # Classify domain type
            freemail_domains = {'gmail.com', 'yahoo.com', 'hotmail.com', 'outlook.com', 'icloud.com', 'protonmail.com'}
            noreply_patterns = ['noreply', 'no-reply', 'donotreply', 'no_reply']

            if domain in freemail_domains:
                features['sender_domain_type'] = 'freemail'
            elif any(p in sender.lower() for p in noreply_patterns):
                features['sender_domain_type'] = 'noreply'
            else:
                features['sender_domain_type'] = 'corporate'

            features['is_noreply'] = any(p in sender.lower() for p in noreply_patterns)
        else:
            features['sender_domain'] = 'unknown'
            features['sender_domain_type'] = 'unknown'
            features['is_noreply'] = False

        return features

    def _extract_patterns(self, email: Email) -> Dict[str, Any]:
        """Extract hard pattern-based features."""
        features = {}

        body = (email.body or email.body_snippet or "").lower()
        subject = (email.subject or "").lower()
        combined = f"{subject} {body}"

        # Authentication patterns
        features['has_otp_pattern'] = bool(re.search(r'\b\d{4,6}\b', combined))
        features['has_verification'] = 'verification' in combined
        features['has_reset_password'] = 'reset password' in combined

        # Transactional patterns
        features['has_invoice_pattern'] = bool(re.search(r'(invoice|bill|receipt)\s*#?\d+', combined, re.I))
        features['has_price'] = bool(re.search(r'\$[\d,]+\.?\d*', combined))
        features['has_order_number'] = bool(re.search(r'order\s*#?\d+', combined, re.I))
        features['has_tracking'] = bool(re.search(r'tracking\s*(number|#)', combined, re.I))

        # Marketing patterns
        features['has_unsubscribe'] = 'unsubscribe' in combined
        features['has_view_in_browser'] = 'view in browser' in combined
        features['has_promotional'] = any(p in combined for p in ['limited time', 'special offer', 'sale', 'discount'])

        # Meeting patterns
        features['has_meeting'] = bool(re.search(r'(meeting|call|zoom|teams|conference)', combined, re.I))
        features['has_calendar'] = 'calendar' in combined

        # Signature patterns
        features['has_signature'] = bool(re.search(r'(regards|sincerely|best|cheers|thanks)', combined, re.I))

        return features

    def extract_batch(self, emails: List[Email], batch_size: int = 512) -> List[Dict[str, Any]]:
        """
        Extract features from multiple emails with batched embeddings.

        Much faster than calling extract() in a loop because embeddings are batched.
        """
        if not emails:
            return []

        # Extract all non-embedding features first
        all_features = []
        texts_to_embed = []

        for email in emails:
            features = {}
            features['subject'] = email.subject
            features['body_snippet'] = email.body_snippet
            features['full_body'] = email.body
            features.update(self._extract_structural(email))
            features.update(self._extract_sender(email))
            features.update(self._extract_patterns(email))
            all_features.append(features)
            texts_to_embed.append(self._build_embedding_text(email))

        # Batch embed all texts
        if self.embedder:
            try:
                # Process in batches
                embeddings = []
                for i in range(0, len(texts_to_embed), batch_size):
                    batch = texts_to_embed[i:i + batch_size]
                    response = self.embedder.embed(
                        model='all-minilm:l6-v2',
                        input=batch
                    )
                    embeddings.extend(response['embeddings'])

                # Add embeddings to features
                for features, embedding in zip(all_features, embeddings):
                    features['embedding'] = np.array(embedding, dtype=np.float32)

            except Exception as e:
                logger.error(f"Batch embedding failed: {e}, falling back to zeros")
                for features in all_features:
                    features['embedding'] = np.zeros(384)
        else:
            for features in all_features:
                features['embedding'] = np.zeros(384)

        return all_features

    def _extract_embedding(self, email: Email) -> np.ndarray:
        """
        Generate semantic embedding for email using Ollama.

        Uses all-minilm:l6-v2 via Ollama (384 dimensions).
        Falls back to zero vector if Ollama unavailable.
        """
        if not self.embedder:
            return np.zeros(384)

        try:
            # Build structured text for embedding
            text = self._build_embedding_text(email)

            # Get embedding from Ollama (use new embed API)
            response = self.embedder.embed(
                model='all-minilm:l6-v2',
                input=text
            )
            embedding = np.array(response['embeddings'][0], dtype=np.float32)
            return embedding
        except Exception as e:
            logger.error(f"Error generating embedding: {e}")
            return np.zeros(384)

    def _build_embedding_text(self, email: Email) -> str:
        """Build structured text for embedding."""
        # Collect basic patterns
        patterns = self._extract_patterns(email)
        structural = self._extract_structural(email)

        # Build structured text with headers
        text = f"""[EMAIL_METADATA]
sender_type: {structural.get('sender_domain_type', 'unknown')}
time_of_day: {structural.get('time_of_day', 'unknown')}
has_attachments: {structural.get('has_attachments', False)}
attachment_count: {structural.get('attachment_count', 0)}

[DETECTED_PATTERNS]
has_otp: {patterns.get('has_otp_pattern', False)}
has_invoice: {patterns.get('has_invoice_pattern', False)}
has_unsubscribe: {patterns.get('has_unsubscribe', False)}
is_noreply: {structural.get('is_noreply', False)}
has_meeting: {patterns.get('has_meeting', False)}

[CONTENT]
subject: {email.subject[:100]}
body: {email.body_snippet[:300]}
"""
        return text

    def fit_text_vectorizer(self, emails: List[Email]) -> bool:
        """Fit TF-IDF vectorizer on email corpus."""
        if not self.text_vectorizer:
            logger.error("Text vectorizer not available")
            return False

        try:
            texts = [f"{e.subject} {e.body_snippet}" for e in emails]
            self.text_vectorizer.fit(texts)
            logger.info(f"Fitted TF-IDF vectorizer with {len(self.text_vectorizer.vocabulary_)} features")
            return True
        except Exception as e:
            logger.error(f"Error fitting vectorizer: {e}")
            return False

    def get_feature_names(self) -> List[str]:
        """Get list of all feature names."""
        names = [
            'has_attachments', 'attachment_count', 'link_count', 'image_count',
            'body_length', 'subject_length', 'has_reply_prefix',
            'time_of_day', 'day_of_week', 'sender_domain_type', 'is_noreply',
            'has_otp_pattern', 'has_verification', 'has_reset_password',
            'has_invoice_pattern', 'has_price', 'has_order_number', 'has_tracking',
            'has_unsubscribe', 'has_view_in_browser', 'has_promotional',
            'has_meeting', 'has_calendar', 'has_signature'
        ]
        return names
