# 安装依赖（如已安装可跳过）
# hello-agents 框架 + 基础工具，imaplib 是 Python 标准库无需安装
# !pip install -q hello-agents python-dotenv rich

import os
import json
import re
import email
import imaplib
from email.header import decode_header
from datetime import datetime, timedelta, timezone
from typing import Dict, Any, List, Optional
from dataclasses import dataclass, field
from collections import Counter

# HelloAgents
from hello_agents import SimpleAgent, HelloAgentsLLM
from hello_agents.tools import Tool

# Utils
from dotenv import load_dotenv
from rich.console import Console
from rich.table import Table
from rich.panel import Panel
from rich.markdown import Markdown

console = Console()
load_dotenv(override=True)

# Init LLM
llm = HelloAgentsLLM()
model_id = os.getenv("LLM_MODEL_ID", "gpt-4o-mini")
print("✅ 环境配置完成")
print(f"   LLM 模型: {model_id}")

# ========================================
# 模拟邮件数据集（10封覆盖不同场景）
# ========================================

DEMO_EMAILS = [
    {
        "id": "1",
        "from": "[email protected]",
        "subject": "紧急：请今天下班前确认Q3预算方案",
        "date": "2026-07-02 08:30:00",
        "body": "各位部门负责人，附件是Q3预算方案初稿。请各位于今天下班前完成审核并回复确认。如有修改意见，请在明天上午的预算会议上提出。未及时确认的部门将按现有方案执行。\n\n附件：Q3_Budget_v2.xlsx"
    },
    {
        "id": "2",
        "from": "[email protected]",
        "subject": "Re: 合同条款第5条和第8条的修改意见",
        "date": "2026-07-02 09:15:00",
        "body": "你好，我们法务团队审核了合同草案，对第5条（付款条款）和第8条（违约责任）有较大异议。特别是违约金的设定比例需要调整。请尽快安排一次线上沟通，我们的法务希望本周内解决这个问题。"
    },
    {
        "id": "3",
        "from": "[email protected]",
        "subject": "【通知】年度绩效评估系统本周五关闭",
        "date": "2026-07-02 10:00:00",
        "body": "各位同事，本年度绩效自评和互评系统将于本周五（7月4日）17:00正式关闭。目前系统显示你还有2位同事的互评未完成。请尽快登录系统完成评估，逾期将影响年度绩效结果。\n\n评估系统链接：http://hr.company.com/performance"
    },
    {
        "id": "4",
        "from": "[email protected]",
        "subject": "[hello-agents] New PR #156: Add EmailDigestAgent project",
        "date": "2026-07-02 10:30:00",
        "body": "Pull Request #156 opened by WHS.\n\nTitle: Add EmailDigestAgent project\nBranch: feature/email-digest to main\nFiles changed: 5\n\nView on GitHub: https://github.com/datawhalechina/hello-agents/pull/156"
    },
    {
        "id": "5",
        "from": "[email protected]",
        "subject": "Tech Weekly #234: AI Agent 最新趋势与工具盘点",
        "date": "2026-07-02 11:00:00",
        "body": "本周精选：\n1. OpenAI 发布新一代 Agent 框架\n2. LangChain v0.3 重大更新\n3. 深度解析：多智能体协作的最佳实践\n4. 开源推荐：5个值得关注的 Agent 项目\n\n点击阅读全文：https://techweekly.com/234"
    },
    {
        "id": "6",
        "from": "[email protected]",
        "subject": "下午3点的项目例会议程更新",
        "date": "2026-07-02 12:00:00",
        "body": "Hi team, 今天下午3点的项目例会议程有更新：\n\n1. Sprint 回顾（15min）\n2. Bug 修复进展讨论（20min）\n3. 新需求评审——用户权限模块（25min）\n\n请提前查阅 Jira board 准备更新。"
    },
    {
        "id": "7",
        "from": "[email protected]",
        "subject": "年中大促最后一天！全场5折起",
        "date": "2026-07-02 13:00:00",
        "body": "年中狂欢倒计时！最后24小时！\n\n全场商品5折起\n满299包邮\n满599送精美礼品\n\n立即抢购：https://shop.example.com/sale\n退订请回复TD"
    },
    {
        "id": "8",
        "from": "[email protected]",
        "subject": "提醒：明天上午10:00 季度评审会议",
        "date": "2026-07-02 14:00:00",
        "body": "日历提醒：\n\n事件：季度评审会议\n时间：2026年7月3日 10:00-12:00\n地点：3楼会议室A / Zoom\n参与者：全体部门负责人\n\n请提前准备汇报材料。"
    },
    {
        "id": "9",
        "from": "[email protected]",
        "subject": "Congratulations! You've won a FREE iPhone!",
        "date": "2026-07-02 15:00:00",
        "body": "Dear Winner,\n\nYour email address has been selected in our annual lottery! You have won a brand new iPhone 16 Pro Max!\n\nTo claim your prize, click: http://suspicious-link.xyz/claim\n\nHurry! This offer expires in 24 hours!"
    },
    {
        "id": "10",
        "from": "[email protected]",
        "subject": "关于新同事入职培训和 mentor 安排",
        "date": "2026-07-02 16:30:00",
        "body": "大家下午好，下周一我们团队有新同事加入（前端开发，2年经验）。需要安排以下事项：\n\n1. 入职培训计划（HR 已发邮件）\n2. 指定一位 mentor（建议资深前端同事）\n3. 准备开发环境和新手引导文档\n\n请今天下班前在群里确认谁能担任 mentor！"
    }
]

