#!/bin/bash
# Agent版本清理脚本
# 将v2版本替换为正式版本，删除旧版本

set -e

echo "=============================================="
echo "BBAV20311 Agent版本清理"
echo "=============================================="
echo ""

AGENTS_DIR="/home/xia-lab/wutx/00test/bbav20311/gfa_core/mas/agents"
BACKUP_DIR="/home/xia-lab/wutx/00test/bbav20311/gfa_core/mas/agents/backup"

echo "1. 创建备份目录"
mkdir -p "$BACKUP_DIR"
echo "   备份目录: $BACKUP_DIR"
echo ""

echo "2. 备份旧版本文件"
for agent in diagnosis planning execution knowledge learning result_parser validator; do
    if [ -f "$AGENTS_DIR/${agent}.py" ]; then
        cp "$AGENTS_DIR/${agent}.py" "$BACKUP_DIR/${agent}_v1.py.bak"
        echo "   ✓ 备份 ${agent}.py -> backup/${agent}_v1.py.bak"
    fi
done
echo ""

echo "3. 将v2版本替换为正式版本"
for agent in diagnosis planning execution; do
    if [ -f "$AGENTS_DIR/${agent}_v2.py" ]; then
        # 备份v2版本
        cp "$AGENTS_DIR/${agent}_v2.py" "$BACKUP_DIR/${agent}_v2.py.bak"
        
        # 替换：v2 -> 正式版本
        cp "$AGENTS_DIR/${agent}_v2.py" "$AGENTS_DIR/${agent}.py"
        
        # 修改类名：AgentV2 -> Agent
        sed -i "s/class ${agent^}AgentV2/class ${agent^}Agent/g" "$AGENTS_DIR/${agent}.py"
        sed -i "s/name = \"${agent}_v2\"/name = \"${agent}\"/g" "$AGENTS_DIR/${agent}.py"
        
        echo "   ✓ ${agent}_v2.py -> ${agent}.py (类名更新)"
    fi
done
echo ""

echo "4. 检查其他Agent的v2版本"
for agent in knowledge learning result_parser validator; do
    if [ -f "$AGENTS_DIR/${agent}_v2.py" ]; then
        v1_lines=$(wc -l < "$AGENTS_DIR/${agent}.py" 2>/dev/null || echo 0)
        v2_lines=$(wc -l < "$AGENTS_DIR/${agent}_v2.py" 2>/dev/null || echo 0)
        
        if [ "$v2_lines" -gt "$v1_lines" ]; then
            # v2版本更完整，替换
            cp "$AGENTS_DIR/${agent}_v2.py" "$BACKUP_DIR/${agent}_v2.py.bak"
            cp "$AGENTS_DIR/${agent}_v2.py" "$AGENTS_DIR/${agent}.py"
            
            sed -i "s/name = \"${agent}_v2\"/name = \"${agent}\"/g" "$AGENTS_DIR/${agent}.py"
            echo "   ✓ ${agent}_v2.py -> ${agent}.py (v2更完整，$v1_lines -> $v2_lines行)"
        else
            echo "   - ${agent}.py 保持v1版本 ($v1_lines行 > $v2_lines行)"
        fi
    fi
done
echo ""

echo "5. 更新__init__.py导入"
cat > "$AGENTS_DIR/__init__.py" << 'EOF'
from .base import BaseAgent, AgentState, AgentMessage, QueryMode
from .router import RouterAgent
from .planning import PlanningAgent
from .execution import ExecutionAgent
from .diagnosis import DiagnosisAgent
from .knowledge import KnowledgeAgent
from .learning import LearningAgent
from .result_parser import ResultParserAgent
from .validator import ValidatorAgent

__all__ = [
    "BaseAgent",
    "AgentState",
    "AgentMessage",
    "QueryMode",
    "RouterAgent",
    "PlanningAgent",
    "ExecutionAgent",
    "DiagnosisAgent",
    "KnowledgeAgent",
    "LearningAgent",
    "ResultParserAgent",
    "ValidatorAgent",
]
EOF
echo "   ✓ __init__.py 已更新"
echo ""

echo "6. 更新orchestrator.py导入"
ORCHESTRATOR_FILE="/home/xia-lab/wutx/00test/bbav20311/gfa_core/mas/core/orchestrator.py"

# 备份原文件
cp "$ORCHESTRATOR_FILE" "$BACKUP_DIR/orchestrator.py.bak"

# 替换导入
sed -i 's/from gfa_core.mas.agents.diagnosis_v2 import DiagnosisAgentV2/from gfa_core.mas.agents.diagnosis import DiagnosisAgent/g' "$ORCHESTRATOR_FILE"
sed -i 's/from gfa_core.mas.agents.planning_v2 import PlanningAgentV2/from gfa_core.mas.agents.planning import PlanningAgent/g' "$ORCHESTRATOR_FILE"
sed -i 's/from gfa_core.mas.agents.execution_v2 import ExecutionAgentV2/from gfa_core.mas.agents.execution import ExecutionAgent/g' "$ORCHESTRATOR_FILE"

# 替换类实例化
sed -i 's/DiagnosisAgentV2/DiagnosisAgent/g' "$ORCHESTRATOR_FILE"
sed -i 's/PlanningAgentV2/PlanningAgent/g' "$ORCHESTRATOR_FILE"
sed -i 's/ExecutionAgentV2/ExecutionAgent/g' "$ORCHESTRATOR_FILE"

echo "   ✓ orchestrator.py 导入已更新"
echo ""

echo "7. 清理v2文件(保留备份)"
for agent in diagnosis planning execution knowledge learning result_parser validator; do
    if [ -f "$AGENTS_DIR/${agent}_v2.py" ]; then
        mv "$AGENTS_DIR/${agent}_v2.py" "$BACKUP_DIR/${agent}_v2.py"
        echo "   ✓ 移动 ${agent}_v2.py -> backup/"
    fi
done
echo ""

echo "=============================================="
echo "清理完成"
echo "=============================================="
echo ""
echo "备份文件位置: $BACKUP_DIR"
echo ""
echo "清理后的Agent文件:"
ls -la "$AGENTS_DIR/"*.py | grep -v "__pycache__" | grep -v "backup"
echo ""
echo "注意: 请重启服务使更改生效"