#!/usr/bin/env python3
"""
pptagent67/run.py — 简化版 Runner
===================================

用本地工具 stub 替代尚不存在的 MCP 服务，直接通过 lambdagent runtime
执行 pptagent67 多智能体 PPT 生成流程。

使用 Claude Opus 4.6 (1M context) 的原生能力:
  - web_search → 补充知识库之外的信息
  - python_exec → 辅助脚本执行（marp-cli 渲染）
  - local fs    → 替代 fs-mcp

用法:
  export ANTHROPIC_API_KEY=sk-ant-...
  python run.py                           # 完整流程
  python run.py --phase outline           # 只跑某个阶段
  python run.py --topic "AI 在医疗中的应用"  # 指定主题
  python run.py --template business       # 指定模板风格
"""
from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
import time
import shutil
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional

# ── 路径设置 ──────────────────────────────────────────────
PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
AGENT_DIR = Path(__file__).resolve().parent
sys.path.insert(0, str(PROJECT_ROOT))

from lambdagent.fromconfig import from_config
from lambdagent.agentruntime.runtime import Runtime, RuntimeResult
from lambdagent.agentruntime.config import RuntimeConfig
from lambdagent.core import Context


# ═══════════════════════════════════════════════════════════
# 1. 本地工具 Stubs — 替代 MCP 服务
# ═══════════════════════════════════════════════════════════