print(f"✅ 已加载 {len(DEMO_EMAILS)} 封模拟邮件")

preview = "\n".join([
    f"  {i+1}. [{e['from']}] {e['subject'][:50]}..."
    for i, e in enumerate(DEMO_EMAILS)
])
console.print(Panel.fit(preview, title="📧 模拟邮件列表", style="blue"))

# ========================================
# 工具1: 邮件获取工具（IMAP 通用版）
# ========================================

def _decode_mime_header(header_value: str) -> str:
    if not header_value:
        return ""
    parts = decode_header(header_value)
    decoded = ""
    for text, charset in parts:
        if isinstance(text, bytes):
            try:
                decoded += text.decode(charset or 'utf-8', errors='replace')
            except (LookupError, UnicodeDecodeError):
                decoded += text.decode('utf-8', errors='replace')
        else:
            decoded += str(text)
    return decoded


def _parse_email_body(msg) -> str:
    body = ""
    if msg.is_multipart():
        for part in msg.walk():
            content_type = part.get_content_type()
            disposition = str(part.get("Content-Disposition", ""))
            if content_type == "text/plain" and "attachment" not in disposition:
                payload = part.get_payload(decode=True)
                if payload:
                    charset = part.get_content_charset() or 'utf-8'
                    try:
                        body = payload.decode(charset, errors='replace')
                    except (LookupError, UnicodeDecodeError):
                        body = payload.decode('utf-8', errors='replace')
                    break
        if not body:
            for part in msg.walk():
                if (part.get_content_type() == "text/html"
                        and "attachment" not in str(part.get("Content-Disposition", ""))):
                    payload = part.get_payload(decode=True)
                    if payload:
                        charset = part.get_content_charset() or 'utf-8'
                        try:
                            body = payload.decode(charset, errors='replace')
                            body = re.sub(r'<[^>]+>', '', body)
                            body = re.sub(r'\s+', ' ', body).strip()
                        except Exception:
                            body = payload.decode('utf-8', errors='replace')
                    break
    else:
        payload = msg.get_payload(decode=True)
        if payload:
            charset = msg.get_content_charset() or 'utf-8'
            try:
                body = payload.decode(charset, errors='replace')
            except Exception:
                body = payload.decode('utf-8', errors='replace')
    return body[:3000]


