# agent.py
import json
import ollama
from dataclasses import dataclass, field
from typing import Optional


@dataclass
class Message:
    role: str       # "user" | "agent" | "system"
    sender: str     # 發(fā)送者名稱
    content: str


@dataclass
class Agent:
    name: str
    role: str
    system_prompt: str
    model: str = "qwen3:8b"
    history: list = field(default_factory=list)

    def think(self, messages: list[Message]) -> str:
        """思考并回復(fù)"""
        # 構(gòu)建完整的消息列表
        full_messages = [
            {"role": "system", "content": self.system_prompt}
        ]
        for msg in messages:
            full_messages.append({
                "role": "user" if msg.sender != self.name else "assistant",
                "content": f"[來自 {msg.sender}] {msg.content}"
            })

        # 調(diào)用模型
        response = ollama.chat(
            model=self.model,
            messages=full_messages
        )
        return response["message"]["content"]

# agents.py
from agent import Agent

# 研究員 Agent：負(fù)責(zé)技術(shù)調(diào)研
researcher = Agent(
    name="Researcher",
    role="技術(shù)研究員",
    system_prompt="""你是一個(gè)技術(shù)研究員。你的任務(wù)是：
1. 分析用戶需求中涉及的技術(shù)點(diǎn)
2. 給出最佳實(shí)踐和技術(shù)方案建議
3. 列出關(guān)鍵注意事項(xiàng)和潛在風(fēng)險(xiǎn)

輸出格式：
## 技術(shù)分析
（分析內(nèi)容）

## 推薦方案
（方案詳情）

## 注意事項(xiàng)
（風(fēng)險(xiǎn)和注意點(diǎn)）"""
)

# 程序員 Agent：負(fù)責(zé)寫代碼
coder = Agent(
    name="Coder",
    role="高級程序員",
    system_prompt="""你是一個(gè)高級 Python 程序員。你的任務(wù)是：
1. 根據(jù)需求和調(diào)研結(jié)果編寫高質(zhì)量的 Python 代碼
2. 代碼必須包含完整的類型注解和 docstring
3. 包含錯(cuò)誤處理和日志記錄
4. 輸出可直接運(yùn)行的完整代碼

輸出格式：
## 實(shí)現(xiàn)思路
（簡述思路）

## 完整代碼
```python
（代碼）
```"""
)

# 審查員 Agent：負(fù)責(zé)代碼審查
reviewer = Agent(
    name="Reviewer",
    role="代碼審查專家",
    system_prompt="""你是一個(gè)嚴(yán)格的代碼審查專家。你的任務(wù)是：
1. 審查代碼的正確性、安全性和性能
2. 檢查是否遵循最佳實(shí)踐
3. 給出修改建議

評分標(biāo)準(zhǔn)（滿分 10 分）：
- 正確性：代碼邏輯是否正確
- 安全性：是否存在安全漏洞
- 可讀性：命名、注釋、代碼結(jié)構(gòu)
- 性能：是否有明顯的性能問題

輸出格式：
## 評分：X/10
## 優(yōu)點(diǎn)
（列出優(yōu)點(diǎn)）
## 問題
（列出問題）
## 修改建議
（具體建議）"""
)

# coordinator.py
from agent import Agent, Message