class LocalToolkit:
    """
    将 MCP 服务的功能用本地实现替代。
    每个方法签名对齐 MCP tool 的 (input: str) -> str 接口。
    """

    def __init__(self, workspace: Path):
        self.workspace = workspace

    # ── fs-mcp 替代 ──────────────────────────────────────
    def read_file(self, path: str) -> str:
        p = self._resolve(path)
        if not p.exists():
            return f"[ERROR] File not found: {p}"
        return p.read_text(encoding="utf-8")

    def write_file(self, args: str) -> str:
        """args = JSON: {"path": "...", "content": "..."}"""
        d = json.loads(args)
        p = self._resolve(d["path"])
        p.parent.mkdir(parents=True, exist_ok=True)
        p.write_text(d["content"], encoding="utf-8")
        return f"Written {len(d['content'])} bytes to {p}"

    def list_dir(self, path: str) -> str:
        p = self._resolve(path)
        if not p.is_dir():
            return f"[ERROR] Not a directory: {p}"
        entries = sorted(p.iterdir())
        return "\n".join(f"{'[DIR] ' if e.is_dir() else ''}{e.name}" for e in entries)

    def mkdir(self, path: str) -> str:
        p = self._resolve(path)
        p.mkdir(parents=True, exist_ok=True)
        return f"Created: {p}"

    def copy_file(self, args: str) -> str:
        d = json.loads(args)
        src, dst = self._resolve(d["src"]), self._resolve(d["dst"])
        dst.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(src, dst)
        return f"Copied {src} → {dst}"

    # ── Marp 渲染 ─────────────────────────────────────
    def render_marp(self, args: str) -> str:
        """使用 marp-cli 将 Marp Markdown 转换为 .pptx 和 .pdf"""
        d = json.loads(args)
        md_path = self._resolve(d.get("input", "presentation.md"))
        output_dir = self._resolve(d.get("output_dir", "."))

        if not md_path.exists():
            return f"[ERROR] Markdown not found: {md_path}"

        results = []

        # 检测 marp-cli
        marp_cmd = self._find_marp()
        if not marp_cmd:
            return "[ERROR] marp-cli not found. Install: npm install -g @marp-team/marp-cli"

        # 转换为 PPTX
        pptx_path = output_dir / "presentation.pptx"
        try:
            r = subprocess.run(
                [*marp_cmd, "--pptx", str(md_path), "-o", str(pptx_path), "--allow-local-files"],
                capture_output=True, text=True, timeout=120,
                cwd=str(md_path.parent),
            )
            if r.returncode == 0 and pptx_path.exists():
                results.append(f"PPTX: {pptx_path} ({pptx_path.stat().st_size} bytes)")
            else:
                results.append(f"[PPTX_ERROR] {r.stderr[:500]}")
        except Exception as e:
            results.append(f"[PPTX_ERROR] {e}")

        # 转换为 PDF
        pdf_path = output_dir / "presentation.pdf"
        try:
            r = subprocess.run(
                [*marp_cmd, "--pdf", str(md_path), "-o", str(pdf_path), "--allow-local-files"],
                capture_output=True, text=True, timeout=120,
                cwd=str(md_path.parent),
            )
            if r.returncode == 0 and pdf_path.exists():
                results.append(f"PDF: {pdf_path} ({pdf_path.stat().st_size} bytes)")
            else:
                results.append(f"[PDF_WARN] {r.stderr[:300]}")
        except Exception as e:
            results.append(f"[PDF_WARN] {e}")

        return "\n".join(results)

    def _find_marp(self) -> list:
        """查找 marp-cli 可执行文件"""
        # 优先全局安装
        for cmd in ["marp", "npx"]:
            try:
                if cmd == "npx":
                    r = subprocess.run(
                        ["npx", "@marp-team/marp-cli", "--version"],
                        capture_output=True, text=True, timeout=30,
                    )
                    if r.returncode == 0:
                        return ["npx", "@marp-team/marp-cli"]
                else:
                    r = subprocess.run(
                        ["marp", "--version"],
                        capture_output=True, text=True, timeout=10,
                    )
                    if r.returncode == 0:
                        return ["marp"]
            except FileNotFoundError:
                continue
        return []

    def python_exec(self, code: str) -> str:
        """在本地执行 Python 代码"""
        try:
            result = subprocess.run(
                [sys.executable, "-c", code],
                capture_output=True, text=True, timeout=120,
                cwd=str(self.workspace),
            )
            output = result.stdout
            if result.returncode != 0:
                output += f"\n[STDERR] {result.stderr}"
            return output[:5000]
        except subprocess.TimeoutExpired:
            return "[TIMEOUT] Python execution exceeded 120s"
        except Exception as e:
            return f"[EXEC_ERROR] {e}"

    def install_package(self, pkg: str) -> str:
        try:
            result = subprocess.run(
                [sys.executable, "-m", "pip", "install", pkg, "-q"],
                capture_output=True, text=True, timeout=60,
            )
            return result.stdout + result.stderr
        except Exception as e:
            return f"[PIP_ERROR] {e}"

    # ── 特殊本地工具 ─────────────────────────────────────
    def terminate(self, result: str) -> str:
        return result

    def read_topic(self, path: str = "") -> str:
        """读取 PPT 主题文件"""
        topic_path = AGENT_DIR / "TOPIC.md"
        if not topic_path.exists():
            return f"[ERROR] TOPIC.md not found at {topic_path}"
        return topic_path.read_text(encoding="utf-8")

    def save_checkpoint(self, data: str) -> str:
        ckpt_dir = self.workspace / ".checkpoints"
        ckpt_dir.mkdir(parents=True, exist_ok=True)
        ts = datetime.now().strftime("%Y%m%d_%H%M%S")
        path = ckpt_dir / f"checkpoint_{ts}.json"
        path.write_text(data, encoding="utf-8")
        return f"Checkpoint saved: {path}"

    def load_checkpoint(self, pattern: str = "latest") -> str:
        ckpt_dir = self.workspace / ".checkpoints"
        if not ckpt_dir.exists():
            return "[ERROR] No checkpoints found"
        files = sorted(ckpt_dir.glob("checkpoint_*.json"))
        if not files:
            return "[ERROR] No checkpoints found"
        return files[-1].read_text(encoding="utf-8")

    def human_feedback(self, question: str) -> str:
        print(f"\n{'='*60}")
        print(f"[HUMAN FEEDBACK REQUESTED]")
        print(f"{question}")
        print(f"{'='*60}")
        try:
            answer = input("Your response (or press Enter to skip): ").strip()
            return answer if answer else "[NO FEEDBACK]"
        except (EOFError, KeyboardInterrupt):
            return "[NO FEEDBACK]"

    def dispatch_agent(self, args: str) -> str:
        """调度子智能体执行"""
        d = json.loads(args)
        agent_name = d.get("agent", "")
        agent_input = d.get("input", "")

        agent_yml = AGENT_DIR / "agents" / f"{agent_name}.yml"
        if not agent_yml.exists():
            return f"[ERROR] Agent config not found: {agent_yml}"

        print(f"\n  >> Dispatching sub-agent: {agent_name}")
        print(f"     Config: {agent_yml}")

        try:
            result = Runtime.execute(str(agent_yml), agent_input)
            return result.result
        except Exception as e:
            return f"[AGENT_ERROR] {agent_name} failed: {e}"

    # ── 内部辅助 ─────────────────────────────────────────
    def _resolve(self, path: str) -> Path:
        p = Path(path)
        if p.is_absolute():
            return p
        return self.workspace / p

    def get_tool_map(self) -> Dict[str, callable]:
        """返回工具名 → 函数的映射，用于注入 runtime"""
        return {
            # fs
            "read_file": self.read_file,
            "write_file": self.write_file,
            "list_dir": self.list_dir,
            "mkdir": self.mkdir,
            "copy_file": self.copy_file,
            # marp
            "render_marp": self.render_marp,
            "python_exec": self.python_exec,
            "install_package": self.install_package,
            # local
            "terminate": self.terminate,
            "read_topic": self.read_topic,
            "save_checkpoint": self.save_checkpoint,
            "load_checkpoint": self.load_checkpoint,
            "human_feedback": self.human_feedback,
            "dispatch_agent": self.dispatch_agent,
        }


