"""
工具节点模块 - 重构版本

提供各种实用功能的ComfyUI节点，包括：
- 数据库日志记录
- VOD文件下载（使用统一存储抽象层）
- 模型管理
- 文件遍历
- Webhook通知
- 任务ID生成

本模块已经过重构，使用统一的存储抽象层替代直接SDK调用。
"""

import glob
import json
import os
import time
import uuid
from pathlib import Path
from typing import Any, Dict, Tuple

import comfy.model_management
import requests
import server
from loguru import logger
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from ..utils.object_storage import DownloadResult, get_provider
from ..utils.task_table import Task


class LogToDB:
    """
    数据库日志记录节点

    将ComfyUI工作流的执行结果记录到数据库中，支持任务状态跟踪。
    """

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "job_id": ("STRING", {"forceInput": True}),
                "log": ("STRING", {"forceInput": True}),
                "status": ("INT", {"default": 1, "max": 1}),
                "sql_url": (
                    "STRING",
                    {"default": "mysql+pymysql://root:root@example.com:3306/test"},
                ),
            },
            "hidden": {
                "unique_id": "UNIQUE_ID",
            },
        }

    RETURN_TYPES = ("STRING",)
    FUNCTION = "log2db"
    OUTPUT_NODE = True
    OUTPUT_IS_LIST = (True,)
    CATEGORY = "不忘科技-自定义节点🚩/工具"

    def log2db(
        self, log: str, status: int, sql_url: str, unique_id: str, **kwargs
    ) -> Dict[str, Any]:
        """
        记录日志到数据库

        Args:
            log: 日志内容
            status: 状态代码
            sql_url: 数据库连接URL
            unique_id: 唯一ID

        Returns:
            Dict: 包含结果信息的字典
        """
        try:
            # 获取comfy服务器队列信息
            (_, prompt_id, prompt, extra_data, outputs_to_execute) = next(
                iter(
                    server.PromptServer.instance.prompt_queue.currently_running.values()
                )
            )
            job_id = extra_data["client_id"]

            engine = create_engine(sql_url, echo=True)
            session = sessionmaker(bind=engine)()

            # 查询现有任务
            tasks = session.query(Task).filter(Task.prompt_id == prompt_id).all()

            result = {
                "curr_node_id": str(unique_id),
                "last_node_id": list(prompt.keys())[-1],
                "node_output": str(log),
            }

            if len(tasks) == 0:
                # 不存在则插入
                task = Task(
                    prompt_id=prompt_id,
                    job_id=job_id,
                    result=json.dumps(result),
                    status=status,
                )
                session.add(task)
                logger.info(f"新增任务记录: {prompt_id}")
            elif len(tasks) == 1:
                # 存在则更新
                session.query(Task).filter(Task.prompt_id == prompt_id).update(
                    {"result": json.dumps(result), "status": status}
                )
                logger.info(f"更新任务记录: {prompt_id}")
            else:
                # 异常情况
                session.rollback()
                raise RuntimeError("状态数据库prompt_id不唯一, 无法记录状态！")

            session.commit()
            session.close()

            return {"ui": {"text": json.dumps(result)}, "result": (json.dumps(result),)}

        except Exception as e:
            logger.error(f"数据库日志记录失败: {e}")
            return {"ui": {"text": str(e)}, "result": (str(e),)}


