"""Parse Enron dataset for training."""
import logging
import os
import email.message
import email.parser
from pathlib import Path
from typing import List, Optional
from datetime import datetime
from email.utils import parsedate_to_datetime

from src.email_providers.base import Email

logger = logging.getLogger(__name__)


class EnronParser:
    """
    Parse Enron email dataset.

    The Enron dataset is in maildir format:
    - enron_mail_20150507/
      - maildir/
        - lastname-f/
          - folder/
            - 1, 2, 3, ... (individual emails)
    """

    def __init__(self, dataset_path: str):
        """Initialize Enron parser."""
        self.dataset_path = Path(dataset_path)

        if not self.dataset_path.exists():
            raise ValueError(f"Dataset path not found: {self.dataset_path}")

        self.maildir_path = self.dataset_path / "maildir"

        if not self.maildir_path.exists():
            raise ValueError(f"Maildir not found at {self.maildir_path}")

        logger.info(f"Initialized Enron parser: {self.dataset_path}")

    def parse_emails(self, limit: Optional[int] = None) -> List[Email]:
        """
        Parse emails from Enron dataset.

        Args:
            limit: Maximum number of emails to parse

        Returns:
            List of Email objects
        """
        emails = []
        email_count = 0

        logger.info(f"Starting Enron parsing (limit: {limit})")

        # Iterate through users
        for user_dir in sorted(self.maildir_path.iterdir()):
            if not user_dir.is_dir():
                continue

            # Iterate through folders
            for folder_dir in sorted(user_dir.iterdir()):
                if not folder_dir.is_dir():
                    continue

                # Parse emails in folder
                for email_file in sorted(folder_dir.iterdir()):
                    if email_file.is_file():
                        try:
                            parsed_email = self._parse_email_file(email_file)
                            if parsed_email:
                                emails.append(parsed_email)
                                email_count += 1

                                if limit and email_count >= limit:
                                    logger.info(f"Reached limit: {email_count} emails parsed")
                                    return emails

                                if email_count % 1000 == 0:
                                    logger.info(f"Progress: {email_count} emails parsed")

                        except Exception as e:
                            logger.debug(f"Error parsing {email_file}: {e}")

        logger.info(f"Parsing complete: {email_count} emails")
        return emails

    def _parse_email_file(self, filepath: Path) -> Optional[Email]:
        """Parse single email file."""
        try:
            with open(filepath, 'rb') as f:
                msg = email.message_from_bytes(f.read())

            # Extract folder name from filepath
            # filepath structure: maildir/user-name/folder-name/123
            folder_name = filepath.parent.name

            # Extract basic info
            msg_id = str(filepath).replace('/', '_').replace('\\', '_')
            subject = msg.get('subject', 'No Subject')
            sender = msg.get('from', '')
            date_str = msg.get('date')

            # Parse date
            date = None
            if date_str:
                try:
                    date = parsedate_to_datetime(date_str)
                except Exception:
                    pass

            # Extract body
            body = self._extract_body(msg)
            body_snippet = body[:500] if body else ""

            return Email(
                id=msg_id,
                subject=subject,
                sender=sender,
                date=date,
                body=body,
                body_snippet=body_snippet,
                has_attachments=self._has_attachments(msg),
                provider='enron',
                headers={'X-Folder': folder_name}
            )

        except Exception as e:
            logger.debug(f"Error parsing email file: {e}")
            return None

    def _extract_body(self, msg: email.message.Message) -> str:
        """Extract email body."""
        body = ""

        if msg.is_multipart():
            for part in msg.walk():
                if part.get_content_type() == 'text/plain':
                    try:
                        payload = part.get_payload(decode=True)
                        if payload:
                            body = payload.decode('utf-8', errors='ignore')
                            break
                    except Exception:
                        pass
        else:
            try:
                payload = msg.get_payload(decode=True)
                if payload:
                    body = payload.decode('utf-8', errors='ignore')
                else:
                    body = msg.get_payload(decode=False)
                    if isinstance(body, str):
                        pass
                    else:
                        body = str(body)
            except Exception:
                pass

        return body.strip() if isinstance(body, str) else ""

    def _has_attachments(self, msg: email.message.Message) -> bool:
        """Check if email has attachments."""
        if msg.is_multipart():
            for part in msg.iter_attachments():
                if part.get_filename():
                    return True
        return False
