"""Email tool: read emails from IMAP mailbox."""

import asyncio
import imaplib
import ssl
from pathlib import Path
from datetime import date
from email import policy
from email.header import decode_header, make_header
from email.parser import BytesParser
from email.utils import parseaddr
from typing import Any

from nanobot.agent.tools.base import Tool


class EmailTool(Tool):
    """Read emails from configured IMAP mailbox."""
    
    name = "read_emails"
    description = (
        "USE THIS TOOL FOR OWNER-INITIATED MAILBOX QUERIES. When the user explicitly asks about their inbox "
        "(e.g. latest email, email sender, unread emails, search inbox, attachments, etc.), you SHOULD call "
        "read_emails(). DO NOT use mcp_gmail_mcp_read_email for emails received via the email channel - use "
        "read_emails instead. DO NOT use exec() with mail/tail/awk commands. DO NOT use read_file() on /var/mail "
        "or memory files. DO NOT try alternative methods. This is the ONLY way to read emails from IMAP - it connects "
        "to IMAP and fetches real-time data. For 'latest email' or 'last email received' queries, use limit=1. "
        "When the user asks to download attachments, use download_attachments=true. When the user asks to find "
        "emails with a specific attachment (e.g., 'find email with attachment Rubiks'), use attachment_name='Rubiks'. "
        "CRITICAL: When the user asks for specific fields like 'From and Subject' or 'sender and subject', extract "
        "and return ONLY those fields from the tool output. Do NOT summarize or analyze the email body content unless "
        "the user specifically asks for it. If the user asks 'give me the from and subject', respond with just: "
        "'From: [email] Subject: [subject]'. VERY IMPORTANT: When replying on the email channel to a single incoming "
        "email (content typically starts with 'Email received. From: ... Subject: ...'), treat that content as the "
        "user's message and compose a direct reply. Do NOT call read_emails in that case unless the human explicitly "
        "asks you to inspect or search the mailbox. Parameters: limit (1-50, default 10, use 1 for latest), "
        "unread_only (bool, default false), mark_seen (bool, default false), download_attachments (bool, default false "
        "- set to true to download all attachments to workspace), attachment_name (string, optional - filter emails by "
        "attachment filename, case-insensitive partial match). Returns formatted email list with sender, subject, date, "
        "attachments (if any), downloaded file paths (if downloaded), and body. After you receive this output, your "
        "reply to the user must address their email question using only this data—no unrelated topics."
    )
    
    def __init__(self, email_config: Any = None, workspace: Path | None = None):
        """
        Initialize email tool with email configuration.
        
        Args:
            email_config: Optional EmailConfig instance. If None, loads from config.
            workspace: Directory for downloaded attachments (defaults to config workspace_path).
        """
        self._email_config = email_config
        self._workspace = workspace
    
    @property
    def config(self) -> Any:
        """Lazy load email config if not provided."""
        if self._email_config is None:
            from nanobot.config.loader import load_config
            config = load_config()
            self._email_config = config.channels.email
        return self._email_config
    
    def coerce_params(self, params: dict[str, Any]) -> dict[str, Any]:
        """Coerce parameters, handling common name mismatches."""
        # Handle nested "parameters" key (some LLMs wrap params this way)
        if "parameters" in params and isinstance(params["parameters"], dict):
            # Extract nested parameters and merge with top-level
            nested = params.pop("parameters")
            params = {**params, **nested}
        
        # Remove common non-parameter keys that LLMs sometimes include
        params = params.copy()
        params.pop("function", None)
        params.pop("functionName", None)
        params.pop("function_name", None)
        params.pop("action", None)  # Some LLMs use "action" instead of function name
        
        coerced = super().coerce_params(params)
        # Map 'count' to 'limit' if limit not present
        if 'count' in coerced and 'limit' not in coerced:
            try:
                coerced['limit'] = int(coerced.pop('count'))
            except (ValueError, TypeError):
                pass
        # Remove unsupported parameters
        supported = {'limit', 'unread_only', 'mark_seen', 'download_attachments', 'attachment_name'}
        coerced = {k: v for k, v in coerced.items() if k in supported}
        return coerced
    
    @property
    def parameters(self) -> dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "limit": {
                    "type": "integer",
                    "description": "Number of emails to return. REQUIRED for 'latest email' queries - always use limit=1. For multiple emails use limit=5, limit=10, etc. (default: 100, max: 100)",
                    "minimum": 1,
                    "maximum": 100,
                },
                "unread_only": {
                    "type": "boolean",
                    "description": "If true, only return unread emails. If false, returns all emails including read ones (default: false)",
                },
                "mark_seen": {
                    "type": "boolean",
                    "description": "If true, mark emails as read after fetching. If false, leave read/unread status unchanged (default: false)",
                },
                "download_attachments": {
                    "type": "boolean",
                    "description": "If true, download all attachments from the emails to the workspace directory (default: false)",
                },
                "attachment_name": {
                    "type": "string",
                    "description": "Optional filter: only return emails that have at least one attachment whose filename contains this string (case-insensitive). Example: 'Rubiks' will match emails with attachments like 'Rubiks_SolutionGuide.pdf'. If not provided, all emails are returned.",
                },
            },
            "required": [],
        }
    
    async def execute(self, limit: int = 100, unread_only: bool = False, mark_seen: bool = False, download_attachments: bool = False, attachment_name: str | None = None, **kwargs: Any) -> str:
        """
        Read emails from IMAP mailbox.
        
        Args:
            limit: Maximum number of emails to return (use limit=1 for latest email)
            unread_only: If true, only fetch unread emails
            mark_seen: If true, mark emails as read after fetching
            **kwargs: Ignore any extra parameters (like count, sort_by, direction)
            
        Returns:
            Formatted string with email information
        """
        # Convert limit to int if it's a string (from JSON parsing)
        if isinstance(limit, str):
            try:
                limit = int(limit)
            except (ValueError, TypeError):
                limit = 10
        
        # Convert boolean parameters from strings if needed (from JSON parsing)
        if isinstance(download_attachments, str):
            download_attachments = download_attachments.lower() in ("true", "1", "yes", "on")
        if isinstance(unread_only, str):
            unread_only = unread_only.lower() in ("true", "1", "yes", "on")
        if isinstance(mark_seen, str):
            mark_seen = mark_seen.lower() in ("true", "1", "yes", "on")
        
        # Normalize attachment_name (empty string or None means no filter)
        if attachment_name is not None:
            attachment_name = str(attachment_name).strip()
            if not attachment_name:
                attachment_name = None
        
        # Handle common parameter name mismatches (agent sometimes uses 'count' instead of 'limit')
        # Also handle if count is passed as a positional argument via kwargs
        if 'count' in kwargs:
            try:
                limit = int(kwargs['count'])
            except (ValueError, TypeError):
                pass
        # Also check if limit was passed in kwargs (in case it wasn't a named parameter)
        if 'limit' in kwargs:
            try:
                limit = int(kwargs['limit'])
            except (ValueError, TypeError):
                pass
        # Ignore unsupported parameters like sort_by, direction, reverse, etc.
        try:
            config = self.config
        except Exception as e:
            return f"Error loading email configuration: {str(e)}"
        
        if not config:
            return "Error: Email configuration not found"
        
        if not hasattr(config, 'enabled') or not config.enabled:
            return "Error: Email channel is not enabled in configuration. Set NANOBOT_CHANNELS__EMAIL__ENABLED=true"
        
        if not hasattr(config, 'consent_granted') or not config.consent_granted:
            return "Error: Email access consent not granted. Set NANOBOT_CHANNELS__EMAIL__CONSENT_GRANTED=true"
        
        if not hasattr(config, 'imap_host') or not config.imap_host:
            return "Error: IMAP host not configured. Set NANOBOT_CHANNELS__EMAIL__IMAP_HOST"
        
        if not hasattr(config, 'imap_username') or not config.imap_username:
            return "Error: IMAP username not configured. Set NANOBOT_CHANNELS__EMAIL__IMAP_USERNAME"
        
        if not hasattr(config, 'imap_password') or not config.imap_password:
            return "Error: IMAP password not configured. Set NANOBOT_CHANNELS__EMAIL__IMAP_PASSWORD"
        
        # Limit to reasonable maximum
        try:
            limit = min(max(1, int(limit)), 100)
        except (ValueError, TypeError):
            limit = 100
        
        try:
            messages = await asyncio.to_thread(
                self._fetch_messages,
                unread_only=unread_only,
                mark_seen=mark_seen,
                limit=limit,
                download_attachments=download_attachments,
                attachment_name=attachment_name,
            )
            
            if not messages:
                if attachment_name:
                    return f"No emails found with attachments matching '{attachment_name}'. Try increasing the limit or checking if the attachment name is correct."
                elif unread_only:
                    return "No unread emails found in your inbox."
                else:
                    return f"No emails found in your inbox. The mailbox appears to be empty or there was an issue retrieving emails."
            
            result_parts = [f"Found {len(messages)} email(s):\n"]
            for i, msg in enumerate(messages, 1):
                result_parts.append(f"\n--- Email {i} ---")
                result_parts.append(f"From: {msg['sender']}")
                result_parts.append(f"Subject: {msg['subject']}")
                result_parts.append(f"Date: {msg['metadata']['date']}")
                if msg['metadata'].get('attachments'):
                    att_list = ', '.join([a['filename'] for a in msg['metadata']['attachments']])
                    result_parts.append(f"Attachments: {att_list}")
                if msg['metadata'].get('downloaded_files'):
                    dl_list = ', '.join(msg['metadata']['downloaded_files'])
                    result_parts.append(f"Downloaded to: {dl_list}")
                # Only include body content if specifically requested, otherwise keep it brief
                result_parts.append(f"\nBody: {msg['content'][:500]}..." if len(msg['content']) > 500 else f"\nBody: {msg['content']}")
            
            return "\n".join(result_parts)
            
        except Exception as e:
            import traceback
            error_details = traceback.format_exc()
            return f"Error reading emails: {str(e)}\n\nDetails: {error_details}"
    
    def _fetch_messages(
        self,
        unread_only: bool,
        mark_seen: bool,
        limit: int,
        download_attachments: bool = False,
        attachment_name: str | None = None,
    ) -> list[dict[str, Any]]:
        """Fetch messages from IMAP mailbox."""
        messages: list[dict[str, Any]] = []
        mailbox = self.config.imap_mailbox or "INBOX"
        
        # Build search criteria
        if unread_only:
            search_criteria = ("UNSEEN",)
        else:
            search_criteria = ("ALL",)
        
        # Connect to IMAP server
        try:
            if self.config.imap_use_ssl:
                client = imaplib.IMAP4_SSL(self.config.imap_host, self.config.imap_port)
            else:
                client = imaplib.IMAP4(self.config.imap_host, self.config.imap_port)
        except Exception as e:
            raise Exception(f"Failed to connect to IMAP server {self.config.imap_host}:{self.config.imap_port}: {str(e)}")
        
        try:
            client.login(self.config.imap_username, self.config.imap_password.strip())
        except imaplib.IMAP4.error as e:
            error_msg = str(e)
            if "AUTHENTICATE" in error_msg.upper() or "LOGIN" in error_msg.upper():
                raise Exception(
                    f"IMAP authentication failed. Please check:\n"
                    f"1. Your email username: {self.config.imap_username}\n"
                    f"2. Your password/app password is correct\n"
                    f"3. For Gmail: Enable 2-Step Verification and create an App Password at https://myaccount.google.com/apppasswords\n"
                    f"4. IMAP is enabled in your email account settings\n"
                    f"Original error: {error_msg}"
                )
            raise
        
        try:
            status, _ = client.select(mailbox)
            if status != "OK":
                return messages
            
            status, data = client.search(None, *search_criteria)
            if status != "OK" or not data:
                return messages
            
            ids = data[0].split()
            if limit > 0 and len(ids) > limit:
                # Get most recent emails (last N)
                ids = ids[-limit:]
            
            for imap_id in ids:
                status, fetched = client.fetch(imap_id, "(BODY.PEEK[] UID)")
                if status != "OK" or not fetched:
                    continue
                
                raw_bytes = self._extract_message_bytes(fetched)
                if raw_bytes is None:
                    continue
                
                parsed = BytesParser(policy=policy.default).parsebytes(raw_bytes)
                sender = parseaddr(parsed.get("From", ""))[1].strip().lower()
                if not sender:
                    # Try to get display name if email not found
                    from_addr = parsed.get("From", "")
                    sender = from_addr if from_addr else "unknown"
                
                subject = self._decode_header_value(parsed.get("Subject", ""))
                date_value = parsed.get("Date", "")
                message_id = parsed.get("Message-ID", "").strip()
                body = self._extract_text_body(parsed)
                attachments = self._extract_attachments(parsed)
                
                # Download attachments if requested
                downloaded_files = []
                if download_attachments and attachments:
                    import logging
                    logger = logging.getLogger(__name__)
                    if self._workspace is not None:
                        workspace = Path(self._workspace).expanduser().resolve()
                    else:
                        from nanobot.config.loader import load_config

                        workspace = load_config().workspace_path.expanduser().resolve()
                    workspace.mkdir(parents=True, exist_ok=True)
                    
                    # Build a map of attachment parts by decoded filename for efficient lookup
                    attachment_parts = {}
                    for part in parsed.walk():
                        if part.get_content_disposition() == "attachment":
                            part_filename = part.get_filename()
                            if part_filename:
                                try:
                                    from email.header import decode_header, make_header
                                    decoded_part_filename = str(make_header(decode_header(part_filename)))
                                except Exception:
                                    decoded_part_filename = part_filename
                                attachment_parts[decoded_part_filename] = part
                                logger.debug(f"Found attachment part: '{decoded_part_filename}' (original: '{part_filename}')")
                    
                    logger.debug(f"Total attachment parts found: {len(attachment_parts)}, requested attachments: {len(attachments)}")
                    if attachments:
                        logger.debug(f"Requested attachment filenames: {[a['filename'] for a in attachments]}")
                    
                    # Download each attachment
                    for att_info in attachments:
                        filename = att_info['filename']
                        matched_filename = filename  # Will be updated if we match by base name
                        try:
                            # Try exact match first
                            part = attachment_parts.get(filename)
                            
                            # If no exact match, try case-insensitive and normalized matching
                            if part is None:
                                filename_lower = filename.lower().strip()
                                for decoded_name, part_candidate in attachment_parts.items():
                                    decoded_lower = decoded_name.lower().strip()
                                    if decoded_lower == filename_lower:
                                        part = part_candidate
                                        matched_filename = decoded_name
                                        logger.debug(f"Matched attachment '{filename}' using case-insensitive match with '{decoded_name}'")
                                        break
                            
                            # If still no match, try matching by base filename (strip common prefixes like Gmail attachment IDs)
                            if part is None:
                                # Extract base filename (everything after last underscore or use full name)
                                # Gmail sometimes adds prefixes like "65afea09c4f7a02afbb9d876_filename.pdf"
                                base_filename = filename
                                if '_' in filename:
                                    # Try to match the part after the last underscore if it looks like a hash prefix
                                    parts = filename.rsplit('_', 1)
                                    if len(parts) == 2 and len(parts[0]) >= 20 and parts[0].isalnum():
                                        # Looks like a hash prefix, use the base name
                                        base_filename = parts[1]
                                
                                base_lower = base_filename.lower().strip()
                                for decoded_name, part_candidate in attachment_parts.items():
                                    decoded_lower = decoded_name.lower().strip()
                                    # Check if decoded name matches base filename
                                    if decoded_lower == base_lower:
                                        part = part_candidate
                                        matched_filename = decoded_name
                                        logger.debug(f"Matched attachment '{filename}' by base filename '{base_filename}' with '{decoded_name}'")
                                        break
                                    # Also check if decoded name ends with base filename (in case it has its own prefix)
                                    if decoded_lower.endswith(base_lower) or base_lower.endswith(decoded_lower):
                                        part = part_candidate
                                        matched_filename = decoded_name
                                        logger.debug(f"Matched attachment '{filename}' by partial match: base '{base_filename}' with '{decoded_name}'")
                                        break
                            
                            if part is None:
                                logger.warning(f"Could not find attachment part for filename: {filename}")
                                logger.debug(f"Available attachment filenames: {list(attachment_parts.keys())}")
                                continue
                            
                            # Save the attachment using the matched filename (cleaner, without prefixes)
                            att_data = part.get_payload(decode=True)
                            if att_data:
                                # Sanitize filename but preserve extension
                                safe_filename = "".join(c for c in matched_filename if c.isalnum() or c in "._- ")
                                safe_filename = safe_filename.replace(" ", "_")
                                file_path = workspace / safe_filename
                                file_path.write_bytes(att_data)
                                downloaded_files.append(str(file_path))
                                logger.info(f"Downloaded attachment '{filename}' (matched as '{matched_filename}') to {file_path}")
                            else:
                                logger.warning(f"Attachment '{filename}' has no data to save")
                        except Exception as e:
                            logger.error(f"Error downloading attachment '{filename}': {str(e)}", exc_info=True)
                
                if not body:
                    body = "(empty email body)"
                
                # Limit body length
                max_chars = getattr(self.config, 'max_body_chars', 12000)
                body = body[:max_chars]
                
                content = (
                    f"Email received.\n"
                    f"From: {sender}\n"
                    f"Subject: {subject}\n"
                    f"Date: {date_value}\n"
                )
                if attachments:
                    content += f"Attachments: {', '.join([a['filename'] for a in attachments])}\n"
                if downloaded_files:
                    content += f"Downloaded attachments to: {', '.join(downloaded_files)}\n"
                content += f"\n{body}"
                
                # Filter by attachment name if specified
                if attachment_name:
                    attachment_name_lower = attachment_name.lower().strip()
                    # Check if any attachment filename contains the search term (case-insensitive)
                    has_matching_attachment = False
                    if attachments:
                        for att in attachments:
                            att_filename_lower = att['filename'].lower()
                            if attachment_name_lower in att_filename_lower:
                                has_matching_attachment = True
                                break
                    
                    # Skip this email if it doesn't have a matching attachment
                    if not has_matching_attachment:
                        continue
                
                metadata = {
                    "message_id": message_id,
                    "subject": subject,
                    "date": date_value,
                    "sender_email": sender,
                    "attachments": attachments,
                    "downloaded_files": downloaded_files,
                }
                
                messages.append({
                    "sender": sender,
                    "subject": subject,
                    "message_id": message_id,
                    "content": content,
                    "metadata": metadata,
                })
                
                if mark_seen:
                    client.store(imap_id, "+FLAGS", "\\Seen")
        finally:
            try:
                client.logout()
            except Exception:
                pass
        
        return messages
    
    @staticmethod
    def _extract_message_bytes(fetched: list[Any]) -> bytes | None:
        """Extract raw message bytes from IMAP fetch response."""
        for item in fetched:
            if isinstance(item, tuple) and len(item) >= 2 and isinstance(item[1], (bytes, bytearray)):
                return bytes(item[1])
        return None
    
    @staticmethod
    def _decode_header_value(value: str) -> str:
        """Decode email header value (handles encoded words)."""
        if not value:
            return ""
        try:
            return str(make_header(decode_header(value)))
        except Exception:
            return value
    
    @staticmethod
    def _extract_attachments(msg: Any) -> list[dict[str, str]]:
        """Extract attachment information from email message."""
        attachments = []
        if msg.is_multipart():
            for part in msg.walk():
                disposition = part.get_content_disposition()
                if disposition == "attachment":
                    filename = part.get_filename()
                    if filename:
                        # Decode filename if needed
                        try:
                            from email.header import decode_header, make_header
                            decoded_filename = str(make_header(decode_header(filename)))
                        except Exception:
                            decoded_filename = filename
                        attachments.append({
                            "filename": decoded_filename,
                            "content_type": part.get_content_type(),
                            "size": len(part.get_payload(decode=True) or b"") if part.get_payload(decode=True) else 0,
                        })
        return attachments
    
    @staticmethod
    def _extract_text_body(msg: Any) -> str:
        """Extract readable text body from email message."""
        if msg.is_multipart():
            plain_parts: list[str] = []
            html_parts: list[str] = []
            for part in msg.walk():
                if part.get_content_disposition() == "attachment":
                    continue
                content_type = part.get_content_type()
                try:
                    payload = part.get_content()
                except Exception:
                    payload_bytes = part.get_payload(decode=True) or b""
                    charset = part.get_content_charset() or "utf-8"
                    payload = payload_bytes.decode(charset, errors="replace")
                if not isinstance(payload, str):
                    continue
                if content_type == "text/plain":
                    plain_parts.append(payload)
                elif content_type == "text/html":
                    html_parts.append(payload)
            if plain_parts:
                return "\n\n".join(plain_parts).strip()
            if html_parts:
                # Simple HTML to text conversion
                import re
                import html
                text = re.sub(r"<\s*br\s*/?>", "\n", "\n\n".join(html_parts), flags=re.IGNORECASE)
                text = re.sub(r"<\s*/\s*p\s*>", "\n", text, flags=re.IGNORECASE)
                text = re.sub(r"<[^>]+>", "", text)
                return html.unescape(text).strip()
            return ""
        
        try:
            payload = msg.get_content()
        except Exception:
            payload_bytes = msg.get_payload(decode=True) or b""
            charset = msg.get_content_charset() or "utf-8"
            payload = payload_bytes.decode(charset, errors="replace")
        if not isinstance(payload, str):
            return ""
        if msg.get_content_type() == "text/html":
            import re
            import html
            text = re.sub(r"<\s*br\s*/?>", "\n", payload, flags=re.IGNORECASE)
            text = re.sub(r"<\s*/\s*p\s*>", "\n", text, flags=re.IGNORECASE)
            text = re.sub(r"<[^>]+>", "", text)
            return html.unescape(text).strip()
        return payload.strip()

