#!/bin/bash
# 数据库初始化和数据迁移脚本

set -e

echo "Agent Manager - Database Initialization and Migration"
echo "======================================================"

# 1. 备份现有数据（如果存在）
if [ -f "agent_manager.db" ]; then
    echo "Backing up existing database..."
    cp agent_manager.db agent_manager.db.backup.$(date +%Y%m%d_%H%M%S)
fi

# 2. 初始化数据库（如果不存在）
echo "Initializing database..."
python3 << 'EOF'
from database import init_db, SessionLocal, Template, Quota, AgentType
from datetime import datetime

# 创建所有表
init_db()
print("✓ Database tables created")

# 创建示例模板数据
db = SessionLocal()

try:
    # 检查是否已有模板
    existing_templates = db.query(Template).count()
    
    if existing_templates == 0:
        print("\nCreating default templates...")
        
        # MySQL Agent (Platform)
        mysql_template = Template(
            name="mysql_agent",
            display_name="MySQL Query Agent",
            description="Platform MySQL database query agent using LangChain",
            agent_type=AgentType.PLATFORM,
            image="agnettaiji.azurecr.io/mysql_agent:latest",
            port=None,
            env_requirements={
                "required": {
                    "MYSQL_HOST": "MySQL server hostname",
                    "MYSQL_USER": "MySQL username",
                    "MYSQL_PASSWORD": "MySQL password",
                    "MYSQL_DATABASE": "MySQL database name",
                    "OPENAI_API_KEY": "OpenAI API key for LangChain"
                },
                "optional": {
                    "MYSQL_PORT": "MySQL port (default: 3306)"
                }
            },
            cpu_request="100m",
            cpu_limit="500m",
            memory_request="256Mi",
            memory_limit="512Mi",
            min_replicas=1,
            max_replicas=3,
            target_cpu_utilization=80
        )
        db.add(mysql_template)
        
        # PostgreSQL Agent (Custom)
        pg_template = Template(
            name="postgresql_agent",
            display_name="PostgreSQL Query Agent",
            description="Custom PostgreSQL database query agent using LangChain",
            agent_type=AgentType.CUSTOM,
            image="agnettaiji.azurecr.io/postgresql_agent:latest",
            port=None,
            env_requirements={
                "required": {
                    "POSTGRES_HOST": "PostgreSQL server hostname",
                    "POSTGRES_USER": "PostgreSQL username",
                    "POSTGRES_PASSWORD": "PostgreSQL password",
                    "POSTGRES_DATABASE": "PostgreSQL database name",
                    "OPENAI_API_KEY": "OpenAI API key for LangChain"
                },
                "optional": {
                    "POSTGRES_PORT": "PostgreSQL port (default: 5432)"
                }
            },
            cpu_request="100m",
            cpu_limit="500m",
            memory_request="256Mi",
            memory_limit="512Mi",
            min_replicas=1,
            max_replicas=5,
            target_cpu_utilization=80
        )
        db.add(pg_template)
        
        # Jina Search Agent (Custom with HTTP service)
        jina_template = Template(
            name="jina_search_agent",
            display_name="Jina Search Agent",
            description="Custom web search agent using Jina Reader API",
            agent_type=AgentType.CUSTOM,
            image="agnettaiji.azurecr.io/jina_search_agent:latest",
            port=8080,
            env_requirements={
                "required": {
                    "JINA_API_KEY": "Jina API key"
                },
                "optional": {
                    "SERVICE_PORT": "HTTP service port (default: 8080)",
                    "SERVICE_HOST": "HTTP service host (default: 0.0.0.0)"
                }
            },
            cpu_request="100m",
            cpu_limit="500m",
            memory_request="128Mi",
            memory_limit="256Mi",
            min_replicas=1,
            max_replicas=5,
            target_cpu_utilization=80
        )
        db.add(jina_template)
        
        # Intelligent Search Agent (Custom with HTTP service)
        search_template = Template(
            name="search_agent",
            display_name="Intelligent Search Agent",
            description="智能搜索代理，提供多源搜索和答案生成功能",
            agent_type=AgentType.CUSTOM,
            agent_framework="langchain",
            image="agnettaiji.azurecr.io/ai-agents/search-agent:v1.0",
            port=8080,
            env_requirements={
                "required": {
                    "LLM_BASE_URL": "LLM API基础URL",
                    "LLM_API_KEY": "LLM API密钥",
                    "SERPER_API_KEY": "Serper API密钥（Google搜索）",
                    "JINA_API_KEY": "Jina API密钥（内容提取和重排序）"
                },
                "optional": {
                    "LLM_MODEL": "LLM模型名称，默认xchat52",
                    "MAX_ITERATIONS": "最大迭代次数，默认3",
                    "MAX_RESULTS_PER_QUERY": "每次搜索最大结果数，默认10",
                    "CONTENT_MAX_LENGTH": "内容最大长度，默认5000",
                    "LOG_LEVEL": "日志级别，默认INFO",
                    "TIMEOUT": "超时时间(秒)，默认30"
                }
            },
            cpu_request="500m",
            cpu_limit="1000m",
            memory_request="512Mi",
            memory_limit="1Gi",
            min_replicas=1,
            max_replicas=3,
            target_cpu_utilization=70
        )
        db.add(search_template)
        
        # Echo Agent (Platform - simple example)
        echo_template = Template(
            name="echo_agent",
            display_name="Echo Agent",
            description="Simple platform echo agent for testing",
            agent_type=AgentType.PLATFORM,
            image="busybox:latest",
            port=None,
            env_requirements={},
            cpu_request="50m",
            cpu_limit="100m",
            memory_request="64Mi",
            memory_limit="128Mi",
            min_replicas=1,
            max_replicas=2,
            target_cpu_utilization=80
        )
        db.add(echo_template)
        
        db.commit()
        print("✓ Created 5 default templates")
        
        # 创建默认配额
        print("\nCreating default quotas...")
        
        # 默认管理员配额
        admin_quota = Quota(
            owner_type="admin",
            owner_id="admin",
            platform_pod_quota=100,
            platform_pod_used=0,
            custom_cpu_quota=50.0,
            custom_cpu_used=0.0,
            custom_memory_quota=102400.0,  # 100GB
            custom_memory_used=0.0
        )
        db.add(admin_quota)
        
        # 默认租户配额
        tenant_quota = Quota(
            owner_type="tenant",
            owner_id="default_tenant",
            platform_pod_quota=10,
            platform_pod_used=0,
            custom_cpu_quota=5.0,
            custom_cpu_used=0.0,
            custom_memory_quota=10240.0,  # 10GB
            custom_memory_used=0.0
        )
        db.add(tenant_quota)
        
        db.commit()
        print("✓ Created default quotas")
        
    else:
        print(f"\n✓ Database already contains {existing_templates} templates")
    
    print("\n✅ Database initialization completed successfully!")
    
except Exception as e:
    print(f"\n❌ Error: {str(e)}")
    db.rollback()
    raise
finally:
    db.close()

EOF

echo ""
echo "======================================================"
echo "Database initialized at: agent_manager.db"
echo ""
echo "Next steps:"
echo "1. Review database content: sqlite3 agent_manager.db"
echo "2. Start the service: python app_new.py"
echo "3. Test API: curl http://localhost:8000/templates"
echo "======================================================"