# skills/subagent-manager/subagent_manager.py

from typing import Dict, List, Optional
from dataclasses import dataclass
from enum import Enum
import asyncio

class AgentRole(Enum):
    """Agent 角色枚举"""
    PM = "pm"
    ARCHITECT = "architect"
    DEVELOPER = "developer"
    TESTER = "tester"
    WRITER = "writer"
    DEVOPS = "devops"

@dataclass
class AgentConfig:
    """Agent 配置"""
    role: AgentRole
    label: str
    model: str = "qwen3.5-plus"
    timeout: int = 1800  # 30 分钟
    workspace: str = ""
    read_only: List[str] = None
    thread: bool = True
    mode: str = "session"

class SubAgentManager:
    """子 Agent 管理器"""
    
    def __init__(self):
        self.agents: Dict[str, AgentConfig] = {}
        self.sessions: Dict[str, str] = {}  # role -> session_key
        
    def spawn(self, config: AgentConfig, task: str) -> str:
        """
        生成子 Agent
        
        Args:
            config: Agent 配置
            task: 任务描述
            
        Returns:
            session_key: 会话密钥
        """
        # 1. 构建工作区配置
        workspace = config.workspace or f"private/{config.role.value}/"
        
        # 2. 构建只读资源列表
        read_only = config.read_only or [
            "memory/MEMORY.md",
            "skills/",
            "docs/",
            "projects/"
        ]
        
        # 3. 构建任务描述
        full_task = self._build_task_description(config.role, task)
        
        # 4. 调用 sessions_spawn
        session_key = sessions_spawn(
            label=config.label,
            runtime="subagent",
            mode=config.mode,
            task=full_task,
            model=config.model,
            thread=config.thread,
            timeout_seconds=config.timeout,
            cleanup="keep"
        )
        
        # 5. 记录 Agent 信息
        self.agents[session_key] = config
        self.sessions[config.role.value] = session_key
        
        return session_key
    
    def _build_task_description(self, role: AgentRole, task: str) -> str:
        """构建任务描述"""
        
        role_instructions = {
            AgentRole.PM: """
【PM 角色】产品需求分析

**工作规则**：
1. 工作文件：private/pm/work/
2. 私有笔记：private/pm/notes/
3. 提交输出：private/pm/output/
4. 读取共享：memory/, skills/, docs/, projects/

**输出要求**：
- 需求文档（PRD）
- 用户故事（User Stories）
- 任务清单与优先级
""",
            AgentRole.ARCHITECT: """
【Architect 角色】系统架构设计

**工作规则**：
1. 工作文件：private/architect/work/
2. 私有笔记：private/architect/notes/
3. 提交输出：private/architect/output/
4. 读取共享：memory/, skills/, docs/, projects/

**输出要求**：
- 架构图（Mermaid/draw.io）
- 技术选型文档
- API 设计规范
""",
            AgentRole.DEVELOPER: """
【Developer 角色】编码实现

**工作规则**：
1. 工作文件：private/developer/work/
2. 私有笔记：private/developer/notes/
3. 提交输出：private/developer/output/
4. 读取共享：memory/, skills/, docs/, projects/

**输出要求**：
- 源代码（含注释）
- 单元测试
- 代码审查通过
""",
            AgentRole.TESTER: """
【Tester 角色】测试验证

**工作规则**：
1. 工作文件：private/tester/work/
2. 私有笔记：private/tester/notes/
3. 提交输出：private/tester/output/
4. 读取共享：memory/, skills/, docs/, projects/

**输出要求**：
- 测试用例
- 测试报告
- 缺陷报告
""",
            AgentRole.WRITER: """
【Writer 角色】文档编写

**工作规则**：
1. 工作文件：private/writer/work/
2. 私有笔记：private/writer/notes/
3. 提交输出：private/writer/output/
4. 读取共享：memory/, skills/, docs/, projects/

**输出要求**：
- 技术文档
- API 文档
- 用户手册
""",
            AgentRole.DEVOPS: """
【DevOps 角色】部署配置

**工作规则**：
1. 工作文件：private/devops/work/
2. 私有笔记：private/devops/notes/
3. 提交输出：private/devops/output/
4. 读取共享：memory/, skills/, docs/, projects/

**输出要求**：
- Dockerfile
- Jenkinsfile
- K8s 配置
- 监控仪表盘
"""
        }
        
        base_instruction = role_instructions.get(role, "")
        
        return f"""{base_instruction}

**当前任务**：
{task}

**完成标准**：
- 质量符合项目标准
- 输出到私有 output/ 目录等待审核
- 在共享 memory/ 中记录关键决策
"""
    
    async def wait_all(self, timeout: int = 3600) -> Dict[str, str]:
        """等待所有子 Agent 完成"""
        
        results = {}
        
        for role, session_key in self.sessions.items():
            try:
                # 轮询会话状态
                result = await self._poll_session(session_key, timeout)
                results[role] = result
            except asyncio.TimeoutError:
                results[role] = f"❌ 超时（{timeout}s）"
        
        return results
    
    async def _poll_session(self, session_key: str, timeout: int) -> str:
        """轮询会话完成"""
        
        start_time = time.time()
        
        while time.time() - start_time < timeout:
            status = sessions_list(active_minutes=1)
            
            for session in status:
                if session['key'] == session_key:
                    if session['status'] == 'completed':
                        return session['result']
                    elif session['status'] == 'failed':
                        return f"❌ 失败：{session['error']}"
            
            await asyncio.sleep(10)  # 每 10 秒检查一次
        
        raise asyncio.TimeoutError(f"等待超时：{timeout}s")
    
    def get_outputs(self, role: str) -> List[str]:
        """获取子 Agent 输出"""
        
        output_dir = Path(f"private/{role}/output/")
        
        if not output_dir.exists():
            return []
        
        return [
            str(f) for f in output_dir.glob("*")
            if f.is_file()
        ]
    
    def review_output(self, role: str, output_path: str) -> ReviewResult:
        """审核子 Agent 输出"""
        
        # 1. 读取输出文件
        content = Path(output_path).read_text()
        
        # 2. 质量检查清单
        checklist = self._get_review_checklist(role)
        
        # 3. 逐项检查
        results = {}
        for item in checklist:
            results[item] = self._check_item(content, item)
        
        # 4. 生成审核结果
        passed = all(results.values())
        
        return ReviewResult(
            passed=passed,
            checklist=results,
            suggestions=self._generate_suggestions(results)
        )
    
    def _get_review_checklist(self, role: str) -> List[str]:
        """获取审核清单"""
        
        checklists = {
            "pm": [
                "需求描述清晰",
                "用户故事完整",
                "优先级合理",
                "验收标准明确"
            ],
            "architect": [
                "架构图清晰",
                "技术选型合理",
                "接口设计完整",
                "性能/安全考虑"
            ],
            "developer": [
                "代码规范",
                "单元测试覆盖",
                "注释完整",
                "无安全漏洞"
            ],
            "tester": [
                "测试用例完整",
                "边界条件覆盖",
                "自动化程度",
                "缺陷描述清晰"
            ],
            "writer": [
                "文档结构清晰",
                "内容准确",
                "示例完整",
                "格式规范"
            ],
            "devops": [
                "配置完整",
                "安全性考虑",
                "监控告警",
                "回滚方案"
            ]
        }
        
        return checklists.get(role, [])
    
    def merge_to_shared(self, role: str, output_path: str, target_dir: str):
        """合并输出到共享区"""
        
        # 1. 读取输出
        content = Path(output_path).read_text()
        
        # 2. 确定目标位置
        target_path = Path(target_dir) / Path(output_path).name
        
        # 3. 备份现有文件（如果存在）
        if target_path.exists():
            backup_path = target_path.with_suffix(target_path.suffix + '.backup')
            shutil.copy(target_path, backup_path)
        
        # 4. 复制到共享区
        shutil.copy(output_path, target_path)
        
        # 5. 记录合并日志
        self._log_merge(role, output_path, target_path)

