import imaplib
import smtplib
import ssl
import email
import re
import time
import json
import os
import sys
from email.header import decode_header, Header
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart

IMAP_HOST = 'mail.rrrosmontis.icu'
IMAP_PORT = 993
SMTP_HOST = 'mail.rrrosmontis.icu'
SMTP_PORT = 465


def decode_mime_str(s):
    if not s:
        return ''
    parts = decode_header(s)
    decoded_parts = []
    for part, charset in parts:
        if isinstance(part, bytes):
            try:
                if charset:
                    decoded_parts.append(part.decode(charset, errors='replace'))
                else:
                    decoded_parts.append(part.decode('utf-8', errors='replace'))
            except Exception:
                decoded_parts.append(part.decode('utf-8', errors='replace'))
        else:
            decoded_parts.append(part)
    return ''.join(decoded_parts)


def decode_attachment_names(raw):
    if not raw:
        return []
    parts = decode_header(raw)
    decoded_parts = []
    for part, charset in parts:
        if isinstance(part, bytes):
            try:
                if charset:
                    decoded_parts.append(part.decode(charset, errors='replace'))
                else:
                    decoded_parts.append(part.decode('utf-8', errors='replace'))
            except Exception:
                decoded_parts.append(part.decode('utf-8', errors='replace'))
        else:
            decoded_parts.append(part)
    combined = ''.join(decoded_parts)
    return [name.strip() for name in combined.split(',') if name.strip()]


def extract_email_body(msg):
    if msg.is_multipart():
        for part in msg.walk():
            if part.get_content_type() == "text/html":
                payload = part.get_payload(decode=True)
                charset = part.get_content_charset() or 'utf-8'
                try:
                    return payload.decode(charset, errors='replace')
                except Exception:
                    return payload.decode('utf-8', errors='replace')
        for part in msg.walk():
            if part.get_content_type() == "text/plain":
                payload = part.get_payload(decode=True)
                charset = part.get_content_charset() or 'utf-8'
                try:
                    return payload.decode(charset, errors='replace')
                except Exception:
                    return payload.decode('utf-8', errors='replace')
    else:
        payload = msg.get_payload(decode=True)
        charset = msg.get_content_charset() or 'utf-8'
        try:
            return payload.decode(charset, errors='replace')
        except Exception:
            return payload.decode('utf-8', errors='replace')
    return ''


def imap_connect(email_addr, password):
    context = ssl._create_unverified_context()
    conn = imaplib.IMAP4_SSL(IMAP_HOST, IMAP_PORT, ssl_context=context)
    conn.login(email_addr, password)
    return conn


def smtp_connect(email_addr, password):
    context = ssl._create_unverified_context()
    smtp = smtplib.SMTP_SSL(SMTP_HOST, SMTP_PORT, context=context)
    smtp.ehlo()
    smtp.login(email_addr, password)
    return smtp


def find_sent_folder(conn):
    typ, folder_data = conn.list()
    if typ != 'OK':
        return None
    sent_candidates = ['Sent', 'Sent Messages', 'Sent Items',
                       '\u5df2\u53d1\u9001', '\u5df2\u53d1\u9001\u90ae\u4ef6']
    for line in folder_data:
        line_str = line.decode() if isinstance(line, bytes) else line
        folder_name = None
        match = re.search(r'"([^"]+)"$', line_str)
        if match:
            folder_name = match.group(1)
        else:
            parts = line_str.split()
            if parts:
                folder_name = parts[-1].strip('"')
        if folder_name:
            if any(folder_name.lower() == cand.lower() for cand in sent_candidates):
                return folder_name
            if 'sent' in folder_name.lower() or '\u5df2\u53d1\u9001' in folder_name:
                return folder_name
    return 'Sent'


def list_mailboxes(email_addr, password):
    conn = imap_connect(email_addr, password)
    try:
        typ, data = conn.list()
        if typ != 'OK':
            return []
        mailboxes = []
        for line in data:
            line_str = line.decode() if isinstance(line, bytes) else line
            match = re.search(r'"([^"]+)"$', line_str)
            if match:
                mailboxes.append(match.group(1))
        return mailboxes
    finally:
        conn.logout()


def list_messages(email_addr, password, folder='INBOX', page=1, per_page=20):
    conn = imap_connect(email_addr, password)
    try:
        conn.select(folder)
        result, data = conn.uid('search', None, "ALL")
        if result != 'OK':
            return [], False

        uid_list = data[0].split()
        total = len(uid_list)

        start = max(0, total - page * per_page)
        end = total - (page - 1) * per_page
        page_uids = uid_list[start:end] if start < end else []

        messages = []
        for uid in page_uids:
            result, msg_data = conn.uid('fetch', uid, '(BODY.PEEK[HEADER])')
            if result != 'OK':
                continue
            raw_header = msg_data[0][1]
            msg = email.message_from_bytes(raw_header)
            messages.append({
                'uid': uid.decode() if isinstance(uid, bytes) else uid,
                'from': decode_mime_str(msg.get('From', '')),
                'subject': decode_mime_str(msg.get('Subject', '')),
                'date': msg.get('Date', ''),
                'attachments': decode_attachment_names(msg.get('X-Attachments', ''))
            })

        messages.reverse()
        return messages, page * per_page < total
    finally:
        conn.logout()


