#!/usr/bin/env python3
"""Generate TradingAgents workflow diagrams for the analysis document."""
import sys, os
sys.path.insert(0, '/home/ubuntu/.nanobot/workspace/skills/hand-drawn-diagram')
from generator import DiagramGenerator, Component, Connection, Container

DOCS = '/home/ubuntu/nanocode/TradingAgents/docs'

# ──────────────────────────────────────────────────────────────────
# Diagram 1: High-level pipeline flowchart
# ──────────────────────────────────────────────────────────────────
gen1 = DiagramGenerator(style="sketch")

steps = [
    {"type": "start",     "label": "用户输入\n股票代码 + 日期"},
    {"type": "process",   "label": "并行数据采集\n(4个分析师同时工作)"},
    {"type": "process",   "label": "研究经理汇总\n收集4份分析报告"},
    {"type": "decision",  "label": "多空辩论\nBull vs Bear × N轮"},
    {"type": "process",   "label": "交易员决策\n综合辩论形成草案"},
    {"type": "decision",  "label": "三级风控评审\n激进/保守/中立"},
    {"type": "end",       "label": "最终输出\nBuy / Sell / Hold"},
]

gen1.create_flowchart(
    title="TradingAgents 单次股票分析流程",
    steps=steps,
    output=f"{DOCS}/workflow-pipeline.html",
)

print(f"✅ 1. Pipeline flowchart: {DOCS}/workflow-pipeline.html")

# ──────────────────────────────────────────────────────────────────
# Diagram 2: Detailed architecture with data flows
# ──────────────────────────────────────────────────────────────────
W, H = 1300, 1200
gen2 = DiagramGenerator(style="sketch")

# --- Component definitions ---
center = W // 2

# Row 1: Input layer
user = Component("用户浏览器", "cloud", center-70, 20, 140, 45, "cloud")

# Row 2: Web server
web_srv = Component("Web Server (FastAPI)\n后台异步分析任务", "frontend", center-95, 95, 190, 55)

# Row 3: LangGraph
graph = Component("TradingAgents Graph\n(LangGraph StateGraph)", "backend", center-105, 185, 210, 55, "hexagon")

# Row 4: Four analysts + manager (parallel)
mkt = Component("市场分析师\n行情+技术指标", "backend", 50, 310, 150, 60)
snt = Component("社媒分析师\n舆论情绪", "backend", 240, 310, 140, 60)
nws = Component("新闻分析师\n个股+宏观新闻", "backend", 420, 310, 150, 60)
fnd = Component("基本面分析师\n三表+估值", "backend", 610, 310, 150, 60)
mgr = Component("研究经理\n汇总4份报告", "security", 810, 310, 140, 60)

# Row 5: Debate agents
bull = Component("Bull 多头\n研究员", "frontend", 230, 440, 140, 50)
bear = Component("Bear 空头\n研究员", "security", 420, 440, 140, 50)
dbj = Component("辩论裁判\n评估论据", "backend", 610, 440, 140, 50)

# Row 6: Trader
trader = Component("Trader 交易员\n综合辩论→决策草案", "backend", center-105, 560, 210, 55, "hexagon")

# Row 7: Risk management trio
agg = Component("Aggressive\n激进风控", "frontend", 190, 680, 140, 55)
con = Component("Conservative\n保守风控", "security", 390, 680, 145, 55)
neu = Component("Neutral\n中立风控", "backend", 590, 680, 140, 55)
rsk = Component("Risk Judge\n风控终审", "security", 800, 680, 140, 55)

# Row 8: Final output
output = Component("输出结果\n分析报告 + 交易信号", "cloud", center-110, 800, 220, 55, "cloud")

# Row 9: Data sources
ak = Component("akshare 主数据源\nA股/港股/美股/新闻", "database", 80, 930, 185, 55)
yf = Component("yfinance 备用1\n美股+内部交易", "database", 320, 930, 165, 55)
av = Component("Alpha Vantage 备用2\n美国市场备选", "database", 540, 930, 165, 55)
llm = Component("DeepSeek LLM\n推理引擎", "database", 780, 930, 155, 55)

components = [user, web_srv, graph,
              mkt, snt, nws, fnd, mgr,
              bull, bear, dbj,
              trader,
              agg, con, neu, rsk,
              output,
              ak, yf, av, llm]