class VodToLocalNode:
    """
    腾讯云VOD文件下载节点

    使用统一的存储抽象层从腾讯云视频点播(VOD)服务下载媒体文件到本地。
    支持通过文件ID和子应用ID定位和下载视频文件。
    """

    def __init__(self):
        """
        初始化VOD下载节点

        使用统一的存储管理器获取VOD存储提供者，
        替代直接使用腾讯云VOD SDK的方式。
        """
        try:
            # 使用统一存储管理器获取VOD提供者
            self.vod_provider = get_provider("vod")
            logger.info("VOD下载节点初始化成功")

        except Exception as e:
            logger.error(f"VOD下载节点初始化失败: {e}")
            raise RuntimeError(f"VOD节点初始化失败: {e}")

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "file_id": ("STRING", {"default": ""}),
                "sub_app_id": ("STRING", {"default": ""}),
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("local_path",)
    FUNCTION = "execute"
    CATEGORY = "不忘科技-自定义节点🚩/工具"

    def execute(self, file_id: str, sub_app_id: str) -> Tuple[str]:
        """
        执行VOD文件下载

        Args:
            file_id: VOD文件ID
            sub_app_id: 子应用ID

        Returns:
            tuple: 包含本地文件路径的元组

        Raises:
            Exception: 下载失败时抛出异常
        """
        try:
            # 参数验证
            if not file_id or not file_id.strip():
                raise ValueError("文件ID不能为空")
            if not sub_app_id or not sub_app_id.strip():
                raise ValueError("子应用ID不能为空")

            # 调用下载逻辑
            local_path = self.download_vod(file_id.strip(), sub_app_id.strip())
            logger.info(f"VOD文件下载成功: {local_path}")
            return (local_path,)

        except Exception as e:
            logger.error(f"VOD文件下载执行失败: {e}")
            raise Exception(f"VOD文件下载失败: {str(e)}")

    def create_directory(self, path: str) -> None:
        """
        创建目录（如果不存在）

        Args:
            path: 目录路径
        """
        p = Path(path)
        if not p.exists():
            p.mkdir(parents=True, exist_ok=True)
            logger.info(f"目录已创建: {path}")
        else:
            logger.debug(f"目录已存在: {path}")

    def download_vod(self, file_id: str, sub_app_id: str) -> str:
        """
        下载腾讯云VOD文件到本地

        Args:
            file_id: VOD文件ID
            sub_app_id: 子应用ID

        Returns:
            str: 本地文件路径

        Raises:
            Exception: 下载失败时抛出异常
        """
        try:
            # 生成本地存储路径
            download_dir = os.path.join(
                os.path.dirname(os.path.abspath(__file__)), "download", f"{sub_app_id}"
            )
            self.create_directory(download_dir)

            output_path = os.path.join(download_dir, f"{file_id}.mp4")

            # 如果文件已存在，直接返回
            if os.path.exists(output_path):
                logger.info(f"VOD文件已存在，直接返回: {output_path}")
                return output_path

            logger.info(f"开始下载VOD文件: {file_id} -> {output_path}")

            # 使用VOD提供者下载文件
            result: DownloadResult = self.vod_provider.download_file(
                file_id, output_path, sub_app_id, timeout=600
            )

            if result.success:
                logger.info(f"VOD文件下载成功: {result.local_path}")
                return result.local_path
            else:
                raise Exception(result.message or "VOD下载失败")

        except Exception as e:
            logger.error(f"VOD文件下载异常: {e}")
            raise Exception(f"VOD下载失败: {str(e)}")


class AnyType(str):
    """特殊类型，用于在不等比较中始终相等。Credit to pythongosssss"""

    def __ne__(self, __value: object) -> bool:
        return False


any = AnyType("*")


class UnloadAllModels:
    """
    卸载所有模型节点

    释放GPU内存，卸载所有已加载的模型。
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {"any": (any, {"forceInput": True})},
            "optional": {},
        }

    RETURN_TYPES = ()
    FUNCTION = "unload_models"
    CATEGORY = "不忘科技-自定义节点🚩/工具"
    OUTPUT_NODE = True

    def unload_models(self, any=None) -> Tuple:
        """
        卸载所有已加载的模型

        Args:
            any: 输入参数（任意类型）

        Returns:
            tuple: 空元组
        """
        try:
            logger.info("开始卸载所有模型...")

            # 卸载所有已加载的模型
            comfy.model_management.soft_empty_cache()
            comfy.model_management.unload_all_models()
            comfy.model_management.soft_empty_cache()

            logger.info("所有模型已成功卸载")
            return ()

        except Exception as e:
            logger.error(f"模型卸载失败: {e}")
            # 即使失败也返回空元组，避免中断工作流
            return ()


class TraverseFolder:
    """
    文件夹遍历节点

    遍历指定文件夹，查找符合条件的文件。
    """

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "folder": (
                    "STRING",
                    {"default": r"E:\comfy\ComfyUI\input\s3", "required": True},
                ),
                "subfix": ("STRING", {"default": ".mp4", "required": True}),
                "recursive": ("BOOLEAN", {"default": True, "required": True}),
                "idx": (
                    "INT",
                    {"default": 0, "min": 0, "max": 0xFFFFFF},
                ),
            },
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("文件路径",)
    FUNCTION = "compute"
    CATEGORY = "不忘科技-自定义节点🚩/工具"

    def compute(
        self, folder: str, subfix: str, recursive: bool, idx: int
    ) -> Tuple[str]:
        """
        遍历文件夹查找文件

        Args:
            folder: 文件夹路径
            subfix: 文件后缀
            recursive: 是否递归查找
            idx: 文件索引

        Returns:
            tuple: 包含文件路径的元组

        Raises:
            RuntimeError: 未找到文件时抛出异常
        """
        try:
            # 参数验证
            if not folder or not folder.strip():
                raise ValueError("文件夹路径不能为空")
            if not subfix or not subfix.strip():
                raise ValueError("文件后缀不能为空")

            folder = folder.strip()
            subfix = subfix.strip()

            # 构建搜索模式
            if recursive:
                pattern = os.path.join(folder, f"**/*{subfix}")
            else:
                pattern = os.path.join(folder, f"*{subfix}")

            # 查找文件
            files = glob.glob(pattern, recursive=recursive)

            if len(files) == 0:
                logger.warning(f"在文件夹 {folder} 中未找到后缀为 {subfix} 的文件")
                raise RuntimeError("No Files Found")

            # 选择文件
            selected_file = files[idx % len(files)]
            logger.info(
                f"找到 {len(files)} 个文件，选择第 {idx % len(files)} 个: {selected_file}"
            )

            return (str(selected_file),)

        except Exception as e:
            logger.error(f"文件夹遍历失败: {e}")
            if isinstance(e, RuntimeError):
                raise
            else:
                raise RuntimeError(f"文件夹遍历错误: {str(e)}")


class PlugAndPlayWebhook:
    """
    即插即用Webhook节点

    连上线就能转发数据到指定的Webhook地址。
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "webhook_url": (
                    "STRING",
                    {
                        "default": "http://127.0.0.1:8010/handler/webhook",
                        "placeholder": "https://your-api.com/webhook",
                    },
                ),
                "image_url": ("STRING", {"default": "", "placeholder": "图片的url"}),
            },
            "optional": {
                "prompt_id": ("STRING", {"default": ""}),
            },
            "hidden": {
                "prompt": "PROMPT",
                "extra_pnginfo": "EXTRA_PNGINFO",
                "unique_id": "UNIQUE_ID",
            },
        }

    RETURN_TYPES = ()
    FUNCTION = "send"
    OUTPUT_NODE = True
    CATEGORY = "不忘科技-自定义节点🚩/工具"

    def send(
        self,
        webhook_url: str,
        image_url: str,
        prompt_id: str = "",
        prompt=None,
        extra_pnginfo=None,
        unique_id=None,
    ) -> Tuple:
        """
        发送数据到Webhook

        Args:
            webhook_url: Webhook URL
            image_url: 图片URL
            prompt_id: 提示ID
            prompt: 提示信息
            extra_pnginfo: 额外PNG信息
            unique_id: 唯一ID

        Returns:
            tuple: 空元组

        Raises:
            ValueError: URL为空时抛出异常
        """
        try:
            # 参数验证
            if not webhook_url or not webhook_url.strip():
                raise ValueError("❌ 请填写Webhook URL！")

            # 使用传入的prompt_id，如果没有则用unique_id
            final_prompt_id = prompt_id or unique_id or "unknown"

            # 准备发送的数据
            data = {
                "img_base64": image_url,
                "format": "png",
                "image_url": image_url,
                "prompt_id": final_prompt_id,
                "timestamp": time.time(),
            }

            logger.info(f"准备发送Webhook数据到: {webhook_url}")

            # 发送Webhook
            response = requests.post(webhook_url.strip(), json=data, timeout=30)
            response.raise_for_status()

            logger.info(f"Webhook发送成功，响应状态: {response.status_code}")
            logger.debug(f"发送的数据: {data}")

        except requests.RequestException as e:
            logger.error(f"❌ Webhook发送失败: {str(e)}")
            # 不抛出异常，避免中断工作流
        except Exception as e:
            logger.error(f"❌ Webhook发送异常: {str(e)}")
            # 不抛出异常，避免中断工作流

        # 终端节点，无需返回
        return ()


