import os, smtplib, imaplib, email
from fastmcp import FastMCP
from typing import Dict, Any, List
from email.mime.text import MIMEText
from email.utils import formataddr
from email.mime.multipart import MIMEMultipart
from email.mime.application import MIMEApplication

mcp = FastMCP("Email Client", version = "0.1.0")
class EmailClient:
    def __init__(self):
        self.username = os.environ.get("MAIL_USER_NAME", "")
        self.password = os.environ.get("MAIL_PASS_WORD", "")
        self.smtp_addr = os.environ.get("MAIL_SMTP_ADDR", "")
        self.imap_addr = os.environ.get("MAIL_IMAP_ADDR", "")
        self.save_path = os.environ.get("MAIL_SAVE_PATH", "")
        self.from_addr = os.environ.get("MAIL_FROM_ADDR", self.username)
        self.from_name = os.environ.get("MAIL_FROM_NAME", self.username)

    def check_env_vars(self, required_envs):
        for env in required_envs:
            if env not in os.environ or not os.environ[env]:
                raise RuntimeError(f"环境变量 {env} 未设置")

    def get_smtp_host_port(self):
        if ":" in self.smtp_addr:
            host, port = self.smtp_addr.split(":", 1)
            return host, int(port)
        return self.smtp_addr, 465

    def get_imap_host_port(self):
        if ":" in self.imap_addr:
            host, port = self.imap_addr.split(":", 1)
            return host, int(port)
        return self.imap_addr, 993

    def get_smtp_server(self):
        self.check_env_vars([
            "MAIL_USER_NAME",
            "MAIL_PASS_WORD",
            "MAIL_SMTP_ADDR",
        ])
        host, port = self.get_smtp_host_port()
        server = smtplib.SMTP_SSL(host, port)
        server.login(self.username, self.password)
        return server

    def get_imap_server(self):
        self.check_env_vars([
            "MAIL_USER_NAME",
            "MAIL_PASS_WORD",
            "MAIL_IMAP_ADDR",
        ])
        host, port = self.get_imap_host_port()
        server = imaplib.IMAP4_SSL(host, port)
        server.login(self.username, self.password)
        return server

    def decode_header_field(self, msg, field):
        import email.header
        val = msg.get(field, "")
        if not val:
            return ""
        parts = email.header.decode_header(val)
        result = []
        for v, enc in parts:
            if isinstance(v, bytes):
                try:
                    result.append(v.decode(enc or 'utf-8', errors='replace'))
                except Exception:
                    result.append(v.decode('utf-8', errors='replace'))
            else:
                result.append(str(v))
        return ''.join(result)

    def _decode_filename(self, filename):
        """解码邮件附件文件名，支持多种编码"""
        import email.header
        try:
            decoded_parts = email.header.decode_header(filename)
            decoded_filename = ''
            for content, encoding in decoded_parts:
                if not isinstance(content, bytes):
                    decoded_filename += str(content)
                    continue
                    
                # 处理字节内容
                if encoding:
                    # 如果指定了编码，直接使用
                    decoded_filename += content.decode(encoding, errors='replace')
                else:
                    # 尝试常见编码
                    for enc in ['utf-8', 'gb18030', 'gb2312', 'gbk', 'latin-1']:
                        try:
                            decoded_filename += content.decode(enc, errors='strict')
                            break
                        except UnicodeDecodeError:
                            continue
                    else:  # 所有编码都失败时
                        decoded_filename += content.decode('utf-8', errors='replace')
            return decoded_filename
        except Exception:
            # 解码失败时返回原始文件名
            return filename
    
    def _sanitize_filename(self, filename):
        """清理文件名，移除非法字符并限制长度"""
        import os
        import re
        # 移除非法字符
        safe_filename = re.sub(r'[\\/:*?"<>|]', '_', filename)
        
        # 限制文件名长度
        if len(safe_filename) > 200:  # 大多数文件系统限制在255字符左右
            name_part, ext_part = os.path.splitext(safe_filename)
            safe_filename = name_part[:195] + ext_part  # 保留扩展名但截断名称部分
            
        return safe_filename
    
    def _save_attachment_file(self, part, attach_dir, filename, attachments):
        """保存附件文件，处理错误并提供备用方案"""
        import os
        
        # 清理并获取安全的文件名
        safe_filename = self._sanitize_filename(filename)
        file_path = os.path.join(attach_dir, safe_filename)
        
        try:
            # 尝试保存文件
            with open(file_path, "wb") as f:
                f.write(part.get_payload(decode=True))
            return True
        except Exception as e:
            print(f"Error saving attachment {safe_filename}: {str(e)}")
            # 使用简化文件名作为备用方案
            simple_filename = f"attachment_{hash(filename) % 10000}{os.path.splitext(filename)[1]}"
            fallback_path = os.path.join(attach_dir, simple_filename)
            try:
                with open(fallback_path, "wb") as f:
                    f.write(part.get_payload(decode=True))
                # 更新列表中的文件名
                attachments[-1] = simple_filename
                return True
            except Exception:
                # 所有尝试都失败，跳过此附件
                attachments.pop()  # 从列表中移除失败的附件
                return False
    
    def save_attachments(self, msg, uid):
        """保存邮件附件并返回附件列表"""
        import os
        attachments = []
        
        if not msg.is_multipart():
            return attachments
            
        for part in msg.walk():
            filename = part.get_filename()
            if not filename:
                continue
                
            # 解码文件名
            filename = self._decode_filename(filename)
            attachments.append(filename)
            
            # 如果设置了保存路径，则保存附件
            if self.save_path:
                attach_dir = os.path.join(self.save_path, uid)
                os.makedirs(attach_dir, exist_ok=True)
                self._save_attachment_file(part, attach_dir, filename, attachments)
                
        return attachments