# ═══════════════════════════════════════════════════════════
# 2. Runner — 驱动 PPT 生成流程
# ═══════════════════════════════════════════════════════════

PHASES = [
    ("01_knowledge_extraction", "knowledge-extractor"),
    ("02_outline_generation",   "outline-generator"),
    ("03_content_refinement",   "content-refiner"),
    ("04_ppt_rendering",        "ppt-renderer"),
]


def create_workspace() -> Path:
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    ws = AGENT_DIR / "workspace" / f"run_{ts}"
    ws.mkdir(parents=True, exist_ok=True)
    return ws


def run_phase(
    phase_name: str,
    agent_name: str,
    workspace: Path,
    round_num: int,
    toolkit: LocalToolkit,
    prev_reports: List[str],
    revision_context: str = "",
    topic: str = "",
    template: str = "academic",
) -> Dict[str, Any]:
    """运行单个阶段的子智能体"""
    phase_dir = workspace / f"round_{round_num}" / phase_name
    phase_dir.mkdir(parents=True, exist_ok=True)
    (phase_dir / "artifacts").mkdir(exist_ok=True)

    agent_yml = AGENT_DIR / "agents" / f"{agent_name}.yml"

    # 构建输入上下文
    context_parts = [
        f"workspace: {workspace}",
        f"phase_dir: {phase_dir}",
        f"round: {round_num}",
        f"template: {template}",
        f"dependencies: {json.dumps(prev_reports)}",
    ]
    if revision_context:
        context_parts.append(f"revision_context: {revision_context}")

    # 读取主题
    if topic:
        context_parts.append(f"\n## PPT 主题:\n{topic}")
    else:
        topic_content = toolkit.read_topic()
        context_parts.append(f"\n## TOPIC.md Content:\n{topic_content}")

    # 附上前序阶段的关键结果摘要
    for rp in prev_reports[-3:]:
        rp_path = workspace / rp
        if rp_path.exists():
            try:
                data = json.loads(rp_path.read_text(encoding="utf-8"))
                summary = {k: v for k, v in data.items() if k == "_meta" or not isinstance(v, (list, dict))}
                context_parts.append(f"\n## Previous report ({rp}):\n{json.dumps(summary, ensure_ascii=False, indent=2)[:2000]}")
            except Exception:
                pass

    input_text = "\n".join(context_parts)

    print(f"\n{'━'*60}")
    print(f"  Phase: {phase_name}")
    print(f"  Agent: {agent_name} → {agent_yml.name}")
    print(f"  Output: {phase_dir}")
    print(f"{'━'*60}")

    t0 = time.time()
    try:
        result = Runtime.execute(str(agent_yml), input_text)
        duration = time.time() - t0
        print(f"  ✓ Completed in {duration:.1f}s, {len(result.trace)} trace steps")

        # 保存 trace
        trace_path = phase_dir / "trace.json"
        trace_data = [
            {"step": t.step, "term": t.term_name, "type": t.term_type,
             "duration_ms": t.duration_ms, "output": t.output[:200]}
            for t in result.trace
        ]
        trace_path.write_text(json.dumps(trace_data, indent=2), encoding="utf-8")

        # 尝试解析 report.json
        report_path = phase_dir / "report.json"
        if report_path.exists():
            report = json.loads(report_path.read_text(encoding="utf-8"))
        else:
            try:
                report = json.loads(result.result)
                report_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
            except (json.JSONDecodeError, TypeError):
                report = {"_meta": {"status": "completed"}, "raw_output": result.result[:3000]}
                report_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")

        return report

    except Exception as e:
        duration = time.time() - t0
        print(f"  ✗ Failed after {duration:.1f}s: {e}")
        error_report = {
            "_meta": {"status": "failed", "error": str(e), "duration_seconds": duration},
        }
        (phase_dir / "report.json").write_text(json.dumps(error_report, indent=2), encoding="utf-8")
        return error_report


