#!/usr/bin/env python3
"""
生成包含多个节点模板的复杂Agent
"""
import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from sqlalchemy.orm import Session
from app.core.database import SessionLocal
from app.models.agent import Agent
from app.models.user import User
from app.models.node_template import NodeTemplate
from datetime import datetime
import uuid


def generate_complex_agent_with_templates(db: Session, username: str = "admin"):
    """生成包含多个节点模板的复杂Agent"""
    print("=" * 60)
    print("生成复杂Agent（包含多个节点模板）")
    print("=" * 60)
    print()
    
    # 查找用户
    user = db.query(User).filter(User.username == username).first()
    if not user:
        print(f"❌ 未找到用户 '{username}'，请先创建该用户")
        return
    print(f"✅ 找到用户: {user.username} (ID: {user.id})")
    print()
    
    # 查找可用的节点模板（公开的或用户自己的）
    templates = db.query(NodeTemplate).filter(
        (NodeTemplate.is_public == True) | (NodeTemplate.user_id == user.id)
    ).limit(10).all()
    
    if len(templates) < 3:
        print(f"⚠️  警告: 只找到 {len(templates)} 个节点模板，建议至少3个")
        if len(templates) == 0:
            print("❌ 没有可用的节点模板，请先创建节点模板")
            return
    
    print(f"📋 找到 {len(templates)} 个可用节点模板:")
    for i, template in enumerate(templates[:5], 1):
        print(f"   {i}. {template.name} (ID: {template.id})")
    print()
    
    # 选择3-5个模板用于创建复杂Agent
    selected_templates = templates[:min(5, len(templates))]
    template_ids = [t.id for t in selected_templates]
    
    # 生成复杂的工作流配置
    # 工作流结构：开始 -> 模板1 -> 条件判断 -> 模板2/模板3 -> 转换 -> 模板4 -> 结束
    nodes = []
    edges = []
    
    # 开始节点
    start_node = {
        "id": "start-1",
        "type": "start",
        "position": {"x": 100, "y": 200},
        "data": {"label": "开始"}
    }
    nodes.append(start_node)
    
    # 第一个模板节点
    if len(template_ids) > 0:
        template1_node = {
            "id": "template-1",
            "type": "template",
            "position": {"x": 300, "y": 200},
            "data": {
                "label": selected_templates[0].name,
                "template_id": template_ids[0],
                "provider": selected_templates[0].provider,
                "model": selected_templates[0].model,
                "temperature": selected_templates[0].temperature,
                "max_tokens": selected_templates[0].max_tokens,
                "prompt": selected_templates[0].prompt
            }
        }
        nodes.append(template1_node)
        edges.append({
            "id": "e1",
            "source": "start-1",
            "target": "template-1"
        })
    
    # 条件判断节点
    condition_node = {
        "id": "condition-1",
        "type": "condition",
        "position": {"x": 500, "y": 200},
        "data": {
            "label": "判断处理结果",
            "condition": "{{result}} contains '需要' or {{result}} contains '重要'"
        }
    }
    nodes.append(condition_node)
    if len(template_ids) > 0:
        edges.append({
            "id": "e2",
            "source": "template-1",
            "target": "condition-1"
        })
    
    # 条件分支：如果满足条件，使用模板2；否则使用模板3
    if len(template_ids) > 1:
        template2_node = {
            "id": "template-2",
            "type": "template",
            "position": {"x": 700, "y": 100},
            "data": {
                "label": selected_templates[1].name,
                "template_id": template_ids[1],
                "provider": selected_templates[1].provider,
                "model": selected_templates[1].model,
                "temperature": selected_templates[1].temperature,
                "max_tokens": selected_templates[1].max_tokens,
                "prompt": selected_templates[1].prompt
            }
        }
        nodes.append(template2_node)
        edges.append({
            "id": "e3",
            "source": "condition-1",
            "target": "template-2",
            "sourceHandle": "true"
        })
    
    if len(template_ids) > 2:
        template3_node = {
            "id": "template-3",
            "type": "template",
            "position": {"x": 700, "y": 300},
            "data": {
                "label": selected_templates[2].name,
                "template_id": template_ids[2],
                "provider": selected_templates[2].provider,
                "model": selected_templates[2].model,
                "temperature": selected_templates[2].temperature,
                "max_tokens": selected_templates[2].max_tokens,
                "prompt": selected_templates[2].prompt
            }
        }
        nodes.append(template3_node)
        edges.append({
            "id": "e4",
            "source": "condition-1",
            "target": "template-3",
            "sourceHandle": "false"
        })
    
    # 数据转换节点
    transform_node = {
        "id": "transform-1",
        "type": "transform",
        "position": {"x": 900, "y": 200},
        "data": {
            "label": "数据转换",
            "mode": "mapping",
            "mapping": {
                "final_result": "{{result}}",
                "processed": "true"
            }
        }
    }
    nodes.append(transform_node)
    
    # 连接条件分支到转换节点
    if len(template_ids) > 1:
        edges.append({
            "id": "e5",
            "source": "template-2",
            "target": "transform-1"
        })
    if len(template_ids) > 2:
        edges.append({
            "id": "e6",
            "source": "template-3",
            "target": "transform-1"
        })
    
    # 最后一个模板节点（如果还有模板）
    if len(template_ids) > 3:
        template4_node = {
            "id": "template-4",
            "type": "template",
            "position": {"x": 1100, "y": 200},
            "data": {
                "label": selected_templates[3].name,
                "template_id": template_ids[3],
                "provider": selected_templates[3].provider,
                "model": selected_templates[3].model,
                "temperature": selected_templates[3].temperature,
                "max_tokens": selected_templates[3].max_tokens,
                "prompt": selected_templates[3].prompt
            }
        }
        nodes.append(template4_node)
        edges.append({
            "id": "e7",
            "source": "transform-1",
            "target": "template-4"
        })
        # 结束节点连接到最后一个模板
        end_node = {
            "id": "end-1",
            "type": "end",
            "position": {"x": 1300, "y": 200},
            "data": {"label": "结束"}
        }
        nodes.append(end_node)
        edges.append({
            "id": "e8",
            "source": "template-4",
            "target": "end-1"
        })
    else:
        # 如果没有更多模板，转换节点直接连接到结束节点
        end_node = {
            "id": "end-1",
            "type": "end",
            "position": {"x": 1100, "y": 200},
            "data": {"label": "结束"}
        }
        nodes.append(end_node)
        edges.append({
            "id": "e7",
            "source": "transform-1",
            "target": "end-1"
        })
    
    # 检查是否已存在同名Agent
    agent_name = "复杂模板Agent（多节点）"
    existing = db.query(Agent).filter(
        Agent.name == agent_name,
        Agent.user_id == user.id
    ).first()
    
    if existing:
        print(f"⚠️  Agent '{agent_name}' 已存在，将更新它...")
        existing.workflow_config = {"nodes": nodes, "edges": edges}
        existing.description = f"包含 {len(selected_templates)} 个节点模板的复杂Agent，支持条件分支和数据转换"
        existing.updated_at = datetime.now()
        agent = existing
    else:
        # 创建Agent
        agent = Agent(
            name=agent_name,
            description=f"包含 {len(selected_templates)} 个节点模板的复杂Agent，支持条件分支和数据转换",
            workflow_config={"nodes": nodes, "edges": edges},
            status="published",
            user_id=user.id,
            version=1
        )
        db.add(agent)
    
    try:
        db.commit()
        db.refresh(agent)
        print()
        print("=" * 60)
        print(f"✅ 成功创建/更新复杂Agent: {agent.name}")
        print(f"   ID: {agent.id}")
        print(f"   状态: {agent.status}")
        print(f"   节点数量: {len(nodes)}")
        print(f"   连接数量: {len(edges)}")
        print(f"   使用的模板节点: {len(selected_templates)}")
        print("=" * 60)
        print()
        print("工作流结构:")
        print("  开始 -> 模板1 -> 条件判断 -> 模板2/模板3 -> 数据转换 -> 模板4 -> 结束")
        print()
        return agent
    except Exception as e:
        db.rollback()
        print()
        print("=" * 60)
        print(f"❌ 创建失败: {e}")
        print("=" * 60)
        import traceback
        traceback.print_exc()
        raise


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="生成包含多个节点模板的复杂Agent")
    parser.add_argument(
        "--username",
        type=str,
        default="admin",
        help="创建Agent的用户名（默认: admin）"
    )
    
    args = parser.parse_args()
    
    db = SessionLocal()
    try:
        generate_complex_agent_with_templates(db, username=args.username)
    except Exception as e:
        print(f"❌ 发生错误: {e}")
        import traceback
        traceback.print_exc()
    finally:
        db.close()