async def distribute_task(main_task: str):
    """分发任务到子 Agent"""
    
    manager = SubAgentManager()
    
    # 1. PM 分析需求
    pm_session = manager.spawn(
        AgentConfig(
            role=AgentRole.PM,
            label="pm-analyst",
            timeout=1800
        ),
        task=f"分析以下需求，输出 PRD 和用户故事：\n\n{main_task}"
    )
    
    # 等待 PM 完成
    pm_result = await manager.wait_for(pm_session)
    
    # 2. Architect 设计架构
    arch_session = manager.spawn(
        AgentConfig(
            role=AgentRole.ARCHITECT,
            label="arch-designer",
            timeout=2700
        ),
        task=f"基于以下 PRD 设计系统架构：\n\n{pm_result}"
    )
    
    # 3. Developer 编码实现
    dev_session = manager.spawn(
        AgentConfig(
            role=AgentRole.DEVELOPER,
            label="dev-coder",
            timeout=5400
        ),
        task=f"基于以下架构设计实现代码：\n\n{arch_result}"
    )
    
    # 4. Tester 测试验证
    test_session = manager.spawn(
        AgentConfig(
            role=AgentRole.TESTER,
            label="tester-qa",
            timeout=2700
        ),
        task=f"为以下代码设计测试用例并验证：\n\n{dev_result}"
    )
    
    # 5. Writer 编写文档
    write_session = manager.spawn(
        AgentConfig(
            role=AgentRole.WRITER,
            label="writer-docs",
            timeout=1800
        ),
        task=f"为以下功能编写技术文档：\n\n{dev_result}"
    )
    
    # 6. DevOps 部署配置
    ops_session = manager.spawn(
        AgentConfig(
            role=AgentRole.DEVOPS,
            label="devops-deploy",
            timeout=1800
        ),
        task=f"为以下功能配置 CI/CD 和部署：\n\n{dev_result}"
    )
    
    # 等待所有子 Agent 完成
    all_results = await manager.wait_all()
    
    return all_results