class EmailFetchTool(Tool):
    def __init__(self, use_demo: bool = True, demo_emails=None):
        super().__init__(
            name="email_fetch",
            description="获取收件箱中的未读邮件"
        )
        self.use_demo = use_demo
        self.demo_emails = demo_emails or DEMO_EMAILS

    def _read_imap_config(self) -> tuple:
        server = os.getenv("IMAP_SERVER") or ""
        port = int(os.getenv("IMAP_PORT") or "0")
        username = os.getenv("IMAP_USERNAME") or ""
        password = os.getenv("IMAP_PASSWORD") or ""
        if not (username and password and server):
            try:
                with open("config/email_config.json", "r", encoding="utf-8") as f:
                    cfg = json.load(f)
                imap_cfg = cfg.get("imap", {})
                server = server or imap_cfg.get("server", "imap.qq.com")
                port = port or imap_cfg.get("port", 993)
                username = username or imap_cfg.get("username", "")
                password = password or imap_cfg.get("password", "")
            except FileNotFoundError:
                pass
        return server or "imap.qq.com", port or 993, username, password

    def _fetch_via_imap(self, max_emails: int = 50, hours: int = 24) -> list:
        server, port, username, password = self._read_imap_config()
        if not username or not password:
            raise ValueError("未配置邮箱信息")

        console.print(f"[dim]  连接 {server}:{port}  用户: {username}...[/dim]")
        mail = imaplib.IMAP4_SSL(server, port)
        mail.login(username, password)

        # 选择收件箱 — 依次尝试: 无引号, 带引号, readonly
        select_status = None
        for name, rdonly in [("INBOX", False), ('"INBOX"', False), ("INBOX", True)]:
            select_status, select_data = mail.select(name, readonly=rdonly)
            console.print(f"[dim]  select({name!r}, readonly={rdonly}) -> {select_status!r}[/dim]")
            if select_status == "OK":
                break
        if select_status != "OK":
            raise ValueError(f"无法选择 INBOX (状态: {select_status!r})")

        since_date = (datetime.now() - timedelta(hours=hours)).strftime("%d-%b-%Y")
        search_criteria = f'(UNSEEN SINCE {since_date})'
        status, message_ids = mail.search(None, search_criteria)
        if status != "OK":
            mail.logout()
            return []

        ids = message_ids[0].split()
        ids = list(reversed(ids))[:max_emails]
        emails = []
        for msg_id in ids:
            try:
                status, msg_data = mail.fetch(msg_id, "(RFC822)")
                if status != "OK":
                    continue
                raw_email = msg_data[0][1]
                msg = email.message_from_bytes(raw_email)
                msg_id_str = msg_id.decode() if isinstance(msg_id, bytes) else msg_id
                emails.append({
                    "id": msg_id_str,
                    "from": _decode_mime_header(msg.get("From", "")),
                    "subject": _decode_mime_header(msg.get("Subject", "(无主题)")),
                    "date": msg.get("Date", ""),
                    "body": _parse_email_body(msg)
                })
            except Exception as e:
                console.print(f"[yellow]  ⚠️ 邮件解析失败: {e}[/yellow]")
                continue

        mail.logout()
        console.print(f"[dim]  已断开连接[/dim]")
        return emails

    def run(self, hours: int = 24, max_emails: int = 50) -> str:
        if self.use_demo:
            emails = self.demo_emails[:max_emails]
        else:
            try:
                emails = self._fetch_via_imap(max_emails=max_emails, hours=hours)
            except imaplib.IMAP4.error as e:
                console.print(f"[red]❌ IMAP 连接失败: {e}[/red]")
                raise RuntimeError(f"IMAP 连接失败: {e}") from e
            except ValueError:
                raise
            except Exception as e:
                console.print(f"[red]❌ 未知错误: {e}[/red]")
                raise

        if not emails:
            return json.dumps({"message": "没有新的未读邮件", "emails": []}, ensure_ascii=False, indent=2)
        return json.dumps({"count": len(emails), "fetch_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"), "emails": emails}, ensure_ascii=False, indent=2)

    def get_parameters(self):
        from hello_agents.tools import ToolParameter
        return []

print("✅ EmailFetchTool defined (IMAP version)")

# ========================================
# 工具2: 邮件摘要与分类工具（LLM驱动）
# ========================================