@mcp.tool()
def send_email(
    to: str,
    subject: str,
    content: str,
    attachments: List[str] = None
) -> Dict[str, Any]:
    """
    Send an email with optional attachments.
    Args:
        to: Recipient email address
        subject: Email subject
        content: Email body (plain text)
        attachments: List of absolute file paths to attach
    Returns:
        Dict with success status and message
    """
    try:
        client = EmailClient()
        server = client.get_smtp_server()
        msg = MIMEMultipart()
        msg['To'] = formataddr((to, to))
        msg['From'] = formataddr((client.from_name, client.from_addr))
        msg['Subject'] = subject
        msg.attach(MIMEText(content, 'plain', 'utf-8'))
        # Attach files
        if attachments:
            for file_path in attachments:
                if not os.path.isfile(file_path):
                    continue
                with open(file_path, 'rb') as f:
                    part = MIMEApplication(f.read(), Name=os.path.basename(file_path))
                    part['Content-Disposition'] = f'attachment; filename="{os.path.basename(file_path)}"'
                    msg.attach(part)
        server.sendmail(client.username, [to], msg.as_string())
        server.quit()
        return {"success": True, "message": "Email sent successfully"}
    except Exception as e:
        return {"success": False, "message": str(e)}

@mcp.tool()
def list_inbox(
    limit: int = 10,
    start: int = 0
) -> Dict[str, Any]:
    """
    List emails in the inbox with brief information.
    Args:
        limit: Maximum number of emails to return
        start: Start index (0-based, from the latest email)
    Returns:
        Dict with a list of emails (uid, subject, snippet, date, seen, attachment)
    """
    emails: List[Dict[str, str]] = []
    try:
        client = EmailClient()
        mail = client.get_imap_server()
        mail.select('INBOX')
        typ, data = mail.search(None, 'ALL')
        mail_ids = data[0].split()
        mail_ids = mail_ids[::-1][start:start+limit]
        for num in mail_ids:
            typ, msg_data = mail.fetch(num, '(RFC822 FLAGS UID)')
            msg = email.message_from_bytes(msg_data[0][1])
            subject = client.decode_header_field(msg, 'Subject')
            date_ = client.decode_header_field(msg, 'Date')
            snippet = ""
            attachment = 0
            if msg.is_multipart():
                for part in msg.walk():
                    if part.get_content_type() == "text/plain":
                        try:
                            charset = part.get_content_charset() or 'utf-8'
                            snippet = part.get_payload(decode=True).decode(charset, errors='ignore')
                            snippet = snippet.strip().replace('\r', '').replace('\n', ' ')
                        except Exception:
                            continue
                    if part.get_filename():
                        attachment += 1
            else:
                try:
                    charset = part.get_content_charset() or 'utf-8'
                    snippet = msg.get_payload(decode=True).decode(charset, errors='ignore')
                    snippet = snippet.strip().replace('\r', '').replace('\n', ' ')
                except Exception:
                    snippet = ""
            snippet = snippet[:120]
            typ, flag_data = mail.fetch(num, '(FLAGS)')
            seen = b'\\Seen' in flag_data[0] if flag_data and flag_data[0] else False
            typ, uid_data = mail.fetch(num, '(UID)')
            import re
            uid_match = re.search(br'UID (\d+)', uid_data[0])
            uid = uid_match.group(1).decode() if uid_match else ""
            emails.append({
                "uid": uid, "subject": subject, "snippet": snippet,
                "date": date_, "seen": seen, "attachment": attachment
            })
        mail.logout()
        return {"emails": emails}
    except Exception as e:
        return {"emails": [], "error": str(e)}