def review_all_outputs(manager: SubAgentManager) -> ReviewReport:
    """审核所有子 Agent 输出"""
    
    report = ReviewReport()
    
    for role in AgentRole:
        # 获取输出文件
        outputs = manager.get_outputs(role.value)
        
        for output_path in outputs:
            # 审核
            result = manager.review_output(role.value, output_path)
            
            if result.passed:
                # 审核通过，合并到共享区
                target_dir = get_target_directory(role)
                manager.merge_to_shared(role.value, output_path, target_dir)
                report.passed.append(output_path)
            else:
                # 审核不通过，返回修改
                report.failed.append({
                    'file': output_path,
                    'role': role.value,
                    'issues': result.checklist,
                    'suggestions': result.suggestions
                })
    
    return report

# Main Agent 分解任务

pm_task = """
分析 CrystalForge v2.3.0 需求：

**背景**：
基于 v2.2.0 盈利模式优化，需要实现：
1. 代际继承版税（Lineage Royalties）
2. 算力众筹池（Compute Crowdfunding Pool）
3. Crystal Pass 预售系统

**输出**：
1. PRD 文档
2. 用户故事（每个功能至少 5 个）
3. 优先级矩阵（MoSCoW）
4. 验收标准
"""

arch_task = """
设计 CrystalForge v2.3.0 架构：

**输入**：PM 输出的 PRD

**要求**：
1. 系统架构图（Mermaid）
2. 数据库设计（ER 图）
3. API 接口设计（OpenAPI）
4. 性能设计（目标：版税计算 < 100ms）
5. 安全设计（防欺诈）
"""

dev_task = """
实现 CrystalForge v2.3.0：

**输入**：架构设计文档

**要求**：
1. 后端：Spring Boot 3.2 + Java 17
2. 前端：Vue 3.4 + Vite 5.1
3. 单元测试覆盖率 > 90%
4. 代码审查通过
"""

test_task = """
测试 CrystalForge v2.3.0：

**输入**：实现代码

**要求**：
1. 测试用例设计（覆盖所有用户故事）
2. 自动化测试脚本
3. 性能测试（500 并发）
4. 测试报告
"""

writer_task = """
编写 CrystalForge v2.3.0 文档：

**输入**：所有输出物

**要求**：
1. 技术文档（架构/实现）
2. API 文档（Swagger）
3. 用户手册
4. 发布说明（Release Notes）
"""

ops_task = """
配置 CrystalForge v2.3.0 部署：

**输入**：实现代码

**要求**：
1. Dockerfile 更新
2. Jenkins Pipeline 配置
3. K8s 部署文件
4. 监控仪表盘
"""

# Main Agent 分配任务

# Writer 1：P1 文章（架构师级别）
spawn(
    label="writer-p1",
    runtime="subagent",
    task="""
【Writer - P1 文章创作】

**任务**：创作 10 篇 P1 级别 OpenClaw 文章

**主题**：
1. OpenClaw K8s 部署实践
2. AI Agent 记忆系统设计
3. Feishu 机器人集成指南
4. OpenClaw 技能开发指南
5. OpenClaw 子 Agent 管理实战
6. OpenClaw 配置文件详解
7. OpenClaw 浏览器自动化
8. OpenClaw 记忆系统优化
9. OpenClaw 性能调优
10. OpenClaw 安全加固

**质量标准**：
- 每篇 20KB+
- 5 个 +Mermaid 架构图
- 3 个 + 实战案例
- 完整代码示例
- 踩坑记录

**输出**：private/writer/output/p1-articles/
"""
)