class EmailDigestTool(Tool):
    """用 LLM 对邮件列表进行分类、摘要和优先级判断"""

    CATEGORIES = ["工作", "客户", "个人", "通知", "促销", "垃圾"]
    PRIORITIES = ["高", "中", "低"]

    def __init__(self, llm=None, use_demo: bool = True, demo_emails=None):
        super().__init__(
            name="email_fetch",
            description="获取收件箱中的未读邮件。支持演示模式和IMAP真实接入模式。"
        )
        self.llm = llm
        self.use_demo = use_demo
        self.demo_emails = demo_emails or DEMO_EMAILS

    def _build_prompt(self, emails_json: str, count: int) -> str:
        cats = ", ".join(self.CATEGORIES)
        prios = ", ".join(self.PRIORITIES)
        return f"""你是一个专业的邮件分析助手。请分析以下邮件列表并返回结构化结果。

邮件数据（JSON格式）：
{emails_json}

分析规则：
1. 分类 - 从 [{cats}] 中选择最合适的一类
2. 优先级 - 从 [{prios}] 中选择：
   高: 需要立即处理（老板指令、客户需求、合同、截止日期）
   中: 需要处理但不紧急（项目讨论、同事沟通、HR通知）
   低: 可忽略（新闻订阅、促销广告、系统通知、垃圾邮件）
3. 一句话摘要 - 用中文提炼核心内容，包含关键动作或截止时间
4. 关键行动 - 若需要回复或行动，用1-2句话说明；否则填"无"

严格按以下JSON格式返回（只返回JSON，不加其他文字）：
{{
  "analyzed_at": "当前时间",
  "total": {count},
  "items": [
    {{
      "id": "邮件ID",
      "from": "发件人",
      "subject": "主题",
      "category": "分类",
      "priority": "优先级",
      "summary": "一句话摘要",
      "action": "关键行动"
    }}
  ]
}}"""

    def run(self, emails_json: str) -> str:
        try:
            data = json.loads(emails_json)
            count = data.get("count", len(data.get("emails", [])))
            prompt = self._build_prompt(emails_json, count)
            response = self.llm.invoke([{"role": "user", "content": prompt}])
            result = response.content.strip()
            
            if result.startswith("```"):
                result = re.sub(r'^```(?:json)?\s*', '', result)
                result = re.sub(r'\s*```$', '', result)

            parsed = json.loads(result)
            return json.dumps(parsed, ensure_ascii=False, indent=2)

        except json.JSONDecodeError as e:
            return json.dumps({
                "error": "LLM 返回格式异常",
                "detail": str(e)
            }, ensure_ascii=False, indent=2)
        except Exception as e:
            return json.dumps({
                "error": "分析过程中出错",
                "detail": str(e)
            }, ensure_ascii=False, indent=2)
    def get_parameters(self):
        from hello_agents.tools import ToolParameter
        return []


print("✅ EmailDigestTool defined")

# ========================================
# 邮件日报生成 Pipeline
# ========================================