def run_full_pipeline(
    topic: str = "",
    template: str = "academic",
    start_phase: Optional[str] = None,
    max_rounds: int = 3,
):
    """运行完整的 PPT 生成流程"""

    if not os.environ.get("ANTHROPIC_API_KEY"):
        print("ERROR: ANTHROPIC_API_KEY environment variable not set.")
        print("  export ANTHROPIC_API_KEY=sk-ant-...")
        sys.exit(1)

    workspace = create_workspace()
    toolkit = LocalToolkit(workspace)

    # 读取主题
    if not topic:
        topic_path = AGENT_DIR / "TOPIC.md"
        if topic_path.exists():
            topic = topic_path.read_text(encoding="utf-8")
        else:
            print("ERROR: No topic provided. Use --topic or create TOPIC.md")
            sys.exit(1)

    print(f"""
╔══════════════════════════════════════════════════════════╗
║         PPTAgent-67 — 知识驱动 PPT 生成                  ║
║         Model: Claude Opus 4.6 (1M context)              ║
║         Template: {template:<39s} ║
╚══════════════════════════════════════════════════════════╝
  Workspace: {workspace}
  Topic: {topic[:60]}...
  Max rounds: {max_rounds}
""")

    # 写入项目计划
    project_plan = f"""# PPTAgent-67 Project Plan

## Topic
{topic[:500]}

## Template
{template}

## Model
Claude Opus 4.6 (1M context) — claude-opus-4-6

## Phases
{chr(10).join(f'{i+1}. {name} → {agent}' for i, (name, agent) in enumerate(PHASES))}

## Target
Content quality score ≥ 0.7

## Created
{datetime.now(timezone.utc).isoformat()}
"""
    (workspace / "project_plan.md").write_text(project_plan, encoding="utf-8")

    # 保存 TOPIC.md 到工作区
    (workspace / "TOPIC.md").write_text(topic, encoding="utf-8")

    progress = {
        "project_id": workspace.name,
        "topic": topic[:200],
        "template": template,
        "current_round": 1,
        "current_phase": PHASES[0][0],
        "max_rounds": max_rounds,
        "best_score": 0.0,
        "rounds": [],
    }

    for round_num in range(1, max_rounds + 1):
        print(f"\n{'═'*60}")
        print(f"  ROUND {round_num} / {max_rounds}")
        print(f"{'═'*60}")

        progress["current_round"] = round_num
        round_reports = []
        prev_reports = []

        # 确定本轮需要跑哪些阶段
        phases_to_run = PHASES
        if round_num > 1 and "revision_target" in progress:
            target = progress["revision_target"]
            start_idx = next((i for i, (name, _) in enumerate(PHASES) if name == target), 0)
            phases_to_run = PHASES[start_idx:]
            print(f"  Revision: starting from {target}")

        if start_phase and round_num == 1:
            start_idx = next((i for i, (name, _) in enumerate(PHASES) if start_phase in name), 0)
            phases_to_run = PHASES[start_idx:]

        for phase_name, agent_name in phases_to_run:
            progress["current_phase"] = phase_name
            _save_progress(workspace, progress)

            report = run_phase(
                phase_name, agent_name, workspace, round_num, toolkit,
                prev_reports,
                revision_context=progress.get("revision_context", ""),
                topic=topic,
                template=template,
            )

            report_rel = f"round_{round_num}/{phase_name}/report.json"
            prev_reports.append(report_rel)
            round_reports.append({"phase": phase_name, "status": report.get("_meta", {}).get("status", "unknown")})

            # 在精修阶段检查质量分数
            if phase_name == "03_content_refinement":
                quality = report.get("quality_score", 0)
                progress["best_score"] = max(progress.get("best_score", 0), quality)

                print(f"\n  Quality Score: {quality:.2f}")
                print(f"  Best Score:   {progress['best_score']:.2f}")

                if quality >= 0.7:
                    print(f"  ✓ Quality threshold met! Proceeding to render.")
                else:
                    # 决定回退
                    dims = report.get("dimension_scores", {})
                    if dims:
                        lowest = min(dims.items(), key=lambda x: x[1])
                        dimension = lowest[0]
                        target_map = {
                            "structure": "02_outline_generation",
                            "clarity": "03_content_refinement",
                            "depth": "01_knowledge_extraction",
                            "visual": "03_content_refinement",
                            "narrative": "02_outline_generation",
                        }
                        if round_num < max_rounds:
                            progress["revision_target"] = target_map.get(dimension, "03_content_refinement")
                            suggestions = report.get("revision_suggestions", [])
                            progress["revision_context"] = json.dumps(suggestions[:3], ensure_ascii=False)
                            print(f"  Lowest dimension: {dimension} → Will retry from {progress['revision_target']}")
                            # 不进入渲染阶段，直接跳到下一轮
                            break

        progress["rounds"].append({"round": round_num, "phases": round_reports})

        # 检查是否已经渲染完成
        pptx_report = workspace / f"round_{round_num}" / "04_ppt_rendering" / "report.json"
        if pptx_report.exists():
            print(f"\n  ✓ PPT rendered successfully!")
            _finalize(workspace, round_num, toolkit, progress)
            break

    else:
        # 达到 max_rounds，强制渲染当前最佳版本
        print(f"\n  Max rounds ({max_rounds}) reached. Rendering current best version.")
        # 执行最终渲染
        run_phase(
            "04_ppt_rendering", "ppt-renderer", workspace, round_num, toolkit,
            prev_reports, topic=topic, template=template,
        )
        _finalize(workspace, round_num, toolkit, progress)

    _save_progress(workspace, progress)
    print(f"\n{'═'*60}")
    print(f"  Pipeline complete. Workspace: {workspace}")
    print(f"{'═'*60}\n")