@mcp.tool()
def list_sent(
    limit: int = 10,
    start: int = 0
) -> Dict[str, Any]:
    """
    List emails in the sent mailbox with brief information.
    Args:
        limit: Maximum number of emails to return
        start: Start index (0-based, from the latest email)
    Returns:
        Dict with a list of emails (uid, subject, snippet, date, seen, attachment)
    """
    emails: List[Dict[str, str]] = []
    try:
        client = EmailClient()
        mail = client.get_imap_server()
        sent_folders = ["Sent", "Sent Mail", "Sent Messages", "已发送", "已发送邮件"]
        selected = False
        for folder in sent_folders:
            try:
                typ, _ = mail.select(folder)
                if typ == 'OK':
                    selected = True
                    break
            except Exception:
                continue
        if not selected:
            mail.logout()
            return {"emails": [], "error": "No sent folder found"}
        typ, data = mail.search(None, 'ALL')
        mail_ids = data[0].split()
        mail_ids = mail_ids[::-1][start:start+limit]
        for num in mail_ids:
            typ, msg_data = mail.fetch(num, '(RFC822 FLAGS UID)')
            msg = email.message_from_bytes(msg_data[0][1])
            subject = client.decode_header_field(msg, 'Subject')
            date_ = client.decode_header_field(msg, 'Date')
            snippet = ""
            attachment = 0
            if msg.is_multipart():
                for part in msg.walk():
                    if part.get_content_type() == "text/plain":
                        try:
                            charset = part.get_content_charset() or 'utf-8'
                            snippet = part.get_payload(decode=True).decode(charset, errors='ignore')
                            snippet = snippet.strip().replace('\r', '').replace('\n', ' ')
                        except Exception:
                            continue
                    if part.get_filename():
                        attachment += 1
            else:
                try:
                    snippet = msg.get_payload(decode=True).decode(msg.get_content_charset() or 'utf-8', errors='ignore')
                    snippet = snippet.strip().replace('\r', '').replace('\n', ' ')
                except Exception:
                    snippet = ""
            snippet = snippet[:120]
            typ, flag_data = mail.fetch(num, '(FLAGS)')
            seen = b'\\Seen' in flag_data[0] if flag_data and flag_data[0] else False
            typ, uid_data = mail.fetch(num, '(UID)')
            import re
            uid_match = re.search(br'UID (\d+)', uid_data[0])
            uid = uid_match.group(1).decode() if uid_match else ""
            emails.append({
                "uid": uid, "subject": subject, "snippet": snippet,
                "date": date_, "seen": seen, "attachment": attachment
            })
        mail.logout()
        return {"emails": emails}
    except Exception as e:
        return {"emails": [], "error": str(e)}

