#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
任务执行器模块
提供任务执行的具体实现
"""

import json
import uuid
import logging
import asyncio
from datetime import datetime
from typing import Dict, List, Any, Optional
from sqlalchemy import select, update
from sqlalchemy.ext.asyncio import AsyncSession

from data.models.taskdef import VWEDTaskDef
from data.models.taskrecord import VWEDTaskRecord
from data.models.blockrecord import VWEDBlockRecord
from data.session import get_async_session
from services.execution.task_context import TaskContext
from services.execution.block_executor import BlockExecutor
from data.enum.task_record_enum import TaskStatus
from utils.logger import get_logger

# 获取日志记录器
logger = get_logger("services.execution.task_executor")

class TaskExecutor:
    """
    任务执行器类
    负责执行具体的任务，处理任务状态管理、任务块执行和错误处理
    """
    
    def __init__(self, task_record_id: str):
        """
        初始化任务执行器
        
        Args:
            task_record_id: 任务记录ID
        """
        self.task_record_id = task_record_id
        self.task_record = None
        self.task_def = None
        self.block_executor = None
        self.task_context = None
        self.is_running = False
        self.error_message = None
        self.timeout = 3600*10  # 默认超时时间：10小时（秒）
        self.is_canceled = False
    
    def set_timeout(self, timeout_seconds: int) -> None:
        """
        设置任务执行超时时间
        
        Args:
            timeout_seconds: 超时时间（秒）
        """
        if timeout_seconds > 0:
            self.timeout = timeout_seconds
            logger.debug(f"任务 {self.task_record_id} 设置超时时间: {timeout_seconds}秒")
    
    async def initialize(self) -> bool:
        """
        初始化任务执行器，加载任务记录和任务定义
        
        Returns:
            bool: 初始化是否成功
        """
        try:
            async with get_async_session() as session:
                # 查询任务记录
                result = await session.execute(
                    select(VWEDTaskRecord).where(VWEDTaskRecord.id == self.task_record_id)
                )
                self.task_record = result.scalars().first()
                
                if not self.task_record:
                    logger.error(f"任务记录不存在: {self.task_record_id}")
                    return False
                
                # 查询任务定义
                result = await session.execute(
                    select(VWEDTaskDef).where(VWEDTaskDef.id == self.task_record.def_id)
                )
                self.task_def = result.scalars().first()
                
                if not self.task_def:
                    logger.error(f"任务定义不存在: {self.task_record.def_id}")
                    await self._update_task_status(session, TaskStatus.FAILED, "任务定义不存在", task_detail=self.task_def.detail)
                    return False
                                
                # 解析任务输入参数
                input_params = json.loads(self.task_record.input_params) if self.task_record.input_params else {}
                self.task_def.map_id
                
                # 创建任务上下文
                self.task_context = TaskContext(
                    task_record_id=self.task_record_id,
                    task_def_id=self.task_def.id,
                    input_params=input_params,
                    variables={},
                    token=self.task_def.user_token,
                    map_id=self.task_def.map_id
                )
                
                # 创建块执行器
                self.block_executor = BlockExecutor(self.task_context)
                
                # 设置任务记录首次执行时间
                if not self.task_record.first_executor_time:
                    self.task_record.first_executor_time = datetime.now()
                    await session.commit()
                
                return True
                
        except Exception as e:
            logger.error(f"初始化任务执行器失败: {str(e)}")
            self.error_message = f"初始化任务执行器失败: {str(e)}"
            return False
    
    def check_task_def_select_agv(self, detail_json: str) -> int:
        """
        检查任务定义详情中是否包含SELECT_AGV块类型
        
        Args:
            detail_json: 任务定义详情JSON字符串
            
        Returns:
            int: 存在返回1，不存在返回0
        """
        from services.execution.handlers.model.block_name import RobotBlockName
        try:
            # 解析JSON字符串
            if not detail_json:
                return 0
                
            detail = json.loads(detail_json)
            if not detail:
                return 0

            # 选择AGV块类型常量
            SELECT_AGV_TYPE = RobotBlockName.SELECT_AGV  # "CSelectAgvBp"
            
            # 检查根块
            root_block = detail.get("rootBlock")
            if not root_block:
                return 0
                
            # 递归检查块类型
            def check_block(block):
                # 检查当前块类型
                if block.get("blockType") == SELECT_AGV_TYPE:
                    return True
                    
                # 递归检查子块
                children = block.get("children", {})
                for child_key, child_list in children.items():
                    if isinstance(child_list, list):
                        for child in child_list:
                            if check_block(child):
                                return True
                                
                return False
                
            # 从根块开始检查
            if check_block(root_block):
                return 1
                
            return 0
            
        except json.JSONDecodeError:
            logger.error(f"解析任务定义详情JSON失败: {detail_json[:100]}...")
            return 0
        except Exception as e:
            logger.error(f"检查任务定义中SELECT_AGV类型异常: {str(e)}")
            return 0
        
    async def execute(self) -> Dict[str, Any]:
        """
        执行任务
        
        Returns:
            Dict[str, Any]: 任务执行结果
        """
        try:
            # 如果还没有初始化，先进行初始化
            if not self.task_context or not self.block_executor:
                if not await self.initialize():
                    return {
                        "success": False,
                        "message": self.error_message or "任务初始化失败",
                        "status": TaskStatus.FAILED,
                        "taskRecordId": self.task_record_id
                    }
            
            self.is_running = True
            
            # 开始计时
            start_time = datetime.now()
            
            # 获取根块
            task_detail = json.loads(self.task_def.detail)
            root_block = task_detail.get("rootBlock", {})
            # 更新任务状态为执行中
            async with get_async_session() as session:
                await self._update_task_status(session, TaskStatus.RUNNING, "任务执行中", task_detail=self.task_def.detail)
            # 执行根块，增加超时控制
            try:
                # 直接执行块，不添加超时控制
                result = await self.block_executor.execute_block(root_block)
            except Exception as e:
                logger.error(f"任务 {self.task_record_id} 执行异常: {str(e)}")
                # 更新任务状态为失败
                async with get_async_session() as session:
                    await self._update_task_status(
                        session, 
                        TaskStatus.FAILED, 
                        f"任务执行异常: {str(e)}",
                        task_detail=self.task_def.detail
                    )
                return {
                    "success": False,
                    "message": f"任务执行异常: {str(e)}",
                    "status": TaskStatus.FAILED,
                    "taskRecordId": self.task_record_id,
                    "error": {"type": "exception", "message": str(e)}
                }
            
            # 计算执行时间
            end_time = datetime.now()
            executor_time = int((end_time - start_time).total_seconds() * 1000)  # 转换为毫秒
            
            # 更新任务状态为完成
            async with get_async_session() as session:
                # print("self.is_canceled", self.is_canceled, "-------------------------------------------")
                if result.get("success", False) and not self.is_canceled:
                    await self._update_task_status(
                        session, 
                        TaskStatus.COMPLETED, 
                        "任务执行成功", 
                        executor_time,
                        task_detail=self.task_def.detail
                    )
                    
                    # 返回成功结果
                    return {
                        "success": True,
                        "message": "任务执行成功",
                        "status": TaskStatus.COMPLETED,
                        "taskRecordId": self.task_record_id,
                        "executorTime": executor_time,
                        "result": result.get("output", {})
                    }
                elif result.get("success", False) and self.is_canceled:
                    await self._update_task_status(
                        session, 
                        TaskStatus.CANCELED, 
                        "任务取消", 
                        executor_time,
                        task_detail=self.task_def.detail
                    )
                    return {
                        "success": True,
                        "message": "任务取消",
                        "status": TaskStatus.CANCELED,
                        "taskRecordId": self.task_record_id,
                    }
                else:
                    error_msg = result.get("message", "任务执行失败")
                    
                    if result.get("is_canceled", False):
                        await self._update_task_status(session, TaskStatus.CANCELED, error_msg, executor_time, task_detail=self.task_def.detail)
                    else:
                        # 更新任务状态为失败
                        await self._update_task_status(session, TaskStatus.FAILED, error_msg, executor_time, task_detail=self.task_def.detail)
                    
                    # 返回失败结果
                    return {
                        "success": False,
                        "message": error_msg,
                        "status": TaskStatus.FAILED,
                        "taskRecordId": self.task_record_id,
                        "executorTime": executor_time,
                        "error": result.get("error", {})
                    }
                    
        except Exception as e:
            logger.error(f"任务执行异常: {str(e)}")
            
            # 更新任务状态为异常
            async with get_async_session() as session:
                await self._update_task_status(session, TaskStatus.FAILED, f"任务执行异常: {str(e)}", task_detail=self.task_def.detail)
            
            return {
                "success": False,
                "message": f"任务执行异常: {str(e)}",
                "status": TaskStatus.FAILED,
                "taskRecordId": self.task_record_id
            }
        finally:
            self.is_running = False
    
    async def cancel(self) -> Dict[str, Any]:
        """
        取消正在执行的任务
        
        Returns:
            Dict[str, Any]: 取消结果
        """
        if not self.is_running:
            return {
                "success": False,
                "message": "任务未在执行中，无法取消",
                "taskRecordId": self.task_record_id
            }
        
        # 更新任务状态为取消
        async with get_async_session() as session:
            await self._update_task_status(session, TaskStatus.CANCELED, "任务被取消", task_detail=self.task_def.detail)
        
        # 尝试取消块执行器
        if self.block_executor:
            self.block_executor.cancel()
        
        self.is_running = False
        self.is_canceled = True
        return {
            "success": True,
            "message": "任务已取消",
            "taskRecordId": self.task_record_id
        }
    
    async def _update_task_status(
        self, 
        session: AsyncSession, 
        status: int, 
        reason: str = None,
        executor_time: int = None,
        task_detail: str = None
    ) -> None:
        """
        更新任务状态
        
        Args:
            session: 数据库会话
            status: 任务状态码
            reason: 状态原因
            executor_time: 执行时间(毫秒)
        """
        try:
            from services.sync_service import set_task_in_progress, set_task_completed, set_task_terminated, set_task_failed
            # print(status, )
            # print("=======================================================")
            # print(self.check_task_def_select_agv(task_detail), )

            if not self.check_task_def_select_agv(task_detail):
                if status == TaskStatus.RUNNING:
                    await set_task_in_progress(self.task_record_id, self.task_def.user_token)
                elif status == TaskStatus.COMPLETED:
                    await set_task_completed(self.task_record_id, self.task_def.user_token)
                elif status == TaskStatus.FAILED:
                    await set_task_failed(self.task_record_id, self.task_def.user_token)
                elif status == TaskStatus.CANCELED:
                    await set_task_terminated(self.task_record_id, self.task_def.user_token)
            # else:

            if status in [TaskStatus.FAILED, TaskStatus.CANCELED, TaskStatus.COMPLETED]:
                update_values = {
                    "status": status,
                    "allow_restart_same_location": True
                }
            else:
                update_values = {
                    "status": status
                }
            
            if reason:
                update_values["ended_reason"] = reason
            
            if executor_time:
                update_values["executor_time"] = executor_time
            
            # 如果状态为完成、失败或取消，设置结束时间
            if status in [TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELED]:
                update_values["ended_on"] = datetime.now()
            
            # 更新任务记录
            await session.execute(
                update(VWEDTaskRecord)
                .where(VWEDTaskRecord.id == self.task_record_id)
                .values(**update_values)
            )
            
            await session.commit()
                        
        except Exception as e:
            logger.error(f"更新任务状态失败: {str(e)}")
            await session.rollback()
            raise 