#!/usr/bin/env python3
"""
测试任务统计修复功能
"""

import sqlite3
import json
import uuid
from datetime import datetime, timezone

def create_mixed_test_data():
    """创建混合的测试数据（单个任务和批量任务）"""
    db_path = r"C:\Users\imeep\AppData\Roaming\mixvideo\mixvideoV2.db"
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    try:
        print("=== 创建混合测试数据 ===")
        
        # 清理旧的测试数据
        cursor.execute("DELETE FROM uni_comfyui_task WHERE workflow_name = 'test_statistics'")
        cursor.execute("DELETE FROM uni_comfyui_batch_task WHERE workflow_name = 'test_statistics'")
        cursor.execute("DELETE FROM uni_comfyui_batch_item WHERE batch_id LIKE 'stats_test_%'")
        
        # 创建单个任务（5个）
        single_tasks = []
        for i in range(5):
            workflow_run_id = f"single_task_{uuid.uuid4().hex[:8]}"
            status = ['completed', 'failed', 'running', 'queued', 'completed'][i]
            
            cursor.execute("""
                INSERT INTO uni_comfyui_task 
                (workflow_run_id, workflow_name, task_type, status, input_params, 
                 result_data, error_message, created_at, started_at, completed_at, server_url)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """, (
                workflow_run_id, "test_statistics", "single", status,
                json.dumps({"test": f"single_{i}"}),
                json.dumps({"result": f"single_result_{i}"}) if status == 'completed' else None,
                f"Single task error {i}" if status == 'failed' else None,
                datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S.%f'),
                datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S.%f') if status != 'queued' else None,
                datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S.%f') if status in ['completed', 'failed'] else None,
                "http://test"
            ))
            
            single_tasks.append({
                'workflow_run_id': workflow_run_id,
                'status': status,
                'task_id': cursor.lastrowid
            })
        
        print(f"✅ 创建了 {len(single_tasks)} 个单个任务")
        for task in single_tasks:
            print(f"   - {task['workflow_run_id']}: {task['status']}")
        
        # 创建批量任务（3个批量任务，每个包含2个子任务）
        batch_tasks = []
        batch_subtasks = []
        
        for batch_idx in range(3):
            batch_id = f"stats_test_batch_{batch_idx}"
            
            # 创建批量任务记录
            cursor.execute("""
                INSERT INTO uni_comfyui_batch_task 
                (batch_id, workflow_name, total_count, completed_count, failed_count, status, created_at)
                VALUES (?, ?, ?, ?, ?, ?, ?)
            """, (
                batch_id, "test_statistics", 2, 1, 1, 'running',
                datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S.%f')
            ))
            
            batch_tasks.append({
                'batch_id': batch_id,
                'status': 'running'
            })
            
            # 为每个批量任务创建2个子任务
            for sub_idx in range(2):
                workflow_run_id = f"batch_task_{batch_idx}_{sub_idx}_{uuid.uuid4().hex[:8]}"
                status = 'completed' if sub_idx == 0 else 'failed'
                
                cursor.execute("""
                    INSERT INTO uni_comfyui_task 
                    (workflow_run_id, workflow_name, task_type, status, input_params, 
                     result_data, error_message, created_at, started_at, completed_at, server_url)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """, (
                    workflow_run_id, "test_statistics", "batch", status,
                    json.dumps({"test": f"batch_{batch_idx}_{sub_idx}"}),
                    json.dumps({"result": f"batch_result_{batch_idx}_{sub_idx}"}) if status == 'completed' else None,
                    f"Batch task error {batch_idx}_{sub_idx}" if status == 'failed' else None,
                    datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S.%f'),
                    datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S.%f'),
                    datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S.%f'),
                    "http://test"
                ))
                
                task_id = cursor.lastrowid
                
                # 创建批量任务子项关联
                cursor.execute("""
                    INSERT INTO uni_comfyui_batch_item (batch_id, task_id, item_index)
                    VALUES (?, ?, ?)
                """, (batch_id, task_id, sub_idx))
                
                batch_subtasks.append({
                    'workflow_run_id': workflow_run_id,
                    'batch_id': batch_id,
                    'status': status,
                    'task_id': task_id
                })
        
        print(f"✅ 创建了 {len(batch_tasks)} 个批量任务，包含 {len(batch_subtasks)} 个子任务")
        for batch in batch_tasks:
            print(f"   - {batch['batch_id']}: {batch['status']}")
        
        conn.commit()
        
        print("\n📊 测试数据统计:")
        print(f"   - 单个任务: 5个 (2个completed, 1个failed, 1个running, 1个queued)")
        print(f"   - 批量任务: 3个批量任务，6个子任务 (3个completed, 3个failed)")
        print(f"   - 总任务数: 11个")
        
        return True
        
    except Exception as e:
        conn.rollback()
        print(f"❌ 创建测试数据失败: {e}")
        import traceback
        traceback.print_exc()
        return False
    finally:
        conn.close()