def read_message(email_addr, password, uid, folder='INBOX'):
    conn = imap_connect(email_addr, password)
    try:
        conn.select(folder)
        result, data = conn.uid('fetch', uid, '(RFC822)')
        if result != 'OK':
            return None
        raw_email = data[0][1]
        msg = email.message_from_bytes(raw_email)
        return {
            'uid': uid,
            'body': extract_email_body(msg),
            'subject': decode_mime_str(msg.get('Subject', '')),
            'from': decode_mime_str(msg.get('From', '')),
            'to': decode_mime_str(msg.get('To', '')),
            'date': msg.get('Date', ''),
            'attachments': decode_attachment_names(msg.get('X-Attachments', ''))
        }
    finally:
        conn.logout()


def list_sent(email_addr, password, page=1, per_page=20):
    conn = imap_connect(email_addr, password)
    try:
        sent_folder = find_sent_folder(conn)
        if not sent_folder:
            return [], False
        try:
            conn.select(sent_folder)
        except Exception:
            return [], False

        result, data = conn.uid('search', None, "ALL")
        if result != 'OK':
            return [], False

        uid_list = data[0].split()
        total = len(uid_list)

        start = max(0, total - page * per_page)
        end = total - (page - 1) * per_page
        page_uids = uid_list[start:end] if start < end else []

        messages = []
        for uid in page_uids:
            result, msg_data = conn.uid('fetch', uid, '(BODY.PEEK[HEADER])')
            if result != 'OK':
                continue
            raw_header = msg_data[0][1]
            msg = email.message_from_bytes(raw_header)
            messages.append({
                'uid': uid.decode() if isinstance(uid, bytes) else uid,
                'from': decode_mime_str(msg.get('From', '')),
                'to': decode_mime_str(msg.get('To', '')),
                'subject': decode_mime_str(msg.get('Subject', '')),
                'date': msg.get('Date', ''),
                'attachments': decode_attachment_names(msg.get('X-Attachments', ''))
            })

        messages.reverse()
        return messages, page * per_page < total
    finally:
        conn.logout()


def send_email(email_addr, password, to_addr, subject, body):
    smtp = smtp_connect(email_addr, password)
    try:
        msg = MIMEMultipart()
        msg['From'] = email_addr
        msg['To'] = to_addr
        msg['Subject'] = subject
        msg.attach(MIMEText(body, 'html', 'utf-8'))
        smtp.send_message(msg)
        smtp.quit()

        # 保存到已发送
        imap_conn = imap_connect(email_addr, password)
        try:
            sent_folder = find_sent_folder(imap_conn)
            if sent_folder:
                try:
                    imap_conn.select(sent_folder)
                except Exception:
                    try:
                        imap_conn.create(sent_folder)
                        imap_conn.select(sent_folder)
                    except Exception:
                        return True
                imap_conn.append(sent_folder, '\\Seen',
                                 imaplib.Time2Internaldate(time.time()),
                                 msg.as_bytes())
                imap_conn.close()
        finally:
            imap_conn.logout()
        return True
    except Exception as e:
        return str(e)


def delete_message(email_addr, password, uid, folder='INBOX'):
    conn = imap_connect(email_addr, password)
    try:
        conn.select(folder)
        result, _ = conn.uid('store', uid, '+FLAGS', '\\Deleted')
        if result != 'OK':
            return False
        conn.expunge()
        return True
    finally:
        conn.logout()


# ============================================================
# MCP JSON-RPC 处理
# ============================================================