class Coordinator:
    """協(xié)調(diào)者：負(fù)責(zé)調(diào)度多個(gè) Agent 協(xié)作"""

    def __init__(self, researcher, coder, reviewer):
        self.researcher = researcher
        self.coder = coder
        self.reviewer = reviewer
        self.log = []

    def _log(self, phase: str, content: str):
        """記錄執(zhí)行日志"""
        self.log.append({"phase": phase, "content": content})
        print(f"\n{'='*60}")
        print(f"?? [{phase}]")
        print(f"{'='*60}")
        print(content[:500] + "..." if len(content) > 500 else content)

    def run(self, task: str, max_revisions: int = 2) -> dict:
        """
        執(zhí)行完整的多智能體協(xié)作流程

        Args:
            task: 用戶需求描述
            max_revisions: 最大修改輪數(shù)

        Returns:
            包含所有階段輸出的字典
        """
        results = {"task": task, "phases": []}

        # ========== 第一階段：需求分析 + 技術(shù)調(diào)研 ==========
        print("\n?? 第一階段：技術(shù)調(diào)研")
        research_msg = Message(
            role="user", sender="Coordinator",
            content=f"請分析以下需求的技術(shù)方案：\n{task}"
        )
        research_result = self.researcher.think([research_msg])
        self._log("技術(shù)調(diào)研", research_result)
        results["phases"].append({
            "phase": "research", "output": research_result
        })

        # ========== 第二階段：代碼生成 ==========
        print("\n?? 第二階段：代碼生成")
        code_msg = Message(
            role="user", sender="Coordinator",
            content=f"""請根據(jù)以下需求和調(diào)研結(jié)果編寫代碼：

## 需求
{task}

## 技術(shù)調(diào)研結(jié)果
{research_result}

請直接輸出完整可運(yùn)行的 Python 代碼。"""
        )
        code_result = self.coder.think([code_msg])
        self._log("代碼生成", code_result)
        results["phases"].append({
            "phase": "code", "output": code_result
        })

        # ========== 第三階段：代碼審查（可循環(huán)） ==========
        current_code = code_result
        revision = 0

        while revision <= max_revisions:
            print(f"\n?? 第三階段：代碼審查（第 {revision + 1} 次）")
            review_msg = Message(
                role="user", sender="Coordinator",
                content=f"請審查以下代碼：\n{current_code}"
            )
            review_result = self.reviewer.think([review_msg])
            self._log(f"代碼審查（第{revision+1}次）", review_result)

            # 檢查評分（簡單解析）
            score = self._extract_score(review_result)
            print(f"\n?? 當(dāng)前評分：{score}/10")

            if score >= 8 or revision == max_revisions:
                results["phases"].append({
                    "phase": "review",
                    "output": review_result,
                    "score": score
                })
                break

            # 分?jǐn)?shù)不夠，讓程序員修改
            print(f"\n?? 評分不足，進(jìn)入修改...")
            fix_msg = Message(
                role="user", sender="Coordinator",
                content=f"""請根據(jù)審查意見修改代碼：

## 審查意見
{review_result}

## 原始代碼
{current_code}

請輸出修改后的完整代碼。"""
            )
            current_code = self.coder.think([fix_msg])
            self._log(f"代碼修改（第{revision+1}次）", current_code)
            revision += 1

        results["final_code"] = current_code
        results["revisions"] = revision + 1

        return results

    def _extract_score(self, review_text: str) -> int:
        """從審查結(jié)果中提取評分"""
        for line in review_text.split("\n"):
            if "評分" in line and "/" in line:
                try:
                    score = int(line.split(":")[-1].strip().split("/")[0])
                    return score
                except (ValueError, IndexError):
                    pass
        return 5  # 默認(rèn)評分

# main.py
from agents import researcher, coder, reviewer
from coordinator import Coordinator
import json


def main():
    print("?? AI 研發(fā)小隊(duì)已就緒！")
    print("=" * 60)

    # 初始化協(xié)調(diào)者
    coordinator = Coordinator(researcher, coder, reviewer)

    # 示例任務(wù)
    task = "用 Python 寫一個(gè)文件批量重命名工具，支持按日期、序號、正則匹配重命名，需要命令行參數(shù)解析和日志記錄"

    # 執(zhí)行任務(wù)
    results = coordinator.run(task, max_revisions=2)

    # 輸出最終結(jié)果
    print("\n" + "=" * 60)
    print("? 任務(wù)完成！")
    print("=" * 60)
    print(f"\n?? 總修改輪數(shù)：{results['revisions']}")
    print(f"\n?? 最終代碼：\n{results['final_code']}")

    # 保存結(jié)果
    with open("output.json", "w", encoding="utf-8") as f:
        json.dump(results, f, ensure_ascii=False, indent=2)
    print("\n?? 完整結(jié)果已保存到 output.json")


if __name__ == "__main__":
    main()

# tools.py
import os
import subprocess


def write_file(path: str, content: str) -> str:
    """寫入文件"""
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        f.write(content)
    return f"文件已寫入: {path}"


def run_python(filepath: str) -> str:
    """運(yùn)行 Python 文件并返回結(jié)果"""
    try:
        result = subprocess.run(
            ["python", filepath],
            capture_output=True, text=True, timeout=30
        )
        output = result.stdout
        if result.returncode != 0:
            output += f"\n? 錯(cuò)誤:\n{result.stderr}"
        return output
    except subprocess.TimeoutExpired:
        return "執(zhí)行超時(shí)（30秒）"


# 將工具注冊到 Agent
AVAILABLE_TOOLS = {
    "write_file": {
        "func": write_file,
        "description": "將內(nèi)容寫入指定文件",
        "params": ["path", "content"]
    },
    "run_python": {
        "func": run_python,
        "description": "運(yùn)行 Python 文件",
        "params": ["filepath"]
    }
}