#!/usr/bin/env python3
"""
CTO总控官调度系统 - 完整版
包含核心调度器、工作流管理器和系统监控
"""

import sqlite3
import json
import os
import time
import threading
import logging
from datetime import datetime, timedelta
from pathlib import Path
from enum import Enum
from dataclasses import dataclass, asdict
from typing import Dict, List, Optional, Any, Tuple
import queue

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('/root/data/disk/system/logs/cto/scheduler.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger('cto_scheduler')

# ============ 枚举定义 ============

class TaskStatus(Enum):
    """任务状态枚举"""
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    RETRYING = "retrying"

class AgentStatus(Enum):
    """Agent状态枚举"""
    IDLE = "idle"
    BUSY = "busy"
    ERROR = "error"
    OFFLINE = "offline"

class TaskPriority(Enum):
    """任务优先级枚举"""
    CRITICAL = 100  # 紧急项目，立即执行
    HIGH = 75      # 高优先级
    NORMAL = 50    # 普通优先级（默认）
    LOW = 25       # 低优先级

# ============ 数据库管理器 ============

class DatabaseManager:
    """数据库管理器"""
    
    def __init__(self):
        self.db_paths = {
            'projects': '/root/data/disk/system/database/projects.db',
            'agents': '/root/data/disk/system/database/agents.db'
        }
    
    def connect(self, db_name: str) -> sqlite3.Connection:
        """连接到指定数据库"""
        db_path = self.db_paths.get(db_name)
        if not db_path:
            raise ValueError(f"未知数据库: {db_name}")
        
        conn = sqlite3.connect(db_path)
        conn.row_factory = sqlite3.Row
        return conn
    
    def get_project(self, project_id: str) -> Optional[Dict]:
        """获取项目信息"""
        try:
            conn = self.connect('projects')
            cursor = conn.cursor()
            
            cursor.execute('''
                SELECT * FROM projects WHERE project_id = ?
            ''', (project_id,))
            
            row = cursor.fetchone()
            conn.close()
            
            if row:
                return dict(row)
            return None
            
        except Exception as e:
            logger.error(f"获取项目信息失败 {project_id}: {e}")
            return None
    
    def create_task(self, task_data: Dict) -> bool:
        """创建新任务"""
        try:
            conn = self.connect('agents')
            cursor = conn.cursor()
            
            cursor.execute('''
                INSERT INTO agent_tasks 
                (task_id, project_id, agent_name, task_type, priority, status, 
                 input_data, expected_outputs, dependencies, retry_count, max_retries,
                 created_at, timeout_seconds)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            ''', (
                task_data['task_id'], task_data['project_id'], task_data['agent_name'],
                task_data['task_type'], task_data.get('priority', 50),
                task_data.get('status', 'pending'),
                json.dumps(task_data.get('input_data', {}), ensure_ascii=False),
                json.dumps(task_data.get('expected_outputs', []), ensure_ascii=False),
                json.dumps(task_data.get('dependencies', []), ensure_ascii=False),
                task_data.get('retry_count', 0),
                task_data.get('max_retries', 3),
                task_data.get('created_at', datetime.now().isoformat()),
                task_data.get('timeout_seconds', 1800)
            ))
            
            conn.commit()
            conn.close()
            logger.info(f"创建任务成功: {task_data['task_id']}")
            return True
            
        except Exception as e:
            logger.error(f"创建任务失败 {task_data['task_id']}: {e}")
            return False
    
    def update_task_status(self, task_id: str, status: str, **kwargs) -> bool:
        """更新任务状态"""
        try:
            conn = self.connect('agents')
            cursor = conn.cursor()
            
            if status == 'running' and 'started_at' in kwargs:
                cursor.execute('''
                    UPDATE agent_tasks 
                    SET status = ?, started_at = ?, error_message = ?
                    WHERE task_id = ?
                ''', (status, kwargs['started_at'], kwargs.get('error_message', ''), task_id))
            elif status == 'completed' and 'completed_at' in kwargs:
                cursor.execute('''
                    UPDATE agent_tasks 
                    SET status = ?, completed_at = ?, error_message = ?
                    WHERE task_id = ?
                ''', (status, kwargs['completed_at'], kwargs.get('error_message', ''), task_id))
            else:
                cursor.execute('''
                    UPDATE agent_tasks 
                    SET status = ?, error_message = ?
                    WHERE task_id = ?
                ''', (status, kwargs.get('error_message', ''), task_id))
            
            conn.commit()
            conn.close()
            logger.info(f"更新任务状态: {task_id} -> {status}")
            return True
            
        except Exception as e:
            logger.error(f"更新任务状态失败 {task_id}: {e}")
            return False
    
    def get_agent_state(self, agent_name: str) -> Optional[Dict]:
        """获取Agent状态"""
        try:
            conn = self.connect('agents')
            cursor = conn.cursor()
            
            cursor.execute('''
                SELECT * FROM agent_status WHERE agent_name = ?
            ''', (agent_name,))
            
            row = cursor.fetchone()
            conn.close()
            
            if row:
                return dict(row)
            return None
            
        except Exception as e:
            logger.error(f"获取Agent状态失败 {agent_name}: {e}")
            return None
    
    def update_agent_state(self, agent_name: str, state_data: Dict) -> bool:
        """更新Agent状态"""
        try:
            conn = self.connect('agents')
            cursor = conn.cursor()
            
            cursor.execute('''
                INSERT OR REPLACE INTO agent_status 
                (agent_name, status, current_task_id, last_heartbeat, performance_metrics)
                VALUES (?, ?, ?, ?, ?)
            ''', (
                agent_name,
                state_data.get('status', 'idle'),
                state_data.get('current_task_id'),
                state_data.get('last_heartbeat', datetime.now().isoformat()),
                json.dumps(state_data.get('performance_metrics', {}), ensure_ascii=False)
            ))
            
            conn.commit()
            conn.close()
            return True
            
        except Exception as e:
            logger.error(f"更新Agent状态失败 {agent_name}: {e}")
            return False
    
    def get_pending_tasks(self, limit: int = 100) -> List[Dict]:
        """获取待处理任务"""
        try:
            conn = self.connect('agents')
            cursor = conn.cursor()
            
            cursor.execute('''
                SELECT * FROM agent_tasks 
                WHERE status = 'pending'
                ORDER BY priority DESC, created_at ASC
                LIMIT ?
            ''', (limit,))
            
            tasks = [dict(row) for row in cursor.fetchall()]
            conn.close()
            return tasks
            
        except Exception as e:
            logger.error(f"获取待处理任务失败: {e}")
            return []
    
    def update_project_status(self, project_id: str, status: str, workflow_step: str = None) -> bool:
        """更新项目状态"""
        try:
            conn = self.connect('projects')
            cursor = conn.cursor()
            
            if workflow_step:
                cursor.execute('''
                    UPDATE projects 
                    SET status = ?, workflow_step = ?, updated_at = CURRENT_TIMESTAMP
                    WHERE project_id = ?
                ''', (status, workflow_step, project_id))
            else:
                cursor.execute('''
                    UPDATE projects 
                    SET status = ?, updated_at = CURRENT_TIMESTAMP
                    WHERE project_id = ?
                ''', (status, project_id))
            
            conn.commit()
            conn.close()
            logger.info(f"更新项目状态: {project_id} -> {status}")
            return True
            
        except Exception as e:
            logger.error(f"更新项目状态失败 {project_id}: {e}")
            return False

# ============ 任务调度器 ============

class TaskScheduler:
    """任务调度器"""
    
    def __init__(self, db_manager: DatabaseManager):
        self.db = db_manager
        self.task_queue = queue.PriorityQueue()
        self.running_tasks: Dict[str, Dict] = {}
        
        # 调度配置
        self.config = {
            'max_parallel_tasks': 10,
            'agent_concurrency': {
                'collector': 2,
                'writer': 1,
                'producer': 1,
                'publisher': 1,
                'analyzer': 2
            },
            'heartbeat_timeout': {
                'collector': 300,    # 5分钟
                'writer': 600,       # 10分钟
                'producer': 900,     # 15分钟
                'publisher': 300,    # 5分钟
                'analyzer': 300      # 5分钟
            }
        }
        
        self.running = False
        self.scheduler_thread = None
        self.monitor_thread = None
    
    def start(self):
        """启动调度器"""
        if self.running:
            logger.warning("调度器已经在运行")
            return
        
        self.running = True
        self.scheduler_thread = threading.Thread(target=self._scheduler_loop, daemon=True)
        self.monitor_thread = threading.Thread(target=self._monitor_loop, daemon=True)
        
        self.scheduler_thread.start()
        self.monitor_thread.start()
        
        logger.info("CTO调度系统启动成功")
    
    def stop(self):
        """停止调度器"""
        self.running = False
        if self.scheduler_thread:
            self.scheduler_thread.join(timeout=5)
        if self.monitor_thread:
            self.monitor_thread.join(timeout=5)
        
        logger.info("CTO调度系统停止")
    
    def schedule_task(self, task_data: Dict) -> bool:
        """调度新任务"""
        try:
            # 计算优先级
            priority = self._calculate_priority(task_data)
            
            # 保存到数据库
            if not self.db.create_task(task_data):
                return False
            
            # 添加到内存队列
            self.task_queue.put((-priority, task_data['task_id'], task_data))
            logger.info(f"任务已调度: {task_data['task_id']} (优先级: {priority})")
            
            return True
            
        except Exception as e:
            logger.error(f"调度任务失败 {task_data.get('task_id')}: {e}")
            return False
    
    def _calculate_priority(self, task_data: Dict) -> int:
        """计算任务优先级"""
        priority = task_data.get('priority', 50)
        
        # 项目优先级加成
        project = self.db.get_project(task_data['project_id'])
        if project:
            if project.get('priority') == 'urgent':
                priority += 40
            elif project.get('priority') == 'high':
                priority += 20
        
        # Agent可用性加成
        agent_state = self.db.get_agent_state(task_data['agent_name'])
        if agent_state and agent_state.get('status') == 'idle':
            priority += 20
        
        return min(priority, 100)
    
    def _scheduler_loop(self):
        """调度器主循环"""
        logger.info("调度器主循环启动")
        
        while self.running:
            try:
                # 检查队列
                if not self.task_queue.empty():
                    # 获取最高优先级任务
                    _, task_id, task_data = self.task_queue.get()
                    
                    # 检查是否可以执行
                    if self._can_execute_task(task_data):
                        # 执行任务
                        self._execute_task(task_data)
                    else:
                        # 放回队列
                        priority = self._calculate_priority(task_data)
                        self.task_queue.put((-priority, task_id, task_data))
                
                time.sleep(1)
                
            except Exception as e:
                logger.error(f"调度器循环异常: {e}")
                time.sleep(5)
    
    def _can_execute_task(self, task_data: Dict) -> bool:
        """检查任务是否可以执行"""
        # 检查Agent状态
        agent_state = self.db.get_agent_state(task_data['agent_name'])
        if not agent_state or agent_state.get('status') != 'idle':
            return False
        
        # 检查并发数
        agent_name = task_data['agent_name']
        current_concurrent = sum(1 for t in self.running_tasks.values() 
                                if t.get('agent_name') == agent_name)
        max_concurrent = self.config['agent_concurrency'].get(agent_name, 1)
        
        if current_concurrent >= max_concurrent:
            return False
        
        # 检查总任务数
        if len(self.running_tasks) >= self.config['max_parallel_tasks']:
            return False
        
        return True
    
    def _execute_task(self, task_data: Dict):
        """执行任务"""
        try:
            task_id = task_data['task_id']
            
            # 更新任务状态
            started_at = datetime.now().isoformat()
            if not self.db.update_task_status(task_id, 'running', started_at=started_at):
                return
            
            # 更新Agent状态
            agent_state = {
                'status': 'busy',
                'current_task_id': task_id,
                'last_heartbeat': started_at
            }
            self.db.update_agent_state(task_data['agent_name'], agent_state)
            
            # 添加到运行中任务列表
            self.running_tasks[task_id] = task_data
            
            # 启动任务线程
            thread = threading.Thread(
                target=self._execute_task_thread,
                args=(task_data,),
                daemon=True
            )
            thread.start()
            
            logger.info(f"任务开始执行: {task_id}")
            
        except Exception as e:
            logger.error(f"执行任务失败 {task_data.get('task_id')}: {e}")
    
    def _execute_task_thread(self, task_data: Dict):
        """任务执行线程"""
        task_id = task_data['task_id']
        agent_name = task_data['agent_name']
        
        try:
            logger.info(f"执行任务 {task_id}: {task_data['task_type']}")
            
            # 这里应该调用实际的Agent执行
            # 目前先模拟执行
            time.sleep(10)  # 模拟任务执行
            
            # 任务完成
            completed_at = datetime.now().isoformat()
            self.db.update_task_status(task_id, 'completed', completed_at=completed_at)
            
            # 更新Agent状态
            agent_state = {
                'status': 'idle',
                'current_task_id': None,
                'last_heartbeat': completed_at
            }
            self.db.update_agent_state(agent_name, agent_state)
            
            # 从运行中任务列表移除
            if task_id in self.running_tasks:
                del self.running_tasks[task_id]
            
            # 更新项目状态
            self._update_project_workflow(task_data['project_id'], task_data['task_type'])
            
            logger.info(f"任务完成: {task_id}")
            
        except Exception as e:
            logger.error(f"任务执行异常 {task_id}: {e}")
            self._handle_task_failure(task_data, str(e))
    
    def _handle_task_failure(self, task_data: Dict, error_message: str):
        """处理任务失败"""
        task_id = task_data['task_id']
        retry_count = task_data.get('retry_count', 0)
        max_retries = task_data.get('max_retries', 3)
        
        if retry_count < max_retries:
            # 重试
            task_data['retry_count'] = retry_count + 1
            task_data['priority'] = max(task_data.get('priority', 50) - 10, 25)
            
            logger.info(f"任务重试 {task_id}: 第{task_data['retry_count']}次")
            self.db.update_task_status(task_id, 'retrying', 
                                      error_message=f"重试第{task_data['retry_count']}次")
            
            # 重新调度
            self.schedule_task(task_data)
            
        else:
            # 失败
            logger.error(f"任务失败 {task_id}: 重试次数用尽")
            self.db.update_task_status(task_id, 'failed', 
                                      error_message=f"重试{max_retries}次后失败")
            
            # 更新Agent状态
            agent_state = {
                'status': 'idle',
                'current_task_id': None,
                'last_heartbeat': datetime.now().isoformat()
            }
            self.db.update_agent_state(task_data['agent_name'], agent_state)
            
            # 从运行中任务列表移除
            if task_id in self.running_tasks:
                del self.running_tasks[task_id]
            
            # 发送通知
            self._send_failure_notification(task_data, error_message)
    
    def _update_project_workflow(self, project_id: str, task_type: str):
        """更新项目工作流状态"""
        stage_map = {
            'collect_hot_topics': ('collecting', 'collection'),
            'collect_materials': ('collecting', 'collection'),
            'create_concept': ('creating', 'creation'),
            'write_script': ('creating', 'creation'),
            'generate_images': ('producing', 'production'),
            'generate_video': ('producing', 'production'),
            'review_content': ('reviewing', 'review'),
            'publish_content': ('reviewing', 'review'),
            'analyze_performance': ('analyzing', 'analysis')
        }
        
        if task_type in stage_map:
            status, step = stage_map[task_type]
            self.db.update_project_status(project_id, status, step)
    
    def _send_failure_notification(self, task_data: Dict, error_message: str):
        """发送失败通知"""
        notification = {
            'timestamp': datetime.now().isoformat(),
            'task_id': task_data['task_id'],
            'project_id': task_data['project_id'],
            'agent_name': task_data['agent_name'],
            'error': error_message,
            'retry_count': task_data.get('retry_count', 0)
        }
        
        try:
            log_file = '/root/data/disk/system/logs/cto/notifications.log'
            with open(log_file, 'a') as f:
                f.write(json.dumps(notification, ensure_ascii=False) + '\n')
        except Exception as e:
            logger.error(f"记录通知失败: {e}")
    
    def _monitor_loop(self):
        """监控循环"""
        logger.info("监控循环启动")
        
        while self.running:
            try:
                self._check_task_timeouts()
                self._check_agent_heartbeats()
                time.sleep(30)
                
            except Exception as e:
                logger.error(f"监控循环异常: {e}")
                time.sleep(60)
    
    def _check_task_timeouts(self):
        """检查任务超时"""
        current_time = datetime.now()
        
        for task_id, task_data in list(self.running_tasks.items()):
            started_at_str = task_data.get('started_at')
            if started_at_str:
                try:
                    started_at = datetime.fromisoformat(started_at_str.replace('Z', '+00:00'))
                    timeout = task_data.get('timeout_seconds', 1800)
                    
                    if (current_time - started_at).total_seconds() > timeout:
                        logger.warning(f"任务超时: {task_id}")
                        self._handle_task_failure(task_data, "任务超时")
                except:
                    pass
    
    def _check_agent_heartbeats(self):
        """检查Agent心跳"""
        agents = ['collector', 'writer', 'producer', 'publisher', 'analyzer']
        
        for agent_name in agents:
            agent_state = self.db.get_agent_state(agent_name)
            if agent_state and agent_state.get('last_heartbeat'):
                try:
                    last_heartbeat = datetime.fromisoformat(
                        agent_state['last_heartbeat'].replace('Z', '+00:00')
                    )
                    timeout = self.config['heartbeat_timeout'].get(agent_name, 300)
                    
                    if (datetime.now() - last_heartbeat).total_seconds() > timeout:
                        logger.warning(f"Agent心跳超时: {agent_name}")
                        
                        # 更新Agent状态
                        new_state = {
                            'status': 'error',
                            'last_heartbeat': agent_state['last_heartbeat']
                        }
                        self.db.update_agent_state(agent_name, new_state)
                        
                        # 处理该Agent的任务
                        self._handle_agent_failure(agent_name)
                except:
                    pass
    
    def _handle_agent_failure(self, agent_name: str):
        """处理Agent故障"""
        for task_id, task_data in list(self.running_tasks.items()):
            if task_data.get('agent_name') == agent_name:
                self._handle_task_failure(task_data, f"Agent {agent_name} 故障")
    
    def get_status(self) -> Dict:
        """获取调度器状态"""
        return {
            'running': self.running,
            'queue_size': self.task_queue.qsize(),
            'running_tasks_count': len(self.running_tasks),
            'agents': {
                agent: self.db.get_agent_state(agent) 
                for agent in ['collector', 'writer', 'producer', 'publisher', 'analyzer']
            }
        }

# ============ 工作流管理器 ============

class WorkflowManager:
    """工作流管理器"""
    
    def __init__(self, db_manager: DatabaseManager, scheduler: TaskScheduler):
        self.db = db_manager
        self.scheduler = scheduler
        
        self.workflows = {
            'standard': [
                {'agent': 'collector', 'task': 'collect_hot_topics', 'stage': 'collection'},
                {'agent': 'writer', 'task': 'create_concept', 'stage': 'creation'},
                {'agent': 'writer', 'task': 'write_script', 'stage': 'creation'},
                {'agent': 'producer', 'task': 'generate_images', 'stage': 'production'},
                {'agent': 'producer', 'task': 'generate_video', 'stage': 'production'},
                {'agent': 'publisher', 'task': 'review_content', 'stage': 'review'},
                {'agent': 'publisher', 'task': 'publish_content', 'stage': 'review'},
                {'agent': 'analyzer', 'task': 'analyze_performance', 'stage': 'analysis'}
            ],
            'quick': [
                {'agent': 'collector', 'task': 'collect_hot_topics', 'stage': 'collection'},
                {'agent': 'writer', 'task': 'create_concept', 'stage': 'creation'},
                {'agent': 'producer', 'task': 'generate_video', 'stage': 'production'}
            ]
        }
    
    def start_workflow(self, project_id: str, workflow_type: str = 'standard') -> bool:
        """启动工作流"""
        try:
            project = self.db.get_project(project_id)
            if not project:
                logger.error(f"项目不存在: {project_id}")
                return False
            
            workflow = self.workflows.get(workflow_type)
            if not workflow:
                logger.error(f"未知工作流类型: {workflow_type}")
                return False
            
            # 更新项目状态
            self.db.update_project_status(project_id, 'collecting', 'collection')
            
            # 创建工作流任务
            tasks = []
            dependencies = []
            
            for i, step in enumerate(workflow):
                task_id = f"{project_id}_{step['agent']}_{step['task']}_{int(time.time())}"
                
                task_data = {
                    'task_id': task_id,
                    'project_id': project_id,
                    'agent_name': step['agent'],
                    'task_type': step['task'],
                    'priority': 50,
                    'input_data': {
                        'project_id': project_id,
                        'project_name': project.get('name', ''),
                        'workflow_type': workflow_type,
                        'stage': step['stage']
                    },
                    'expected_outputs': [],
                    'dependencies': dependencies.copy(),
                    'timeout_seconds': 1800
                }
                
                tasks.append(task_data)
                dependencies = [task_id]  # 下一个任务依赖当前任务
            
            # 调度所有任务
            for task_data in tasks:
                if not self.scheduler.schedule_task(task_data):
                    logger.error(f"调度任务失败: {task_data['task_id']}")
                    return False
            
            logger.info(f"工作流启动成功: {project_id} ({workflow_type})")
            return True
            
        except Exception as e:
            logger.error(f"启动工作流失败 {project_id}: {e}")
            return False

# ============ 主应用程序 ============

class CTOApplication:
    """CTO主应用程序"""
    
    def __init__(self):
        self.db_manager = DatabaseManager()
        self.scheduler = TaskScheduler(self.db_manager)
        self.workflow_manager = WorkflowManager(self.db_manager, self.scheduler)
        self.running = False
    
    def start(self):
        """启动应用程序"""
        if self.running:
            logger.warning("应用程序已经在运行")
            return
        
        # 初始化Agent状态
        self._init_agent_states()
        
        # 启动调度器
        self.scheduler.start()
        
        self.running = True
        logger.info("CTO应用程序启动成功")
        
        # 保持主线程运行
        try:
            while self.running:
                time.sleep(1)
        except KeyboardInterrupt:
            self.stop()
    
    def stop(self):
        """停止应用程序"""
        self.running = False
        self.scheduler.stop()
        logger.info("CTO应用程序停止")
    
    def _init_agent_states(self):
        """初始化Agent状态"""
        agents = ['collector', 'writer', 'producer', 'publisher', 'analyzer']
        
        for agent_name in agents:
            state = self.db_manager.get_agent_state(agent_name)
            if not state:
                # 创建初始状态
                agent_state = {
                    'status': 'idle',
                    'current_task_id': None,
                    'last_heartbeat': datetime.now().isoformat(),
                    'performance_metrics': {
                        'tasks_completed': 0,
                        'tasks_failed': 0,
                        'success_rate': 1.0
                    }
                }
                self.db_manager.update_agent_state(agent_name, agent_state)
                logger.info(f"初始化Agent状态: {agent_name}")
    
    def create_test_project(self, project_name: str = "测试项目") -> str:
        """创建测试项目"""
        project_id = datetime.now().strftime("%Y%m%d%H%M%S")
        
        # 在数据库中创建项目
        try:
            conn = self.db_manager.connect('projects')
            cursor = conn.cursor()
            
            cursor.execute('''
                INSERT INTO projects 
                (project_id, name, description, status, created_by)
                VALUES (?, ?, ?, 'draft', 'cto_system')
            ''', (project_id, project_name, "CTO调度系统测试项目"))
            
            conn.commit()
            conn.close()
            
            logger.info(f"创建测试项目: {project_id} - {project_name}")
            return project_id
            
        except Exception as e:
            logger.error(f"创建测试项目失败: {e}")
            return None
    
    def start_test_workflow(self, project_id: str):
        """启动测试工作流"""
        if not project_id:
            logger.error("项目ID为空")
            return False
        
        return self.workflow_manager.start_workflow(project_id, 'quick')
    
    def get_system_status(self) -> Dict:
        """获取系统状态"""
        scheduler_status = self.scheduler.get_status()
        
        return {
            'application_running': self.running,
            'scheduler': scheduler_status,
            'database': {
                'projects': self._count_table('projects'),
                'agent_tasks': self._count_table('agents', 'agent_tasks'),
                'agent_status': self._count_table('agents', 'agent_status')
            }
        }
    
    def _count_table(self, db_name: str, table_name: str = None) -> int:
        """统计表记录数"""
        try:
            conn = self.db_manager.connect(db_name)
            cursor = conn.cursor()
            
            if table_name:
                cursor.execute(f'SELECT COUNT(*) as count FROM {table_name}')
            else:
                # 如果是projects数据库，默认统计projects表
                cursor.execute('SELECT COUNT(*) as count FROM projects')
            
            count = cursor.fetchone()['count']
            conn.close()
            return count
            
        except Exception as e:
            logger.error(f"统计表记录失败: {e}")
            return 0

# ============ 命令行接口 ============

def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description='CTO总控官调度系统')
    parser.add_argument('action', choices=['start', 'stop', 'status', 'test'],
                       help='执行的动作')
    parser.add_argument('--project', help='项目ID')
    parser.add_argument('--name', help='项目名称')
    
    args = parser.parse_args()
    
    app = CTOApplication()
    
    if args.action == 'start':
        print("启动CTO调度系统...")
        app.start()
        
    elif args.action == 'stop':
        print("停止CTO调度系统...")
        app.stop()
        
    elif args.action == 'status':
        status = app.get_system_status()
        print(json.dumps(status, indent=2, ensure_ascii=False))
        
    elif args.action == 'test':
        if args.name:
            project_id = app.create_test_project(args.name)
            if project_id:
                print(f"创建测试项目: {project_id}")
                if app.start_test_workflow(project_id):
                    print(f"启动测试工作流成功")
                else:
                    print(f"启动测试工作流失败")
        else:
            print("请使用 --name 指定项目名称")

if __name__ == '__main__':
    main()