@dataclass
class EmailDigestPipeline:
    """邮件日报 Pipeline: 获取(IMAP/演示) -> LLM分类摘要 -> Markdown日报"""

    llm: any
    use_demo: bool = True
    fetch_tool: EmailFetchTool = field(init=False)
    digest_tool: EmailDigestTool = field(init=False)

    def __post_init__(self):
        self.fetch_tool = EmailFetchTool(use_demo=self.use_demo)
        self.digest_tool = EmailDigestTool(llm=self.llm)

    def run(self, hours: int = 24, max_emails: int = 50) -> str:
        """执行完整 Pipeline，返回 Markdown 日报"""
        mode = "演示数据" if self.use_demo else "IMAP 真实邮箱"
        console.print(Panel.fit(
            f"模式: {mode}  |  时间: 最近{hours}h  |  上限: {max_emails}封",
            title="📬 EmailDigestAgent", style="cyan"
        ))

        # Step 1: Fetch
        console.print("\n📥 [1/4] 获取邮件...", style="cyan")
        raw = self.fetch_tool.run(hours=hours, max_emails=max_emails)
        data = json.loads(raw)
        count = data.get("count", len(data.get("emails", [])))
        if count == 0:
            return self._empty_report()
        console.print(f"   ✅ {count} 封邮件")

        # Step 2: LLM digest
        console.print("\n🧠 [2/4] LLM 分类与摘要...", style="cyan")
        digest_raw = self.digest_tool.run(raw)
        digest = json.loads(digest_raw)
        items = digest.get("items", [])
        if "error" in digest:
            console.print(f"   ⚠️  {digest['error']}")
        console.print(f"   ✅ 已分析 {len(items)} 封")

        # Step 3: Stats
        console.print("\n📊 [3/4] 统计汇总...", style="cyan")
        stats = self._compute_stats(items)
        console.print(f"   🔴高:{stats['priorities'].get('高',0)}  🟡中:{stats['priorities'].get('中',0)}  🟢低:{stats['priorities'].get('低',0)}")

        # Step 4: Report
        console.print("\n📝 [4/4] 生成日报...", style="cyan")
        report = self._generate_report(stats, items)

        os.makedirs("output", exist_ok=True)
        path = f"output/email_digest_{datetime.now().strftime('%Y%m%d')}.md"
        with open(path, 'w', encoding='utf-8') as f:
            f.write(report)
        console.print(f"\n📄 已保存: {path}", style="bold green")
        return report

    def _compute_stats(self, items):
        cats = Counter(i.get("category", "other") for i in items)
        prios = Counter(i.get("priority", "中") for i in items)
        actions = sum(1 for i in items if i.get("action", "无") != "无")
        return {"total": len(items), "categories": dict(cats),
                "priorities": dict(prios), "need_action": actions}

    def _generate_report(self, stats, items):
        now = datetime.now().strftime("%Y-%m-%d %H:%M")
        src = "演示数据" if self.use_demo else "IMAP 收件箱"
        order = {"高": 0, "中": 1, "低": 2}
        sorted_items = sorted(items, key=lambda x: order.get(x.get("priority", "中"), 1))

        lines = [
            f"# 📬 邮件日报",
            f"**生成时间：** {now}  |  **来源：** {src}",
            f"---",
            f"## 📊 概览",
            f"| 指标 | 数值 |",
            f"|------|------|",
            f"| 总计 | {stats['total']} 封 |",
            f"| 🔴 高优先级 | {stats['priorities'].get('高', 0)} |",
            f"| 🟡 中优先级 | {stats['priorities'].get('中', 0)} |",
            f"| 🟢 低优先级 | {stats['priorities'].get('低', 0)} |",
            f"| 📋 需行动 | {stats['need_action']} |",
            f"",
            f"### 分类分布",
        ]
        for cat, cnt in stats['categories'].items():
            lines.append(f"- **{cat}**：{cnt} 封 {'█' * cnt}")
        lines.append(f"---")

        for pri, label in [("高", "🔴 高优先级 — 立即处理"),
                            ("中", "🟡 中优先级 — 今天处理"),
                            ("低", "🟢 低优先级 — 可忽略")]:
            pitems = [i for i in sorted_items if i.get("priority") == pri]
            if not pitems:
                continue
            lines.append(f"## {label}")
            lines.append(f"| # | 发件人 | 主题 | 类型 | 一句话摘要 |")
            lines.append(f"|---|--------|------|------|-----------|")
            for idx, item in enumerate(pitems, 1):
                fr = item.get("from", "")
                subj = item.get("subject", "")
                cat = item.get("category", "-")
                summ = item.get("summary", "-")
                lines.append(f"| {idx} | {fr} | {subj} | {cat} | {summ} |")
            lines.append("")
            if pri in ("高", "中"):
                acts = [i for i in pitems if i.get("action", "无") != "无"]
                if acts:
                    lines.append("**📋 行动项：**")
                    for a in acts:
                        lines.append(f"- [{a.get('category','')}] **{a.get('subject','')}** — {a.get('action','')}")
                    lines.append("")

        lines.extend(["---", "*EmailDigestAgent · HelloAgents*"])
        return "\n".join(lines)

    def _empty_report(self):
        now = datetime.now().strftime("%Y-%m-%d %H:%M")
        return (f"# 📬 邮件日报\n\n**{now}**\n\n## 🎉 收件箱干净！\n\n"
                f"当前没有未读邮件。\n\n---\n*EmailDigestAgent · HelloAgents*")

print("✅ EmailDigestPipeline defined")

# 创建 Pipeline（演示模式）并运行
pipeline = EmailDigestPipeline(llm=llm, use_demo=True)
report = pipeline.run(hours=24, max_emails=20)