# Writer 2：P2 文章（高级开发）
spawn(
    label="writer-p2",
    runtime="subagent",
    task="""
【Writer - P2 文章创作】

**任务**：创作 20 篇 P2 级别技术文章

**主题方向**：
- Spring Boot 最佳实践
- Vue 3 高级技巧
- 微服务架构
- 数据库优化
- 缓存设计
- 消息队列
- API 设计
- 测试策略

**质量标准**：
- 每篇 15KB+
- 3 个 + 架构图
- 2 个 + 实战案例
- 代码示例

**输出**：private/writer/output/p2-articles/
"""
)

# Writer 3：P3 文章（技术教程）
spawn(
    label="writer-p3",
    runtime="subagent",
    task="""
【Writer - P3 文章创作】

**任务**：创作 16 篇 P3 级别技术教程

**主题方向**：
- 工具使用教程
- 环境搭建指南
- 快速入门
- 常见问题解答
- 最佳实践

**质量标准**：
- 每篇 10KB+
- 步骤清晰
- 截图完整
- 可复现

**输出**：private/writer/output/p3-articles/
"""
)

# Reviewer：质量审核
spawn(
    label="reviewer",
    runtime="subagent",
    task="""
【Reviewer - 文章质量审核】

**任务**：审核所有文章质量

**审核清单**：
- [ ] 字数达标
- [ ] 架构图完整
- [ ] 案例真实
- [ ] 代码可运行
- [ ] 格式规范
- [ ] 无错别字

**输出**：private/reviewer/output/review-report.md
"""
)

async def spawn_parallel(agents: List[AgentConfig], task: str):
    """并行生成多个子 Agent"""
    
    # 使用信号量限制并发数
    semaphore = asyncio.Semaphore(5)  # 最多 5 个并发
    
    async def spawn_with_semaphore(config):
        async with semaphore:
            return await spawn_async(config, task)
    
    # 并行生成
    tasks = [spawn_with_semaphore(config) for config in agents]
    sessions = await asyncio.gather(*tasks)
    
    return sessions

def get_workspace_config(role: str) -> dict:
    """获取工作区配置"""
    
    return {
        "read_only": [
            "memory/MEMORY.md",
            "memory/skill-usage-log.md",
            "skills/",
            "docs/",
            "projects/",
            "AGENTS.md",
            "SOUL.md"
        ],
        "write_only": [
            f"private/{role}/work/",
            f"private/{role}/notes/",
            f"private/{role}/output/"
        ],
        "shared_write": [
            "memory/YYYY-MM-DD.md"
        ]
    }

class AgentCommunication:
    """Agent 间通信"""
    
    def __init__(self):
        self.message_queue = asyncio.Queue()
        self.subscribers = defaultdict(list)
    
    def subscribe(self, topic: str, callback):
        """订阅主题"""
        self.subscribers[topic].append(callback)
    
    def publish(self, topic: str, message: dict):
        """发布消息"""
        for callback in self.subscribers[topic]:
            asyncio.create_task(callback(message))
    
    async def send_message(self, from_agent: str, to_agent: str, content: str):
        """发送消息"""
        await self.message_queue.put({
            'from': from_agent,
            'to': to_agent,
            'content': content,
            'timestamp': datetime.now()
        })
    
    async def receive_message(self, agent: str, timeout: int = 60):
        """接收消息"""
        try:
            message = await asyncio.wait_for(
                self.message_queue.get(),
                timeout=timeout
            )
            return message
        except asyncio.TimeoutError:
            return None

# 明确依赖关系
task_dependencies = {
    "developer": ["architect"],
    "tester": ["developer"],
    "writer": ["developer", "architect"],
    "devops": ["developer"]
}

async def check_dependencies(role: str) -> bool:
    """检查依赖是否完成"""
    deps = task_dependencies.get(role, [])
    for dep in deps:
        if not is_completed(dep):
            return False
    return True

def select_model(task_type: str, channel: str, user_preference: str = None):
    """模型选择策略"""
    
    # 1. 用户显式指定
    if user_preference:
        return user_preference
    
    # 2. 按任务类型
    if task_type in config['models']['byTask']:
        return config['models']['byTask'][task_type]
    
    # 3. 按渠道
    if channel in config['models']['byChannel']:
        return config['models']['byChannel'][channel]
    
    # 4. 默认模型
    return config['models']['default']