# --- Connections ---
connections = [
    Connection("用户浏览器", "Web Server (FastAPI)\n后台异步分析任务", "solid", "POST /api/analyze", "#22d3ee"),
    Connection("Web Server (FastAPI)\n后台异步分析任务", "TradingAgents Graph\n(LangGraph StateGraph)", "solid", "启动分析", "#22d3ee"),
    # Graph → 4 analysts
    Connection("TradingAgents Graph\n(LangGraph StateGraph)", "市场分析师\n行情+技术指标", "dashed", "分发", "#fbbf24"),
    Connection("TradingAgents Graph\n(LangGraph StateGraph)", "社媒分析师\n舆论情绪", "dashed", "分发", "#fbbf24"),
    Connection("TradingAgents Graph\n(LangGraph StateGraph)", "新闻分析师\n个股+宏观新闻", "dashed", "分发", "#fbbf24"),
    Connection("TradingAgents Graph\n(LangGraph StateGraph)", "基本面分析师\n三表+估值", "dashed", "分发", "#fbbf24"),
    # Analysts → Manager
    Connection("市场分析师\n行情+技术指标", "研究经理\n汇总4份报告", "solid", "报告", "#34d399"),
    Connection("社媒分析师\n舆论情绪", "研究经理\n汇总4份报告", "solid", "报告", "#34d399"),
    Connection("新闻分析师\n个股+宏观新闻", "研究经理\n汇总4份报告", "solid", "报告", "#34d399"),
    Connection("基本面分析师\n三表+估值", "研究经理\n汇总4份报告", "solid", "报告", "#34d399"),
    # Manager → Debate
    Connection("研究经理\n汇总4份报告", "Bull 多头\n研究员", "solid", "启动辩论", "#a78bfa"),
    Connection("研究经理\n汇总4份报告", "Bear 空头\n研究员", "solid", "启动辩论", "#a78bfa"),
    Connection("Bull 多头\n研究员", "辩论裁判\n评估论据", "solid", "多方论据", "#f87171"),
    Connection("Bear 空头\n研究员", "辩论裁判\n评估论据", "solid", "空方论据", "#f87171"),
    # Judge → Trader
    Connection("辩论裁判\n评估论据", "Trader 交易员\n综合辩论→决策草案", "solid", "辩论摘要", "#fb923c"),
    # Trader → Risk
    Connection("Trader 交易员\n综合辩论→决策草案", "Aggressive\n激进风控", "dashed", "草案", "#38bdf8"),
    Connection("Trader 交易员\n综合辩论→决策草案", "Conservative\n保守风控", "dashed", "草案", "#38bdf8"),
    Connection("Trader 交易员\n综合辩论→决策草案", "Neutral\n中立风控", "dashed", "草案", "#38bdf8"),
    # Risk → Judge
    Connection("Aggressive\n激进风控", "Risk Judge\n风控终审", "solid", "评估", "#34d399"),
    Connection("Conservative\n保守风控", "Risk Judge\n风控终审", "solid", "评估", "#34d399"),
    Connection("Neutral\n中立风控", "Risk Judge\n风控终审", "solid", "评估", "#34d399"),
    # Final
    Connection("Risk Judge\n风控终审", "输出结果\n分析报告 + 交易信号", "solid", "最终决策", "#22d3ee"),
    # Data source connections
    Connection("市场分析师\n行情+技术指标", "akshare 主数据源\nA股/港股/美股/新闻", "dashed", "数据", "#94a3b8"),
    Connection("基本面分析师\n三表+估值", "akshare 主数据源\nA股/港股/美股/新闻", "dashed", "数据", "#94a3b8"),
    Connection("新闻分析师\n个股+宏观新闻", "akshare 主数据源\nA股/港股/美股/新闻", "dashed", "数据", "#94a3b8"),
    Connection("akshare 主数据源\nA股/港股/美股/新闻", "yfinance 备用1\n美股+内部交易", "dashed", "fallback", "#64748b"),
    Connection("yfinance 备用1\n美股+内部交易", "Alpha Vantage 备用2\n美国市场备选", "dashed", "fallback", "#64748b"),
    Connection("Bull 多头\n研究员", "DeepSeek LLM\n推理引擎", "dashed", "推理", "#94a3b8"),
    Connection("Bear 空头\n研究员", "DeepSeek LLM\n推理引擎", "dashed", "推理", "#94a3b8"),
]

# --- Container groups ---
containers = [
    Container("分析师团队 (并行)", 20, 285, 960, 110, "#fbbf24"),
    Container("多空辩论 (可多轮)", 200, 415, 580, 100, "#a78bfa"),
    Container("三级风控", 150, 655, 820, 105, "#f87171"),
    Container("数据源层", 50, 905, 950, 105, "#64748b"),
]

# --- Generate ---
gen2.create_architecture(
    title="TradingAgents 详细架构与数据流",
    components=components,
    connections=connections,
    output=f"{DOCS}/workflow-architecture",
    containers=containers,
    auto_layout=False,
)

print(f"✅ 2. Architecture diagram: {DOCS}/workflow-architecture.html")
print("\nDone!")