def _save_progress(workspace: Path, progress: dict):
    (workspace / "progress.json").write_text(
        json.dumps(progress, ensure_ascii=False, indent=2), encoding="utf-8"
    )


def _finalize(workspace: Path, last_round: int, toolkit: LocalToolkit, progress: dict):
    """归档最终产出到 final/ 目录"""
    final_dir = workspace / "final"
    final_dir.mkdir(exist_ok=True)

    round_dir = workspace / f"round_{last_round}"

    # 复制产出文件
    for src_name, dst_name in [
        ("04_ppt_rendering/artifacts/presentation.pptx", "presentation.pptx"),
        ("04_ppt_rendering/artifacts/presentation.pdf", "presentation.pdf"),
        ("03_content_refinement/artifacts/presentation.md", "presentation.md"),
        ("02_outline_generation/artifacts/outline.json", "outline.json"),
        ("02_outline_generation/artifacts/outline.md", "outline.md"),
    ]:
        src = round_dir / src_name
        if src.exists():
            shutil.copy2(src, final_dir / dst_name)

    # 生成总结
    summary = {
        "topic": progress.get("topic", ""),
        "template_used": progress.get("template", "academic"),
        "total_rounds": last_round,
        "best_quality_score": progress.get("best_score", 0),
        "final_pptx_path": "final/presentation.pptx",
        "workspace": str(workspace),
        "completed_at": datetime.now(timezone.utc).isoformat(),
    }
    (final_dir / "generation_summary.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8"
    )

    progress["final_output"] = str(final_dir)
    print(f"\n  Final output archived to: {final_dir}")