class TaskIdGenerate:
    """
    TaskID生成器

    用户可传入自定义TaskID或自动生成UUID。
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {},
            "optional": {
                "custom_task_id": (
                    "STRING",
                    {"default": "", "placeholder": "留空则自动生成"},
                ),
            },
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("task_id",)
    FUNCTION = "generate_task_id"
    OUTPUT_NODE = False
    CATEGORY = "不忘科技-自定义节点🚩/工具"

    def generate_task_id(self, custom_task_id: str = "") -> Tuple[str]:
        """
        生成或使用自定义TaskID

        Args:
            custom_task_id: 自定义TaskID

        Returns:
            tuple: 包含TaskID的元组
        """
        try:
            if custom_task_id and custom_task_id.strip():
                # 用户输入了自定义ID
                task_id = custom_task_id.strip()
                logger.info(f"📝 使用自定义TaskID: {task_id}")
            else:
                # 自动生成UUID
                task_id = str(uuid.uuid4())
                logger.info(f"🎲 自动生成TaskID: {task_id}")

            return (task_id,)

        except Exception as e:
            logger.error(f"TaskID生成失败: {e}")
            # 发生异常时生成一个基础UUID
            fallback_id = str(uuid.uuid4())
            logger.warning(f"使用备用TaskID: {fallback_id}")
            return (fallback_id,)


# 节点映射字典
NODE_CLASS_MAPPINGS = {
    "LogToDB": LogToDB,
    "VodToLocalNode": VodToLocalNode,
    "UnloadAllModels": UnloadAllModels,
    "TraverseFolder": TraverseFolder,
    "PlugAndPlayWebhook": PlugAndPlayWebhook,
    "TaskIdGenerate": TaskIdGenerate,
}

# 节点显示名称映射
NODE_DISPLAY_NAME_MAPPINGS = {
    "LogToDB": "数据库日志记录",
    "VodToLocalNode": "VOD文件下载",
    "UnloadAllModels": "卸载所有模型",
    "TraverseFolder": "文件夹遍历",
    "PlugAndPlayWebhook": "Webhook通知",
    "TaskIdGenerate": "TaskID生成器",
}