# 渲染日报
console.print("\n" + "=" * 60)
console.print("📬 邮件日报", style="bold cyan")
console.print("=" * 60 + "\n")
console.print(Markdown(report))

import glob
reports = sorted(glob.glob("output/email_digest_*.md"), reverse=True)
if reports:
    print(f"Latest: {reports[0]}")
    with open(reports[0], 'r', encoding='utf-8') as f:
        console.print(Markdown(f.read()))
else:
    print("No reports yet. Run the pipeline first.")

# 分类和优先级统计
emails_data = json.loads(pipeline.fetch_tool.run(hours=24, max_emails=20))
digest_data = json.loads(pipeline.digest_tool.run(json.dumps(emails_data)))
items = digest_data.get("items", [])

table = Table(title="📊 邮件分析统计")
table.add_column("指标", style="cyan")
table.add_column("数值", style="white")
table.add_column("说明", style="dim")
table.add_row("总数", str(len(items)), "")
table.add_row("需立即处理(高)", str(sum(1 for i in items if i.get("priority")=="高")), "🔴")
table.add_row("今天处理(中)", str(sum(1 for i in items if i.get("priority")=="中")), "🟡")
table.add_row("可延后(低)", str(sum(1 for i in items if i.get("priority")=="低")), "🟢")
table.add_row("需要行动", str(sum(1 for i in items if i.get("action","无")!="无")), "📋")
console.print(table)

console.print("\n📂 分类分布：")
cats = Counter(i.get("category","other") for i in items)
for cat, cnt in cats.most_common():
    bar = "█" * (cnt * 3)
    console.print(f"  {cat:8s}  {cnt:2d}  {bar}")

print("=" * 60)
print("📧 真实邮箱接入指南")
print("=" * 60)
print()
print("【各邮箱 IMAP 配置】")
print()
print("  QQ邮箱:")
print("    server: imap.qq.com  port: 993")
print("    步骤: QQ邮箱网页版 -> 设置 -> 账户 -> 开启IMAP/SMTP服务")
print("          -> 验证密保 -> 获取授权码 (注意不是QQ密码)")
print()
print("  网易163:")
print("    server: imap.163.com  port: 993")
print("    步骤: 网易邮箱网页版 -> 设置 -> POP3/SMTP/IMAP -> 开启IMAP")
print("          -> 获取授权码")
print()
print("  网易126:")
print("    server: imap.126.com  port: 993")
print("    步骤: 同163邮箱")
print()
print("  Gmail:")
print("    server: imap.gmail.com  port: 993")
print("    步骤: Google账户 -> 安全性 -> 两步验证 -> 应用专用密码")
print()
print("  Outlook/Hotmail:")
print("    server: outlook.office365.com  port: 993")
print("    步骤: Microsoft账户 -> 安全性 -> 应用密码")
print()
print("【配置方式二选一】")
print()
print("  方式1 (推荐) - 编辑 .env 文件:")
print("    IMAP_SERVER=imap.qq.com")
print("    IMAP_PORT=993")
print("    [email protected]")
print("    IMAP_PASSWORD=你的授权码")
print()
print("  方式2 - 编辑 config/email_config.json:")
print("    修改 imap 段中的 server / port / username / password")
print()
print("【运行真实模式】")
print("  pipeline = EmailDigestPipeline(llm=llm, use_demo=False)")
print("  report = pipeline.run(hours=24, max_emails=50)")
print("  console.print(Markdown(report))")
print()
print("=" * 60)

# ========================================
# 真实 IMAP 模式（取消注释以运行）
# ========================================
#
# 配置好 .env 中的 IMAP_SERVER/IMAP_PORT/IMAP_USERNAME/IMAP_PASSWORD 后运行：
#
# pipeline_real = EmailDigestPipeline(llm=llm, use_demo=False)
# report = pipeline_real.run(hours=24, max_emails=50)
# console.print(Markdown(report))

print("💡 取消上面代码的注释，配置好 .env 中的 IMAP 信息后即可运行真实模式。")
print("如果 IMAP 连接失败，会直接报错并展示原因，方便排查。")