#!/usr/bin/env python3
"""
简化的数据库初始化脚本
"""

import sqlite3
import json
import os
from datetime import datetime

# 数据库配置
DATABASE_CONFIG = {
    "projects": {
        "path": "/root/data/disk/system/database/projects.db",
        "journal_mode": "WAL",
        "cache_size": -2000
    },
    "agents": {
        "path": "/root/data/disk/system/database/agents.db",
        "journal_mode": "WAL"
    },
    "analytics": {
        "path": "/root/data/disk/system/database/analytics.db",
        "journal_mode": "WAL"
    }
}

def init_database(db_name):
    """初始化指定数据库"""
    db_config = DATABASE_CONFIG[db_name]
    db_path = db_config['path']
    db_dir = os.path.dirname(db_path)
    
    # 创建数据库目录
    os.makedirs(db_dir, exist_ok=True)
    
    print(f"初始化数据库: {db_name} -> {db_path}")
    
    # 连接到数据库
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row
    
    try:
        # 设置PRAGMA参数
        cursor = conn.cursor()
        cursor.execute("PRAGMA journal_mode=WAL")
        cursor.execute("PRAGMA busy_timeout=5000")
        cursor.execute("PRAGMA foreign_keys=ON")
        cursor.execute("PRAGMA synchronous=NORMAL")
        
        if db_name == 'projects':
            cursor.execute("PRAGMA cache_size=-2000")
            cursor.execute("PRAGMA temp_store=MEMORY")
        
        cursor.close()
        
        # 创建表结构
        if db_name == 'projects':
            create_projects_tables(conn)
        elif db_name == 'agents':
            create_agents_tables(conn)
        elif db_name == 'analytics':
            create_analytics_tables(conn)
        
        # 插入初始数据
        if db_name == 'agents':
            insert_agent_initial_data(conn)
        
        conn.commit()
        print(f"✅ 数据库 {db_name} 初始化成功")
        
    except Exception as e:
        print(f"❌ 数据库 {db_name} 初始化失败: {e}")
        conn.rollback()
        raise
    finally:
        conn.close()

def create_projects_tables(conn):
    """创建项目数据库表"""
    cursor = conn.cursor()
    
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS projects (
        project_id TEXT PRIMARY KEY,
        name TEXT NOT NULL,
        description TEXT,
        status TEXT CHECK(status IN (
            'draft', 'collecting', 'creating', 'producing', 
            'reviewing', 'publishing', 'analyzing', 'completed', 'failed'
        )),
        workflow_step TEXT,
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        created_by TEXT,
        metadata JSON
    )
    ''')
    
    cursor.execute('''
    CREATE TRIGGER IF NOT EXISTS update_projects_timestamp 
    AFTER UPDATE ON projects
    BEGIN
        UPDATE projects SET updated_at = CURRENT_TIMESTAMP 
        WHERE project_id = NEW.project_id;
    END
    ''')
    
    cursor.execute('CREATE INDEX IF NOT EXISTS idx_projects_status ON projects(status)')
    cursor.execute('CREATE INDEX IF NOT EXISTS idx_projects_created ON projects(created_at)')
    
    cursor.close()

def create_agents_tables(conn):
    """创建Agent数据库表"""
    cursor = conn.cursor()
    
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS agent_tasks (
        task_id TEXT PRIMARY KEY,
        project_id TEXT,
        agent_name TEXT CHECK(agent_name IN (
            'cto', 'collector', 'writer', 'producer', 'publisher', 'analyzer'
        )),
        task_type TEXT,
        priority INTEGER DEFAULT 1,
        status TEXT CHECK(status IN (
            'pending', 'running', 'completed', 'failed', 'retrying'
        )),
        input_data JSON,
        output_data JSON,
        dependencies JSON,
        retry_count INTEGER DEFAULT 0,
        max_retries INTEGER DEFAULT 3,
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        started_at TIMESTAMP,
        completed_at TIMESTAMP,
        error_message TEXT
    )
    ''')
    
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS agent_status (
        agent_name TEXT PRIMARY KEY,
        status TEXT CHECK(status IN ('idle', 'busy', 'error', 'offline')),
        current_task_id TEXT,
        last_heartbeat TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        performance_metrics JSON
    )
    ''')
    
    cursor.execute('CREATE INDEX IF NOT EXISTS idx_agent_tasks_status ON agent_tasks(status)')
    cursor.execute('CREATE INDEX IF NOT EXISTS idx_agent_tasks_agent ON agent_tasks(agent_name)')
    
    cursor.close()

def create_analytics_tables(conn):
    """创建分析数据库表"""
    cursor = conn.cursor()
    
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS ideas (
        idea_id TEXT PRIMARY KEY,
        project_id TEXT,
        title TEXT,
        description TEXT,
        source TEXT,
        hotness_score REAL,
        potential_score REAL,
        tags JSON,
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
    )
    ''')
    
    cursor.execute('''
    CREATE TABLE IF NOT EXISTS publish_records (
        publish_id TEXT PRIMARY KEY,
        project_id TEXT,
        platform TEXT,
        status TEXT CHECK(status IN ('pending', 'published', 'failed', 'removed')),
        publish_url TEXT,
        publish_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        view_count INTEGER DEFAULT 0,
        like_count INTEGER DEFAULT 0,
        comment_count INTEGER DEFAULT 0,
        share_count INTEGER DEFAULT 0,
        feedback_summary TEXT
    )
    ''')
    
    cursor.execute('CREATE INDEX IF NOT EXISTS idx_ideas_project ON ideas(project_id)')
    cursor.execute('CREATE INDEX IF NOT EXISTS idx_publish_records_project ON publish_records(project_id)')
    
    cursor.close()

def insert_agent_initial_data(conn):
    """插入Agent初始数据"""
    cursor = conn.cursor()
    
    agents = [
        ('cto', 'idle', None, datetime.now().isoformat(), '{}'),
        ('collector', 'idle', None, datetime.now().isoformat(), '{}'),
        ('writer', 'idle', None, datetime.now().isoformat(), '{}'),
        ('producer', 'idle', None, datetime.now().isoformat(), '{}'),
        ('publisher', 'idle', None, datetime.now().isoformat(), '{}'),
        ('analyzer', 'idle', None, datetime.now().isoformat(), '{}'),
    ]
    
    cursor.executemany('''
    INSERT OR IGNORE INTO agent_status 
    (agent_name, status, current_task_id, last_heartbeat, performance_metrics)
    VALUES (?, ?, ?, ?, ?)
    ''', agents)
    
    cursor.close()
    print(f"✅ 插入了 {len(agents)} 个Agent状态记录")

def main():
    """主函数"""
    print("=" * 60)
    print("开始初始化慢剧工作台数据库系统")
    print("=" * 60)
    
    databases = ['projects', 'agents', 'analytics']
    
    for db_name in databases:
        print(f"\n📊 处理数据库: {db_name}")
        print("-" * 40)
        
        try:
            init_database(db_name)
        except Exception as e:
            print(f"❌ 数据库 {db_name} 初始化失败: {e}")
            continue
    
    print("\n" + "=" * 60)
    print("数据库初始化完成")
    print("=" * 60)

if __name__ == "__main__":
    main()