def verify_database_statistics():
    """验证数据库中的统计数据"""
    db_path = r"C:\Users\imeep\AppData\Roaming\mixvideo\mixvideoV2.db"
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    try:
        print("\n=== 验证数据库统计 ===")
        
        # 总体统计
        cursor.execute("""
            SELECT 
                COUNT(*) as total,
                SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END) as completed,
                SUM(CASE WHEN status = 'failed' THEN 1 ELSE 0 END) as failed,
                SUM(CASE WHEN status = 'running' THEN 1 ELSE 0 END) as running,
                SUM(CASE WHEN status = 'queued' THEN 1 ELSE 0 END) as queued
            FROM uni_comfyui_task 
            WHERE workflow_name = 'test_statistics'
        """)
        
        total_stats = cursor.fetchone()
        print(f"📊 总体统计: 总计={total_stats[0]}, 完成={total_stats[1]}, 失败={total_stats[2]}, 运行中={total_stats[3]}, 排队={total_stats[4]}")
        
        # 单个任务统计
        cursor.execute("""
            SELECT 
                COUNT(*) as total,
                SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END) as completed,
                SUM(CASE WHEN status = 'failed' THEN 1 ELSE 0 END) as failed,
                SUM(CASE WHEN status = 'running' THEN 1 ELSE 0 END) as running,
                SUM(CASE WHEN status = 'queued' THEN 1 ELSE 0 END) as queued
            FROM uni_comfyui_task 
            WHERE workflow_name = 'test_statistics' AND task_type = 'single'
        """)
        
        single_stats = cursor.fetchone()
        print(f"📊 单个任务统计: 总计={single_stats[0]}, 完成={single_stats[1]}, 失败={single_stats[2]}, 运行中={single_stats[3]}, 排队={single_stats[4]}")
        
        # 批量任务统计
        cursor.execute("""
            SELECT 
                COUNT(*) as total,
                SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END) as completed,
                SUM(CASE WHEN status = 'failed' THEN 1 ELSE 0 END) as failed,
                SUM(CASE WHEN status = 'running' THEN 1 ELSE 0 END) as running,
                SUM(CASE WHEN status = 'queued' THEN 1 ELSE 0 END) as queued
            FROM uni_comfyui_task 
            WHERE workflow_name = 'test_statistics' AND task_type = 'batch'
        """)
        
        batch_stats = cursor.fetchone()
        print(f"📊 批量任务统计: 总计={batch_stats[0]}, 完成={batch_stats[1]}, 失败={batch_stats[2]}, 运行中={batch_stats[3]}, 排队={batch_stats[4]}")
        
        # 验证数据正确性
        expected_single = (5, 2, 1, 1, 1)  # 总计, 完成, 失败, 运行中, 排队
        expected_batch = (6, 3, 3, 0, 0)   # 总计, 完成, 失败, 运行中, 排队
        expected_total = (11, 5, 4, 1, 1)  # 总计, 完成, 失败, 运行中, 排队
        
        print("\n✅ 数据验证:")
        print(f"   单个任务: 期望{expected_single}, 实际{single_stats} - {'✅' if single_stats == expected_single else '❌'}")
        print(f"   批量任务: 期望{expected_batch}, 实际{batch_stats} - {'✅' if batch_stats == expected_batch else '❌'}")
        print(f"   总体统计: 期望{expected_total}, 实际{total_stats} - {'✅' if total_stats == expected_total else '❌'}")
        
        return single_stats == expected_single and batch_stats == expected_batch and total_stats == expected_total
        
    except Exception as e:
        print(f"❌ 验证失败: {e}")
        import traceback
        traceback.print_exc()
        return False
    finally:
        conn.close()

def show_expected_api_response():
    """显示预期的API响应格式"""
    print("\n=== 预期的API响应格式 ===")
    
    expected_response = {
        "total_tasks": 11,
        "completed_tasks": 5,
        "failed_tasks": 4,
        "running_tasks": 1,
        "queued_tasks": 1,
        "success_rate": 45.45,  # 5/11 * 100
        "average_execution_time": 0.0,
        "total_execution_time": 0.0,
        "most_used_workflows": [],
        "daily_stats": [],
        "single_task_stats": {
            "total_tasks": 5,
            "completed_tasks": 2,
            "failed_tasks": 1,
            "running_tasks": 1,
            "queued_tasks": 1,
            "success_rate": 40.0  # 2/5 * 100
        },
        "batch_task_stats": {
            "total_tasks": 6,
            "completed_tasks": 3,
            "failed_tasks": 3,
            "running_tasks": 0,
            "queued_tasks": 0,
            "success_rate": 50.0  # 3/6 * 100
        }
    }
    
    print(json.dumps(expected_response, indent=2, ensure_ascii=False))

def cleanup_test_data():
    """清理测试数据"""
    db_path = r"C:\Users\imeep\AppData\Roaming\mixvideo\mixvideoV2.db"
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    try:
        cursor.execute("DELETE FROM uni_comfyui_batch_item WHERE batch_id LIKE 'stats_test_%'")
        cursor.execute("DELETE FROM uni_comfyui_batch_task WHERE workflow_name = 'test_statistics'")
        cursor.execute("DELETE FROM uni_comfyui_task WHERE workflow_name = 'test_statistics'")
        
        conn.commit()
        print("✅ 测试数据清理完成")
    except Exception as e:
        conn.rollback()
        print(f"❌ 清理失败: {e}")
    finally:
        conn.close()

if __name__ == "__main__":
    print("=== 任务统计修复测试 ===")
    
    # 1. 创建测试数据
    if create_mixed_test_data():
        # 2. 验证数据库统计
        if verify_database_statistics():
            print("\n🎉 数据库统计验证通过!")
            
            # 3. 显示预期的API响应
            show_expected_api_response()
            
            print("\n📝 前端测试说明:")
            print("   现在可以在前端调用 get_task_statistics API")
            print("   应该能看到正确区分的单个任务和批量任务统计")
            print("   - single_task_stats: 单个任务的统计信息")
            print("   - batch_task_stats: 批量任务的统计信息")
        else:
            print("\n❌ 数据库统计验证失败!")
    else:
        print("\n❌ 测试数据创建失败!")
    
    # 询问是否清理测试数据
    print("\n是否清理测试数据? (y/n): ", end="")
    # cleanup_test_data()  # 暂时注释掉，让用户手动决定