TOOLS = [
    {
        "name": "list_mailboxes",
        "description": "列出 IMAP 邮箱中的所有文件夹",
        "inputSchema": {
            "type": "object",
            "properties": {
                "email": {"type": "string", "description": "邮箱地址"},
                "password": {"type": "string", "description": "邮箱密码，不传则从环境变量 ROS_MAIL_MCP_PASSWORD 读取"}
            },
            "required": ["email"]
        }
    },
    {
        "name": "list_messages",
        "description": "列出指定文件夹中的邮件（默认 INBOX），支持分页",
        "inputSchema": {
            "type": "object",
            "properties": {
                "email": {"type": "string", "description": "邮箱地址"},
                "password": {"type": "string", "description": "邮箱密码，不传则从环境变量 ROS_MAIL_MCP_PASSWORD 读取"},
                "folder": {"type": "string", "description": "文件夹名，默认 INBOX"},
                "page": {"type": "integer", "description": "页码，从 1 开始，默认 1"},
                "per_page": {"type": "integer", "description": "每页数量，默认 20"}
            },
            "required": ["email"]
        }
    },
    {
        "name": "read_message",
        "description": "读取指定 UID 的邮件全文",
        "inputSchema": {
            "type": "object",
            "properties": {
                "email": {"type": "string", "description": "邮箱地址"},
                "password": {"type": "string", "description": "邮箱密码，不传则从环境变量 ROS_MAIL_MCP_PASSWORD 读取"},
                "uid": {"type": "string", "description": "邮件 UID"},
                "folder": {"type": "string", "description": "文件夹名，默认 INBOX"}
            },
            "required": ["email", "uid"]
        }
    },
    {
        "name": "list_sent",
        "description": "列出已发送邮件，支持分页",
        "inputSchema": {
            "type": "object",
            "properties": {
                "email": {"type": "string", "description": "邮箱地址"},
                "password": {"type": "string", "description": "邮箱密码，不传则从环境变量 ROS_MAIL_MCP_PASSWORD 读取"},
                "page": {"type": "integer", "description": "页码，从 1 开始，默认 1"},
                "per_page": {"type": "integer", "description": "每页数量，默认 20"}
            },
            "required": ["email"]
        }
    },
    {
        "name": "send_email",
        "description": "发送邮件（HTML 格式），发送后会保存到已发送文件夹",
        "inputSchema": {
            "type": "object",
            "properties": {
                "email": {"type": "string", "description": "发件人邮箱地址"},
                "password": {"type": "string", "description": "发件人邮箱密码"},
                "to": {"type": "string", "description": "收件人邮箱地址"},
                "subject": {"type": "string", "description": "邮件主题"},
                "body": {"type": "string", "description": "邮件正文（HTML 格式）"}
            },
            "required": ["email", "to", "subject", "body"]
        }
    },
    {
        "name": "delete_message",
        "description": "删除指定 UID 的邮件",
        "inputSchema": {
            "type": "object",
            "properties": {
                "email": {"type": "string", "description": "邮箱地址"},
                "password": {"type": "string", "description": "邮箱密码，不传则从环境变量 ROS_MAIL_MCP_PASSWORD 读取"},
                "uid": {"type": "string", "description": "邮件 UID"},
                "folder": {"type": "string", "description": "文件夹名，默认 INBOX"}
            },
            "required": ["email", "uid"]
        }
    }
]

TOOL_MAP = {
    "list_mailboxes": list_mailboxes,
    "list_messages": list_messages,
    "read_message": read_message,
    "list_sent": list_sent,
    "send_email": send_email,
    "delete_message": delete_message
}


def resolve_password(arguments):
    if arguments.get('password'):
        return arguments['password']
    pw = os.environ.get('ROS_MAIL_MCP_PASSWORD', '')
    if pw:
        return pw
    return None


def handle_request(req):
    method = req.get("method", "")
    req_id = req.get("id")

    if method == "initialize":
        return {
            "jsonrpc": "2.0",
            "id": req_id,
            "result": {
                "protocolVersion": "2024-11-05",
                "capabilities": {
                    "tools": {}
                },
                "serverInfo": {
                    "name": "ros-mail-mcp",
                    "version": "1.0.0"
                }
            }
        }

    if method == "ping":
        return {"jsonrpc": "2.0", "id": req_id, "result": {}}

    if method == "notifications/initialized":
        return None

    if method == "tools/list":
        return {"jsonrpc": "2.0", "id": req_id, "result": {"tools": TOOLS}}

    if method == "tools/call":
        tool_name = req["params"]["name"]
        arguments = req["params"].get("arguments", {})
        tool_func = TOOL_MAP.get(tool_name)
        if not tool_func:
            return {"jsonrpc": "2.0", "id": req_id,
                    "error": {"code": -32601, "message": f"未知工具: {tool_name}"}}
        password = resolve_password(arguments)
        if not password:
            return {"jsonrpc": "2.0", "id": req_id,
                    "error": {"code": -32602, "message": "缺少密码: 请在参数中传入 password 或设置环境变量 ROS_MAIL_MCP_PASSWORD"}}
        arguments['password'] = password
        # 参数名映射：schema 中的参数名 -> Python 函数参数名
        if 'email' in arguments and 'email_addr' not in arguments:
            arguments['email_addr'] = arguments.pop('email')
        if 'to' in arguments and 'to_addr' not in arguments:
            arguments['to_addr'] = arguments.pop('to')
        try:
            result = tool_func(**arguments)
            if isinstance(result, tuple):
                result = {"data": result[0], "has_more": result[1]}
            elif isinstance(result, list):
                result = {"mailboxes": result}
            elif isinstance(result, dict):
                pass
            elif isinstance(result, bool):
                result = {"success": result}
            text = json.dumps(result, ensure_ascii=False, indent=2)
            return {"jsonrpc": "2.0", "id": req_id,
                    "result": {"content": [{"type": "text", "text": text}]}}
        except Exception as e:
            return {"jsonrpc": "2.0", "id": req_id,
                    "error": {"code": -32000, "message": str(e)}}

    return {"jsonrpc": "2.0", "id": req_id,
            "error": {"code": -32601, "message": f"未知方法: {method}"}}


def main():
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            req = json.loads(line)
            resp = handle_request(req)
            if resp is None:
                continue
            sys.stdout.write(json.dumps(resp, ensure_ascii=False) + '\n')
            sys.stdout.flush()
        except json.JSONDecodeError:
            pass


if __name__ == '__main__':
    main()