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

"""
脚本输出捕获工具
用于捕获脚本执行过程中的标准输出和错误输出，并通过WebSocket推送给订阅的客户端
"""

import asyncio
import sys
from io import StringIO
from contextlib import contextmanager
from .script_websocket_service import get_websocket_manager
from utils.logger import get_logger

logger = get_logger("services.script_output_capture")


class ScriptOutputCapture:
    """脚本输出捕获类"""
    
    def __init__(self, script_id: str, level: str = "INFO", original_stream=None):
        self.script_id = script_id
        self.level = level
        self.websocket_manager = get_websocket_manager()
        self.buffer = StringIO()
        self.original_stream = original_stream  # 保留原始输出流
    
    def write(self, text):
        """写入文本，同时推送到WebSocket和原始输出流"""
        # 1. 写入原始输出流（如控制台）
        if self.original_stream:
            try:
                self.original_stream.write(text)
                self.original_stream.flush()
            except Exception as e:
                logger.debug(f"写入原始输出流失败: {e}")
        
        # 2. 推送到WebSocket（仅对非空内容）
        if text and text.strip():
            try:
                # 尝试在当前事件循环中创建任务
                asyncio.create_task(self._push_output(text.strip()))
            except RuntimeError:
                # 如果没有运行中的事件循环，记录日志但不报错
                logger.debug(f"无法推送输出到WebSocket，没有运行中的事件循环: [{self.script_id}] {text.strip()}")
        
        # 3. 写入缓冲区保留原有行为
        self.buffer.write(text)
    
    def flush(self):
        """刷新缓冲区"""
        self.buffer.flush()
        if self.original_stream:
            try:
                self.original_stream.flush()
            except Exception:
                pass
    
    async def _push_output(self, message: str):
        """推送输出到WebSocket"""
        try:
            await self.websocket_manager.broadcast_script_log(
                self.script_id, self.level, message
            )
        except Exception as e:
            logger.debug(f"推送脚本输出失败 [{self.script_id}]: {e}")


@contextmanager
def capture_script_output(script_id: str):
    """
    上下文管理器，用于捕获脚本执行过程中的标准输出和错误输出
    
    使用示例:
    with capture_script_output("script_123"):
        print("这条消息会被捕获并推送到WebSocket，同时显示在控制台")
        # 执行脚本代码
    """
    # 备份原始输出流
    original_stdout = sys.stdout
    original_stderr = sys.stderr
    
    # 创建捕获器，传入原始输出流以保持双重输出
    stdout_capture = ScriptOutputCapture(script_id, "INFO", original_stdout)
    stderr_capture = ScriptOutputCapture(script_id, "ERROR", original_stderr)
    
    try:
        # 替换标准输出流
        sys.stdout = stdout_capture
        sys.stderr = stderr_capture
        
        yield {
            'stdout_capture': stdout_capture,
            'stderr_capture': stderr_capture
        }
        
    finally:
        # 恢复原始输出流
        sys.stdout = original_stdout
        sys.stderr = original_stderr


async def push_script_error(script_id: str, error_message: str):
    """推送脚本错误信息到WebSocket"""
    try:
        websocket_manager = get_websocket_manager()
        await websocket_manager.broadcast_script_log(
            script_id, "ERROR", error_message
        )
    except Exception as e:
        logger.debug(f"推送脚本错误信息失败 [{script_id}]: {e}")


async def push_script_info(script_id: str, info_message: str):
    """推送脚本信息到WebSocket"""
    try:
        websocket_manager = get_websocket_manager()
        await websocket_manager.broadcast_script_log(
            script_id, "INFO", info_message
        )
    except Exception as e:
        logger.debug(f"推送脚本信息失败 [{script_id}]: {e}")