"""
节点依赖管理工具
支持节点独立管理自己的依赖，使用 uv pip install 直接安装
支持异步并行安装，提升多节点场景下的性能
"""
import os
import asyncio
import hashlib
import subprocess
import logging
from pathlib import Path
from typing import Dict, Any, Optional, List
from jingrow.utils.path import get_jingrow_root, get_root_path

logger = logging.getLogger(__name__)


def get_node_pyproject_path(node_type: str) -> Optional[Path]:
    """获取节点的 pyproject.toml 路径"""
    jingrow_root = get_jingrow_root()
    node_dir = jingrow_root / "ai" / "nodes" / node_type
    pyproject_path = node_dir / "pyproject.toml"
    return pyproject_path if pyproject_path.exists() else None


def has_node_dependencies(node_type: str) -> bool:
    """检查节点是否有独立的依赖配置"""
    return get_node_pyproject_path(node_type) is not None


def discover_nodes_with_dependencies() -> List[str]:
    """
    自动发现所有有 pyproject.toml 的节点
    
    Returns:
        节点类型列表
    """
    jingrow_root = get_jingrow_root()
    nodes_dir = jingrow_root / "ai" / "nodes"
    
    if not nodes_dir.exists():
        return []
    
    nodes_with_deps = []
    for node_dir in nodes_dir.iterdir():
        if node_dir.is_dir() and not node_dir.name.startswith('_'):
            pyproject_path = node_dir / "pyproject.toml"
            if pyproject_path.exists():
                nodes_with_deps.append(node_dir.name)
    
    return sorted(nodes_with_deps)


def _parse_node_dependencies(pyproject_path: Path) -> tuple[List[str], str]:
    """
    解析节点的依赖列表和 pyproject.toml 内容
    
    Args:
        pyproject_path: pyproject.toml 文件路径
    
    Returns:
        (dependencies列表, pyproject.toml内容)
    """
    dependencies = []
    content = ""
    try:
        with open(pyproject_path, 'r', encoding='utf-8') as f:
            content = f.read()
        
        in_dependencies = False
        bracket_count = 0
        for line in content.split('\n'):
            stripped = line.strip()
            if stripped.startswith('dependencies') and '=' in stripped and '[' in stripped:
                in_dependencies = True
                bracket_count = stripped.count('[') - stripped.count(']')
                if ']' in stripped:
                    in_dependencies = False
                continue
            elif in_dependencies:
                bracket_count += stripped.count('[') - stripped.count(']')
                if bracket_count <= 0 and ']' in stripped:
                    break
                dep = stripped.strip('",\' ,')
                if dep and not dep.startswith('#'):
                    dependencies.append(dep)
    except Exception as e:
        logger.warning(f"解析 pyproject.toml 失败: {str(e)}")
    
    return dependencies, content


def _build_install_command(dependencies: List[str], node_dir: Path) -> List[str]:
    """
    构建安装命令
    
    Args:
        dependencies: 依赖列表
        node_dir: 节点目录
    
    Returns:
        安装命令列表
    """
    if dependencies:
        return ["uv", "pip", "install"] + dependencies
    else:
        return ["uv", "pip", "install", "-e", str(node_dir)]