# ═══════════════════════════════════════════════════════════
# 3. CLI
# ═══════════════════════════════════════════════════════════

def main():
    parser = argparse.ArgumentParser(description="PPTAgent-67 Runner")
    parser.add_argument("--topic", type=str, default="",
                        help="PPT topic (or create TOPIC.md)")
    parser.add_argument("--template", type=str, default="academic",
                        choices=["academic", "business", "minimal", "tech"],
                        help="PPT template style (default: academic)")
    parser.add_argument("--phase", type=str, default=None,
                        help="Start from a specific phase (e.g., outline, rendering)")
    parser.add_argument("--max-rounds", type=int, default=3,
                        help="Maximum iteration rounds (default: 3)")
    parser.add_argument("--test-tools", action="store_true",
                        help="Test local toolkit connectivity and exit")
    args = parser.parse_args()

    if args.test_tools:
        _test_tools()
        return

    run_full_pipeline(
        topic=args.topic,
        template=args.template,
        start_phase=args.phase,
        max_rounds=args.max_rounds,
    )


def _test_tools():
    """测试本地工具是否可用"""
    ws = Path("/tmp/pptagent67_test")
    ws.mkdir(exist_ok=True)
    toolkit = LocalToolkit(ws)

    tests = [
        ("python_exec", lambda: toolkit.python_exec("import sys; print(f'Python {sys.version}')")),
        ("marp-cli", lambda: toolkit.python_exec("import subprocess; r=subprocess.run(['marp','--version'],capture_output=True,text=True); print(r.stdout.strip() or 'not found')")),
        ("write_file", lambda: toolkit.write_file(json.dumps({"path": "test.txt", "content": "hello"}))),
        ("read_file", lambda: toolkit.read_file("test.txt")),
        ("list_dir", lambda: toolkit.list_dir(".")),
    ]

    print("Testing local toolkit...\n")
    for name, fn in tests:
        try:
            t0 = time.time()
            result = fn()
            dt = time.time() - t0
            ok = "[ERROR]" not in result and "[STUB]" not in result
            status = "✓" if ok else "⚠"
            print(f"  {status} {name:30s} ({dt:.1f}s) {result[:80]}")
        except Exception as e:
            print(f"  ✗ {name:30s} ERROR: {e}")

    shutil.rmtree(ws, ignore_errors=True)
    print("\nDone.")


if __name__ == "__main__":
    main()