#!/usr/bin/env python3
"""
AI Dev System - Database Manager
資料庫管理工具
"""

import sqlite3
import json
from datetime import datetime
from pathlib import Path
import os

# 資料庫位置
DB_PATH = Path.home() / ".ai-dev-logs" / "ai-dev-system.db"

def init_database():
    """初始化資料庫結構"""
    DB_PATH.parent.mkdir(parents=True, exist_ok=True)
    
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    # 執行歷史表
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS execution_history (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            project_name TEXT NOT NULL,
            prompt TEXT NOT NULL,
            ai_provider TEXT NOT NULL,
            ai_output TEXT,
            execution_status TEXT DEFAULT 'running',
            iteration_number INTEGER,
            git_commit_hash TEXT,
            error_message TEXT,
            started_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            completed_at TIMESTAMP,
            duration_seconds INTEGER,
            token_count INTEGER,
            cost_estimate REAL
        )
    """)
    
    # 專案配置表
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS project_configs (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            project_name TEXT UNIQUE NOT NULL,
            project_dir TEXT NOT NULL,
            ai_provider TEXT DEFAULT 'claude',
            is_enabled BOOLEAN DEFAULT 1,
            max_iterations INTEGER DEFAULT 50,
            current_iteration INTEGER DEFAULT 0,
            is_completed BOOLEAN DEFAULT 0,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """)
    
    # 環境變數表
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS environment_variables (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            var_name TEXT UNIQUE NOT NULL,
            var_value TEXT NOT NULL,
            var_description TEXT,
            is_sensitive BOOLEAN DEFAULT 0,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """)
    
    # Prompt 範本表
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS prompt_templates (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            template_name TEXT UNIQUE NOT NULL,
            template_content TEXT NOT NULL,
            description TEXT,
            category TEXT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """)
    
    # 系統狀態表
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS system_status (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            current_project TEXT,
            is_running BOOLEAN DEFAULT 0,
            last_execution_time TIMESTAMP,
            total_executions INTEGER DEFAULT 0,
            total_tokens_used INTEGER DEFAULT 0,
            total_cost REAL DEFAULT 0.0,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """)
    
    # 建立索引
    cursor.execute("""
        CREATE INDEX IF NOT EXISTS idx_execution_project 
        ON execution_history(project_name)
    """)
    
    cursor.execute("""
        CREATE INDEX IF NOT EXISTS idx_execution_time 
        ON execution_history(started_at DESC)
    """)
    
    # 初始化系統狀態（如果不存在）
    cursor.execute("SELECT COUNT(*) FROM system_status")
    if cursor.fetchone()[0] == 0:
        cursor.execute("""
            INSERT INTO system_status (current_project, is_running)
            VALUES ('', 0)
        """)
    
    conn.commit()
    conn.close()
    print(f"✅ 資料庫初始化完成: {DB_PATH}")

def record_execution_start(project_name, prompt, ai_provider, iteration):
    """記錄執行開始"""
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    cursor.execute("""
        INSERT INTO execution_history 
        (project_name, prompt, ai_provider, execution_status, iteration_number)
        VALUES (?, ?, ?, 'running', ?)
    """, (project_name, prompt, ai_provider, iteration))
    
    execution_id = cursor.lastrowid
    
    # 更新系統狀態
    cursor.execute("""
        UPDATE system_status 
        SET is_running = 1, 
            last_execution_time = CURRENT_TIMESTAMP,
            total_executions = total_executions + 1
        WHERE id = 1
    """)
    
    conn.commit()
    conn.close()
    
    return execution_id

def record_execution_complete(execution_id, ai_output, git_commit, 
                              duration, status='completed', error=None,
                              token_count=None):
    """記錄執行完成"""
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    # 估算成本（假設 Claude Sonnet 4）
    cost = 0.0
    if token_count:
        # Input: $3/MTok, Output: $15/MTok (粗略估算各半)
        cost = (token_count / 1000000) * 9  # 平均值
    
    cursor.execute("""
        UPDATE execution_history 
        SET ai_output = ?,
            git_commit_hash = ?,
            execution_status = ?,
            error_message = ?,
            completed_at = CURRENT_TIMESTAMP,
            duration_seconds = ?,
            token_count = ?,
            cost_estimate = ?
        WHERE id = ?
    """, (ai_output, git_commit, status, error, duration, token_count, cost, execution_id))
    
    # 更新系統狀態
    cursor.execute("""
        UPDATE system_status 
        SET is_running = 0,
            total_tokens_used = total_tokens_used + ?,
            total_cost = total_cost + ?
        WHERE id = 1
    """, (token_count or 0, cost))
    
    conn.commit()
    conn.close()

def get_project_config(project_name):
    """獲取專案配置"""
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    cursor = conn.cursor()
    
    cursor.execute("""
        SELECT * FROM project_configs WHERE project_name = ?
    """, (project_name,))
    
    result = cursor.fetchone()
    conn.close()
    
    return dict(result) if result else None

def update_project_config(project_name, **kwargs):
    """更新專案配置"""
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    # 先檢查專案是否存在
    cursor.execute("""
        SELECT id FROM project_configs WHERE project_name = ?
    """, (project_name,))
    
    if cursor.fetchone():
        # 更新
        fields = ', '.join([f"{k} = ?" for k in kwargs.keys()])
        values = list(kwargs.values()) + [project_name]
        cursor.execute(f"""
            UPDATE project_configs 
            SET {fields}, updated_at = CURRENT_TIMESTAMP
            WHERE project_name = ?
        """, values)
    else:
        # 插入
        kwargs['project_name'] = project_name
        fields = ', '.join(kwargs.keys())
        placeholders = ', '.join(['?' for _ in kwargs])
        cursor.execute(f"""
            INSERT INTO project_configs ({fields})
            VALUES ({placeholders})
        """, list(kwargs.values()))
    
    conn.commit()
    conn.close()

def get_execution_history(project_name=None, limit=50):
    """獲取執行歷史"""
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    cursor = conn.cursor()
    
    if project_name:
        cursor.execute("""
            SELECT * FROM execution_history 
            WHERE project_name = ?
            ORDER BY started_at DESC 
            LIMIT ?
        """, (project_name, limit))
    else:
        cursor.execute("""
            SELECT * FROM execution_history 
            ORDER BY started_at DESC 
            LIMIT ?
        """, (limit,))
    
    results = [dict(row) for row in cursor.fetchall()]
    conn.close()
    
    return results

def get_system_stats():
    """獲取系統統計"""
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    cursor = conn.cursor()
    
    # 系統狀態
    cursor.execute("SELECT * FROM system_status WHERE id = 1")
    status = dict(cursor.fetchone())
    
    # 專案統計
    cursor.execute("""
        SELECT 
            project_name,
            COUNT(*) as execution_count,
            SUM(CASE WHEN execution_status = 'completed' THEN 1 ELSE 0 END) as completed_count,
            SUM(CASE WHEN execution_status = 'failed' THEN 1 ELSE 0 END) as failed_count,
            SUM(token_count) as total_tokens,
            SUM(cost_estimate) as total_cost,
            MAX(started_at) as last_execution
        FROM execution_history
        GROUP BY project_name
    """)
    
    projects = [dict(row) for row in cursor.fetchall()]
    
    conn.close()
    
    return {
        'status': status,
        'projects': projects
    }

if __name__ == '__main__':
    import sys
    
    if len(sys.argv) > 1 and sys.argv[1] == 'init':
        init_database()
    else:
        print("Usage: python db_manager.py init")