def _format_install_result(node_type: str, success: bool, stdout: str, stderr: str, 
                          after_install_result: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
    """
    格式化安装结果
    
    Args:
        node_type: 节点类型
        success: 是否成功
        stdout: 标准输出
        stderr: 标准错误
        after_install_result: 后安装脚本结果
    
    Returns:
        格式化的结果字典
    """
    if success:
        message = f"节点 {node_type} 依赖安装成功"
        if after_install_result:
            if after_install_result.get("success"):
                message += f" ({after_install_result.get('message', '后安装脚本执行成功')})"
            else:
                logger.warning(f"节点 {node_type} 后安装脚本执行失败: {after_install_result.get('error', '未知错误')}")
                message += f" (后安装脚本警告: {after_install_result.get('error', '未知错误')})"
        
        return {
            "success": True,
            "message": message,
            "output": stdout,
            "post_install": after_install_result
        }
    else:
        error_msg = stderr or stdout
        logger.error(f"安装节点 {node_type} 依赖失败: {error_msg}")
        return {
            "success": False,
            "error": f"依赖安装失败: {error_msg}",
            "stderr": stderr,
            "stdout": stdout
        }


def install_node_dependencies(node_type: str) -> Dict[str, Any]:
    """
    安装节点的依赖（同步版本，用于单个节点安装）
    
    Args:
        node_type: 节点类型
    
    Returns:
        安装结果字典
    """
    try:
        jingrow_root = get_jingrow_root()
        node_dir = jingrow_root / "ai" / "nodes" / node_type
        
        if not node_dir.exists():
            return {"success": False, "error": f"节点目录不存在: {node_type}"}
        
        pyproject_path = node_dir / "pyproject.toml"
        if not pyproject_path.exists():
            # 节点没有独立依赖，返回成功（使用根依赖）
            return {"success": True, "message": "节点使用根依赖，无需安装"}
        
        project_root = jingrow_root.parent
        
        # 解析依赖
        dependencies, content = _parse_node_dependencies(pyproject_path)
        
        # 构建安装命令
        cmd = _build_install_command(dependencies, node_dir)
        
        # 执行安装
        result = subprocess.run(
            cmd,
            cwd=project_root,
            capture_output=True,
            text=True,
            timeout=300
        )
        
        # 执行后安装脚本
        after_install_result = None
        if result.returncode == 0:
            after_install_result = _run_node_after_install(node_type, node_dir, project_root, content)
        
        # 格式化结果
        return _format_install_result(
            node_type,
            result.returncode == 0,
            result.stdout,
            result.stderr,
            after_install_result
        )
            
    except subprocess.TimeoutExpired:
        return {"success": False, "error": "依赖安装超时"}
    except Exception as e:
        logger.error(f"安装节点 {node_type} 依赖时发生异常: {str(e)}")
        return {"success": False, "error": str(e)}


def check_node_dependencies(node_type: str) -> Dict[str, Any]:
    """
    检查节点依赖是否已安装（简单检查，通过尝试导入主要依赖）
    
    Args:
        node_type: 节点类型
    
    Returns:
        检查结果字典
    """
    pyproject_path = get_node_pyproject_path(node_type)
    if not pyproject_path:
        return {"success": True, "message": "节点无独立依赖"}
    
    try:
        # 读取 pyproject.toml 获取依赖列表（使用简单文本解析，避免额外依赖）
        with open(pyproject_path, 'r', encoding='utf-8') as f:
            content = f.read()
        
        # 简单提取 dependencies 列表
        dependencies = []
        in_dependencies = False
        bracket_count = 0
        for line in content.split('\n'):
            stripped = line.strip()
            # 检测 dependencies 数组开始
            if stripped.startswith('dependencies') and '=' in stripped and '[' in stripped:
                in_dependencies = True
                bracket_count = stripped.count('[') - stripped.count(']')
                # 检查是否在同一行结束
                if ']' in stripped:
                    in_dependencies = False
                continue
            elif in_dependencies:
                # 计算括号
                bracket_count += stripped.count('[') - stripped.count(']')
                if bracket_count <= 0 and ']' in stripped:
                    break
                # 提取依赖字符串（去除引号和逗号）
                dep = stripped.strip('",\' ,')
                if dep and not dep.startswith('#'):
                    dependencies.append(dep)
        
        if not dependencies:
            return {"success": True, "message": "节点无依赖需要检查"}
        
        # 简单检查：尝试导入依赖包（提取包名）
        missing = []
        # 包名到模块名的映射（处理特殊情况）
        package_to_module = {
            "beautifulsoup4": "bs4",
            "Pillow": "PIL",
        }
        
        for dep in dependencies:
            # 解析依赖字符串，提取包名（处理版本约束）
            package_name = dep.split(">=")[0].split("==")[0].split("~=")[0].split("<")[0].strip()
            # 处理 extras，如 "package[extra]"
            package_name = package_name.split("[")[0]
            
            # 获取模块名（优先使用映射，否则转换 - 为 _）
            module_name = package_to_module.get(package_name, package_name.replace("-", "_"))
            
            try:
                __import__(module_name)
            except ImportError:
                missing.append(package_name)
        
        if missing:
            return {
                "success": False,
                "missing": missing,
                "message": f"缺少依赖: {', '.join(missing)}"
            }
        else:
            return {"success": True, "message": "所有依赖已安装"}
            
    except Exception as e:
        logger.warning(f"检查节点 {node_type} 依赖时发生异常: {str(e)}")
        # 检查失败时返回 None，让安装流程处理
        return {"success": None, "message": f"检查依赖时出错: {str(e)}"}


def _run_node_after_install(node_type: str, node_dir: Path, project_root: Path, pyproject_content: str) -> Optional[Dict[str, Any]]:
    """
    执行节点的后安装脚本（如果配置了的话）
    
    支持配置格式：
    [tool.jingrow.after_install]
    commands = [
        "python -m playwright install-deps",
        "python -m playwright install"
    ]
    # 或单条命令
    script = "python -m playwright install-deps && python -m playwright install"
    
    Args:
        node_type: 节点类型
        node_dir: 节点目录
        project_root: 项目根目录
        pyproject_content: pyproject.toml 文件内容
    
    Returns:
        执行结果字典，如果没有配置后安装脚本则返回 None
    """
    try:
        commands = []
        
        # 解析 [tool.jingrow.after_install] 配置
        in_after_install = False
        in_commands_array = False
        bracket_count = 0
        
        for line in pyproject_content.split('\n'):
            stripped = line.strip()
            
            # 检测 [tool.jingrow.after_install] 段
            if stripped.startswith('[tool.jingrow.after_install]'):
                in_after_install = True
                continue
            
            # 如果遇到新的段，停止解析
            if in_after_install and stripped.startswith('[') and not stripped.startswith('[tool.jingrow'):
                break
            
            if in_after_install:
                # 查找 commands = [...] 数组格式（推荐，支持多条命令）
                if stripped.startswith('commands') and '=' in stripped and '[' in stripped:
                    in_commands_array = True
                    bracket_count = stripped.count('[') - stripped.count(']')
                    # 检查是否在同一行结束
                    if ']' in stripped:
                        in_commands_array = False
                    # 提取同一行的命令
                    if '[' in stripped and ']' in stripped:
                        # 提取数组内容
                        array_content = stripped.split('[', 1)[1].rsplit(']', 1)[0]
                        for cmd in array_content.split(','):
                            cmd = cmd.strip().strip('"').strip("'")
                            if cmd:
                                commands.append(cmd)
                    continue
                elif in_commands_array:
                    # 继续解析数组内容
                    bracket_count += stripped.count('[') - stripped.count(']')
                    if bracket_count <= 0 and ']' in stripped:
                        in_commands_array = False
                        # 如果这一行只有 ]，直接跳过
                        if stripped.strip() == ']':
                            break
                        # 否则提取 ] 之前的内容
                        stripped = stripped.rsplit(']', 1)[0]
                    
                    # 提取命令（去除引号和逗号）
                    cmd = stripped.strip('",\' ,')
                    # 跳过空行、注释和单独的 ]
                    if cmd and not cmd.startswith('#') and cmd != ']':
                        commands.append(cmd)
                    if bracket_count <= 0:
                        break
                    continue
                # 查找 script = 配置（单条命令）
                elif stripped.startswith('script') and '=' in stripped:
                    script_value = stripped.split('=', 1)[1].strip()
                    # 处理字符串格式（去除引号）
                    if script_value.startswith('"') and script_value.endswith('"'):
                        script_value = script_value[1:-1]
                    elif script_value.startswith("'") and script_value.endswith("'"):
                        script_value = script_value[1:-1]
                    if script_value:
                        commands = [script_value]
                        break
        
        if not commands:
            return None
        
        # 将多条命令合并为一条脚本（用 && 连接）
        script = ' && '.join(commands)
        
        # 检查是否已经执行过相同的脚本
        script_hash = hashlib.md5(script.encode('utf-8')).hexdigest()
        # 将标记文件存储在项目根目录的缓存目录
        cache_dir = get_root_path() / '.cache' / 'node_install'
        cache_dir.mkdir(parents=True, exist_ok=True)
        after_install_marker = cache_dir / f"{node_type}.after_install_done"
        
        # 如果标记文件存在，检查脚本哈希值是否匹配
        if after_install_marker.exists():
            try:
                with open(after_install_marker, 'r', encoding='utf-8') as f:
                    saved_hash = f.read().strip()
                if saved_hash == script_hash:
                    logger.info(f"节点 {node_type} 的后安装脚本已执行过，跳过")
                    return {
                        "success": True,
                        "message": "后安装脚本已执行过，跳过",
                        "skipped": True
                    }
            except Exception as e:
                logger.warning(f"读取后安装标记文件失败: {str(e)}，将重新执行")
        
        logger.info(f"执行节点 {node_type} 的后安装脚本: {script}")
        
        # 执行脚本（使用 shell 执行，支持 && 等操作符）
        # 将脚本中的 python -m 命令替换为 uv run python -m，确保在正确的环境中执行
        import re
        # 替换 python -m 为 uv run python -m（避免重复替换已存在的 uv run）
        script = re.sub(r'(?<!uv run )\bpython -m\b', 'uv run python -m', script)
        
        result = subprocess.run(
            script,
            shell=True,
            cwd=project_root,
            capture_output=True,
            text=True,
            timeout=600
        )
        
        if result.returncode == 0:
            # 保存执行成功的标记
            try:
                with open(after_install_marker, 'w', encoding='utf-8') as f:
                    f.write(script_hash)
                logger.debug(f"已保存节点 {node_type} 的后安装脚本执行标记")
            except Exception as e:
                logger.warning(f"保存后安装标记文件失败: {str(e)}")
            
            return {
                "success": True,
                "message": "后安装脚本执行成功",
                "output": result.stdout
            }
        else:
            error_msg = result.stderr or result.stdout
            return {
                "success": False,
                "error": f"后安装脚本执行失败: {error_msg}",
                "stderr": result.stderr,
                "stdout": result.stdout
            }
            
    except subprocess.TimeoutExpired:
        return {
            "success": False,
            "error": "后安装脚本执行超时"
        }
    except Exception as e:
        logger.warning(f"执行节点 {node_type} 后安装脚本时发生异常: {str(e)}")
        return {
            "success": False,
            "error": f"后安装脚本执行异常: {str(e)}"
        }


def ensure_node_dependencies(node_type: str) -> Dict[str, Any]:
    """
    确保节点依赖已安装（检查并安装）
    
    Args:
        node_type: 节点类型
    
    Returns:
        结果字典
    """
    if not has_node_dependencies(node_type):
        return {"success": True, "message": "节点无独立依赖"}
    
    # 先检查
    check_result = check_node_dependencies(node_type)
    if check_result.get("success") is True:
        return {"success": True, "message": "依赖已安装"}
    
    # 如果检查失败或不确定，尝试安装
    return install_node_dependencies(node_type)


async def install_node_dependencies_async(node_type: str) -> Dict[str, Any]:
    """
    异步安装节点的依赖（内部使用，保持与同步版本相同的逻辑）
    
    Args:
        node_type: 节点类型
    
    Returns:
        安装结果字典
    """
    try:
        jingrow_root = get_jingrow_root()
        node_dir = jingrow_root / "ai" / "nodes" / node_type
        
        if not node_dir.exists():
            return {"success": False, "error": f"节点目录不存在: {node_type}"}
        
        pyproject_path = node_dir / "pyproject.toml"
        if not pyproject_path.exists():
            return {"success": True, "message": "节点使用根依赖，无需安装"}
        
        project_root = jingrow_root.parent
        
        # 解析依赖
        dependencies, content = _parse_node_dependencies(pyproject_path)
        
        # 构建安装命令
        cmd = _build_install_command(dependencies, node_dir)
        
        # 使用异步 subprocess
        process = await asyncio.create_subprocess_exec(
            *cmd,
            cwd=project_root,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )
        
        stdout, stderr = await asyncio.wait_for(process.communicate(), timeout=300)
        stdout_text = stdout.decode('utf-8') if stdout else ''
        stderr_text = stderr.decode('utf-8') if stderr else ''
        
        # 执行后安装脚本（同步执行，因为可能涉及系统级操作）
        after_install_result = None
        if process.returncode == 0:
            after_install_result = _run_node_after_install(node_type, node_dir, project_root, content)
        
        # 格式化结果
        return _format_install_result(
            node_type,
            process.returncode == 0,
            stdout_text,
            stderr_text,
            after_install_result
        )
            
    except asyncio.TimeoutError:
        return {"success": False, "error": "依赖安装超时"}
    except Exception as e:
        logger.error(f"安装节点 {node_type} 依赖时发生异常: {str(e)}")
        return {"success": False, "error": str(e)}


def install_all_nodes_dependencies(max_concurrent: int = 5) -> List[Dict[str, Any]]:
    """
    并行安装所有节点的依赖（同步接口，内部使用异步）
    
    Args:
        max_concurrent: 最大并发数，默认5，避免资源过载
    
    Returns:
        所有节点的安装结果列表
    """
    nodes = discover_nodes_with_dependencies()
    
    if not nodes:
        logger.info("未发现需要安装依赖的节点")
        return []
    
    logger.info(f"发现 {len(nodes)} 个节点需要安装依赖，使用 {max_concurrent} 个并发")
    
    async def install_all():
        semaphore = asyncio.Semaphore(max_concurrent)
        
        async def install_with_limit(node_type: str):
            async with semaphore:
                return await install_node_dependencies_async(node_type)
        
        tasks = [install_with_limit(node) for node in nodes]
        return await asyncio.gather(*tasks, return_exceptions=True)
    
    # 运行异步函数
    results = asyncio.run(install_all())
    
    # 处理异常结果
    processed_results = []
    for i, result in enumerate(results):
        if isinstance(result, Exception):
            node_type = nodes[i]
            logger.error(f"安装节点 {node_type} 依赖时发生异常: {str(result)}")
            processed_results.append({
                "success": False,
                "error": f"异常: {str(result)}",
                "node_type": node_type
            })
        else:
            processed_results.append(result)
    
    # 统计结果
    success_count = sum(1 for r in processed_results if r.get("success"))
    logger.info(f"节点依赖安装完成: {success_count}/{len(nodes)} 成功")
    
    return processed_results

