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

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

import json
import uuid
import logging
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.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
from services.sync_service import set_task_in_progress, set_task_completed, set_task_failed, set_task_terminated, \
    refresh_token_if_needed
from utils.alert_sync import sync_error_alert
from utils.background_task_manager import create_background_task

# 获取日志记录器
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
        self.is_error = 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
                token = await refresh_token_if_needed()
                # 创建任务上下文
                self.task_context = TaskContext(
                    task_record_id=self.task_record_id,
                    task_def_id=self.task_def.id,
                    input_params=input_params,
                    variables={},
                    token=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_record.task_def_detail读取(支持动态更新,如换机器人)
            # 如果task_record.task_def_detail存在,使用它;否则使用task_def.detail
            if self.task_record.task_def_detail:
                task_detail_str = self.task_record.task_def_detail
                logger.info(f"任务 {self.task_record_id} 使用task_record.task_def_detail(已动态更新)")
            else:
                task_detail_str = self.task_def.detail
                logger.info(f"任务 {self.task_record_id} 使用task_def.detail(原始定义)")

            task_detail = json.loads(task_detail_str)
            root_block = task_detail.get("rootBlock", {})
            release_sites = self.task_def.release_sites
            print("root_block:::::::::::", root_block)
            # 更新任务状态为执行中
            async with get_async_session() as session:
                await self._update_task_status(session, TaskStatus.RUNNING, "任务执行中", task_detail=task_detail_str)
            # 执行根块，增加超时控制
            site_list = []
            try:
                # 直接执行块，不添加超时控制
                result = await self.block_executor.execute_block(root_block)
                site_list = [v.get("siteId") for k, v in dict(self.task_context.block_outputs).items()]
            except Exception as e:
                logger.error(f"任务 {self.task_record_id} 执行异常: {str(e)}")

                # 推送错误告警到主系统
                sync_error_alert(
                    logger_name="services.execution.task_executor",
                    error_message=f"任务 {self.task_record_id} 执行异常: {str(e)}",
                    level=logging.ERROR,
                    context_info={
                        "task_id": self.task_record_id,
                        "filename": "task_executor.py",
                        "lineno": 221,
                        "funcName": "execute"
                    }
                )

                # 更新任务状态为失败
                async with get_async_session() as session:
                    await self._update_task_status(
                        session,
                        TaskStatus.FAILED,
                        f"任务执行异常: {str(e)}",
                        task_detail=self.task_def.detail
                    )

                # 根据release_sites配置决定是否释放库位

                if release_sites == 1 and site_list:
                    logger.info(f"具备自动释放库位锁定条件！释放一下库位：{site_list}")
                    await self._release_storage_locations(site_list)

                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:
                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
                    )
                    # 根据release_sites配置决定是否释放库位
                    if release_sites == 1 and site_list:
                        logger.info(f"具备自动释放库位锁定条件！释放一下库位：{site_list}")
                        await self._release_storage_locations(site_list)
                    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:
                        # 推送错误告警到主系统
                        sync_error_alert(
                            logger_name="services.execution.task_executor",
                            error_message=f"任务 {self.task_record_id} 执行失败: {error_msg}",
                            level=logging.ERROR,
                            context_info={
                                "task_id": self.task_record_id,
                                "filename": "task_executor.py",
                                "lineno": 310,
                                "funcName": "execute",
                                "error_detail": result.get("error", {})
                            }
                        )

                        # 更新任务状态为失败
                        await self._update_task_status(session, TaskStatus.FAILED, error_msg, executor_time, task_detail=self.task_def.detail)

                    # 根据release_sites配置决定是否释放库位
                    if release_sites == 1 and site_list:
                        logger.info(f"具备自动释放库位锁定条件！释放一下库位：{site_list}")
                        await self._release_storage_locations(site_list)
                    
                    # 返回失败结果
                    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)
            
            # 如果有库位信息需要释放
            site_list = []
            try:
                site_list = [v.get("siteId") for k, v in dict(self.task_context.block_outputs).items() if v.get("siteId")]
            except:
                pass
            # 推送错误告警到主系统
            sync_error_alert(
                logger_name="services.execution.task_executor",
                error_message=f"任务 {self.task_record_id} 执行失败: {error_msg}",
                level=logging.ERROR,
                context_info={
                    "task_id": self.task_record_id,
                    "filename": "task_executor.py",
                    "lineno": 310,
                    "funcName": "execute",
                    "error_detail": result.get("error", {})
                }
            )
            # 根据release_sites配置决定是否释放库位
            if hasattr(self.task_def, 'release_sites') and self.task_def.release_sites == 1 and site_list:
                logger.info(f"具备自动释放库位锁定条件！释放一下库位：{site_list}")
                await self._release_storage_locations(site_list)
            
            return {
                "success": False,
                "message": f"任务执行异常: {str(e)}",
                "status": TaskStatus.FAILED,
                "taskRecordId": self.task_record_id
            }
        finally:
            self.is_running = False

            # 清理任务执行器，释放内存
            try:
                if self.task_context:
                    if hasattr(self.task_context, 'cleanup'):
                        self.task_context.cleanup()
                    self.task_context = None

                if self.block_executor:
                    self.block_executor.task_context = None
                    self.block_executor = None

                # 清空任务记录引用
                self.task_record = None
                self.task_def = None

                logger.debug(f"任务执行器 {self.task_record_id} 已清理")
            except Exception as cleanup_error:
                logger.error(f"清理任务执行器失败: {str(cleanup_error)}")
    
    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 fail(self, failure_reason: str = "任务设置为失败") -> Dict[str, Any]:
        """
        将任务设置为失败状态，并将失败原因设置到所有块的详情内

        Args:
            failure_reason: 失败原因

        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.FAILED, failure_reason, task_detail=self.task_def.detail)
        
        # 根据release_sites配置决定是否释放库位
        if hasattr(self.task_def, 'release_sites') and self.task_def.release_sites == 1:
            try:
                site_list = [v.get("siteId") for k, v in dict(self.task_context.block_outputs).items() if v.get("siteId")]
                if site_list:
                    await self._release_storage_locations(site_list)
            except Exception as e:
                logger.error(f"任务 {self.task_record_id} 释放库位时出错: {str(e)}")
        
        # 尝试将块执行器设置为失败
        if self.block_executor:
            self.block_executor.fail(failure_reason)
        
        self.is_running = False
        self.is_error = True

        self.error_message = failure_reason
        return {
            "success": True,
            "message": "任务已设置为失败",
            "taskRecordId": self.task_record_id,
            "failureReason": failure_reason
        }
    
    async def _release_storage_locations(self, site_list: List[str]) -> None:
        """
        释放任务执行过程中选择的库位
        
        Args:
            site_list: 需要释放的库位ID列表
        """
        if not site_list:
            logger.debug(f"任务 {self.task_record_id} 没有需要释放的库位")
            return
            
        try:
            from data.session import get_db
            from services.operate_point_service import OperatePointService
            from routes.model.operate_point_model import StorageLocationStatusUpdateRequest, StorageLocationActionEnum
            
            # 获取数据库会话
            db = next(get_db())
            
            success_count = 0
            failed_count = 0
            
            # 逐个释放库位
            for site_id in site_list:
                if not site_id:
                    continue
                    
                try:
                    # 创建释放库位请求
                    request = StorageLocationStatusUpdateRequest(
                        layer_name=site_id,
                        action=StorageLocationActionEnum.UNLOCK,
                        reason=f"任务失败自动释放库位 - 任务ID: {self.task_record_id}"
                    )
                    
                    # 调用服务释放库位
                    result = OperatePointService.update_storage_location_status(db=db, request=request)
                    
                    if result.success:
                        success_count += 1
                        logger.info(f"任务 {self.task_record_id} 成功释放库位: {site_id}")
                    else:
                        failed_count += 1
                        logger.warning(f"任务 {self.task_record_id} 释放库位失败: {site_id}, 原因: {result.message}")
                        
                except Exception as e:
                    failed_count += 1
                    logger.error(f"任务 {self.task_record_id} 释放库位异常: {site_id}, 错误: {str(e)}")
            
            logger.info(f"任务 {self.task_record_id} 库位释放完成: 成功 {success_count} 个, 失败 {failed_count} 个")
            
        except Exception as e:
            logger.error(f"任务 {self.task_record_id} 批量释放库位异常: {str(e)}")
        finally:
            try:
                db.close()
            except:
                pass
    
    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: 执行时间(毫秒)
            task_detail: 任务详情
        """
        try:
            # 先完成数据库更新，立即释放数据库连接
            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()

            # 数据库更新完成后，异步处理外部API调用（不阻塞当前流程）- 使用后台任务管理器
            if not self.is_error:
                create_background_task(
                    self._sync_external_api_async(status, task_detail),
                    name=f"sync_api_{self.task_record_id}_{status}",
                    context=f"TaskRecordId={self.task_record_id}, Status={status}"
                )
                        
        except Exception as e:
            logger.error(f"更新任务状态失败: {str(e)}")
            await session.rollback()
            raise
            
    async def _sync_external_api_async(self, status: int, task_detail: str = None) -> None:
        """
        异步处理外部API同步，不阻塞主流程
        
        Args:
            status: 任务状态码
            task_detail: 任务详情
        """
        try:
            token = await refresh_token_if_needed()
            if not self.check_task_def_select_agv(task_detail):
                if status == TaskStatus.RUNNING:
                    await set_task_in_progress(self.task_record_id, token)
                elif status == TaskStatus.COMPLETED:
                    await set_task_completed(self.task_record_id, token)
                elif status == TaskStatus.FAILED:
                    await set_task_failed(self.task_record_id, token)
                elif status == TaskStatus.CANCELED:
                    await set_task_terminated(self.task_record_id, token)
            else:
                if status == TaskStatus.FAILED:
                    await set_task_failed(self.task_record_id, token)
                    
        except Exception as e:
            logger.error(f"外部API同步失败: {str(e)}")
            # 外部API失败不应该影响主任务流程，只记录日志 