@mcp.tool()
def read_email(
    uid: str,
) -> Dict[str, Any]:
    """
    Read full details of an email by UID, and optionally save attachments.
    Args:
        uid: UID of the email
    Returns:
        Dict with all email fields (subject, snippet, text, html, attachments, uid, seen, from, to, date, message_id, in_reply_to, references)
    """
    import os
    try:
        client = EmailClient()
        mail = client.get_imap_server()
        mail.select('INBOX')
        typ, msg_data = mail.uid('fetch', uid, '(RFC822 FLAGS UID)')
        if not msg_data or not msg_data[0] or not isinstance(msg_data[0], tuple) or not msg_data[0][1]:
            mail.logout()
            return {"error": "Email content is empty"}
        msg = email.message_from_bytes(msg_data[0][1])
        typ, flag_data = mail.uid('fetch', uid, '(FLAGS)')
        seen = b'\\Seen' in flag_data[0] if flag_data and flag_data[0] else False
        uid_val = uid
        snippet = ""
        full_text = ""
        html_body = ""
        plain_body = ""
        attachments = client.save_attachments(msg, uid)
        # 获取 plain_body 和 html_body
        if msg.is_multipart():
            for part in msg.walk():
                ctype = part.get_content_type()
                charset = part.get_content_charset() or 'utf-8'
                if ctype == "text/plain" and not plain_body:
                    try:
                        part_text = part.get_payload(decode=True).decode(charset, errors='ignore')
                        part_text = part_text.strip().replace('\r', '').replace('\n', ' ')
                        plain_body = part_text
                    except Exception:
                        continue
                elif ctype == "text/html" and not plain_body:
                    try:
                        html = part.get_payload(decode=True).decode(charset, errors='ignore')
                        html_body = html.strip()
                    except Exception:
                        continue
        else:
            try:
                ctype = msg.get_content_type()
                charset = msg.get_content_charset() or 'utf-8'
                if ctype == "text/plain":
                    plain_body = msg.get_payload(decode=True).decode(charset, errors='ignore')
                    plain_body = plain_body.strip().replace('\r', '').replace('\n', ' ')
                elif ctype == "text/html":
                    html_body = msg.get_payload(decode=True).decode(charset, errors='ignore')
                    html_body = html_body.strip()
            except Exception:
                pass

        # snippet/正文预览/正文内容优先用 plain_body，没有时用 html_body 转文本
        import re
        if plain_body:
            full_text = plain_body
        elif html_body:
            text = re.sub('<[^<]+?>', '', html_body)
            text = text.strip().replace('\r', '').replace('\n', ' ')
            full_text = text
        else:
            full_text = ""
        snippet = full_text[:120]
        detail = {
            "subject": client.decode_header_field(msg, 'Subject'),
            "snippet": snippet, "text": full_text, "html": html_body,
            "attachments": attachments, "uid": uid_val, "seen": seen,
            "from": client.decode_header_field(msg, 'From'),
            "to": client.decode_header_field(msg, 'To'),
            "date": client.decode_header_field(msg, 'Date'),
            "message_id": client.decode_header_field(msg, 'Message-ID'),
            "in_reply_to": client.decode_header_field(msg, 'In-Reply-To'),
            "references": client.decode_header_field(msg, 'References'),
        }
        mail.logout()
        return detail
    except Exception as e:
        return {"error": str(e)}


@mcp.tool()
def reply_email(
    uid: str,
    subject: str,
    content: str,
    attachments: List[str] = None
) -> Dict[str, Any]:
    """
    Reply to an email, automatically setting reply headers and supporting attachments.
    Args:
        uid: UID of the original email to reply to
        subject: Email subject
        content: Email body (plain text or HTML, auto-detected)
        attachments: List of absolute file paths to attach
    Returns:
        Dict with success status and message
    """
    try:
        client = EmailClient()
        mail = client.get_imap_server()
        mail.select('INBOX')
        typ, msg_data = mail.uid('fetch', uid, '(RFC822)')
        if not msg_data or not msg_data[0] or not isinstance(msg_data[0], tuple) or not msg_data[0][1]:
            mail.logout()
            return {"success": False, "message": "Original email content is empty"}
        msg_orig = email.message_from_bytes(msg_data[0][1])
        original_from = client.decode_header_field(msg_orig, 'From')
        original_subject = client.decode_header_field(msg_orig, 'Subject')
        original_message_id = client.decode_header_field(msg_orig, 'Message-ID')
        mail.logout()
        html_signs = ['<html', '<body', '</p>', '<div', '<span',
                      '<br', '<table', '<tr', '<td', '<!DOCTYPE html']
        is_html = any(sign in content.lower() for sign in html_signs)
        server = client.get_smtp_server()
        msg = MIMEMultipart()
        msg['From'] = formataddr((client.from_name, client.from_addr))
        msg['To'] = original_from
        msg['Subject'] = subject
        if original_message_id:
            msg['In-Reply-To'] = original_message_id
            msg['References'] = original_message_id
        if original_from:
            msg['Reply-To'] = original_from
        if original_subject:
            msg['Thread-Topic'] = original_subject
        if is_html:
            msg.attach(MIMEText(content, 'html', 'utf-8'))
        else:
            msg.attach(MIMEText(content, 'plain', 'utf-8'))
        # Attach files
        if attachments:
            for file_path in attachments:
                if not os.path.isfile(file_path):
                    continue
                with open(file_path, 'rb') as f:
                    part = MIMEApplication(f.read(), Name=os.path.basename(file_path))
                    part['Content-Disposition'] = f'attachment; filename="{os.path.basename(file_path)}"'
                    msg.attach(part)
        server.sendmail(client.username, [original_from], msg.as_string())
        server.quit()
        return {"success": True, "message": "Reply email sent successfully"}
    except Exception as e:
        return {"success": False, "message": str(e)}

def main():
    mcp.run(show_banner=False)

if __name__ == "__main__":
    main()
