"""
Model Metadata Reader Nodes
读取模型文件的元数据信息，包括下载地址等
"""

import os
import json
import safetensors.torch
import folder_paths


class ModelMetadataReader:
    """读取模型文件的元数据信息"""

    @classmethod
    def INPUT_TYPES(cls):
        model_types = ["checkpoints", "loras", "vae", "controlnet", "embeddings",
                      "upscale_models", "clip_vision", "diffusion_models", "text_encoders"]

        return {
            "required": {
                "model_type": (model_types, {"default": "loras"}),
                "model_name": ("STRING", {"default": "", "multiline": False}),
            }
        }

    RETURN_TYPES = ("STRING", "STRING", "STRING")
    RETURN_NAMES = ("metadata_json", "download_url", "model_info")
    FUNCTION = "read_metadata"
    CATEGORY = "utils/model"
    DESCRIPTION = "读取模型文件的元数据信息，包括下载地址等"

    def read_metadata(self, model_type, model_name):
        """读取模型元数据"""
        metadata_json = ""
        download_url = ""
        model_info = ""

        try:
            # 获取模型完整路径
            if not model_name:
                model_info = f"请输入模型名称\n可用模型类型: {model_type}"
                return (metadata_json, download_url, model_info)

            full_path = folder_paths.get_full_path(model_type, model_name)

            if full_path is None:
                # 尝试列出可用的模型文件
                available_models = folder_paths.get_filename_list(model_type)
                model_info = f"未找到模型: {model_name}\n\n可用模型 ({model_type}):\n"
                model_info += "\n".join(available_models[:20])  # 只显示前20个
                if len(available_models) > 20:
                    model_info += f"\n... 还有 {len(available_models) - 20} 个模型"
                return (metadata_json, download_url, model_info)

            # 读取模型元数据
            metadata = None
            file_ext = os.path.splitext(full_path)[1].lower()

            if file_ext in ['.safetensors', '.sft']:
                # 读取safetensors文件的元数据
                with safetensors.safe_open(full_path, framework="pt", device="cpu") as f:
                    metadata = f.metadata()

            if metadata:
                # 转换为JSON字符串
                metadata_json = json.dumps(metadata, indent=2, ensure_ascii=False)

                # 提取下载地址
                download_url = metadata.get("modelspec.download_url",
                              metadata.get("ss_download_url",
                              metadata.get("download_url", "")))

                # 生成模型信息摘要
                model_info = f"模型文件: {model_name}\n"
                model_info += f"完整路径: {full_path}\n"
                model_info += f"文件类型: {file_ext}\n"
                model_info += f"文件大小: {os.path.getsize(full_path) / (1024*1024):.2f} MB\n\n"

                # 添加一些关键元数据
                if "modelspec.title" in metadata:
                    model_info += f"标题: {metadata['modelspec.title']}\n"
                if "modelspec.description" in metadata:
                    model_info += f"描述: {metadata['modelspec.description']}\n"
                if "modelspec.architecture" in metadata:
                    model_info += f"架构: {metadata['modelspec.architecture']}\n"
                if "ss_base_model_version" in metadata:
                    model_info += f"基础模型: {metadata['ss_base_model_version']}\n"
                if "ss_resolution" in metadata:
                    model_info += f"训练分辨率: {metadata['ss_resolution']}\n"
                if "ss_network_dim" in metadata:
                    model_info += f"网络维度: {metadata['ss_network_dim']}\n"
                if "ss_network_alpha" in metadata:
                    model_info += f"网络Alpha: {metadata['ss_network_alpha']}\n"
                if download_url:
                    model_info += f"\n下载地址: {download_url}\n"

                # 显示元数据字段数量
                model_info += f"\n元数据字段总数: {len(metadata)}"
            else:
                model_info = f"模型文件: {model_name}\n"
                model_info += f"完整路径: {full_path}\n"
                model_info += f"文件类型: {file_ext}\n"
                model_info += f"文件大小: {os.path.getsize(full_path) / (1024*1024):.2f} MB\n\n"
                model_info += "该文件格式不支持读取元数据或没有元数据信息"

        except Exception as e:
            model_info = f"读取模型元数据时发生错误:\n{str(e)}"

        return (metadata_json, download_url, model_info)


class ModelListReader:
    """列出指定类型的所有模型"""

    @classmethod
    def INPUT_TYPES(cls):
        model_types = ["checkpoints", "loras", "vae", "controlnet", "embeddings",
                      "upscale_models", "clip_vision", "diffusion_models", "text_encoders"]

        return {
            "required": {
                "model_type": (model_types, {"default": "loras"}),
            }
        }

    RETURN_TYPES = ("STRING", "INT")
    RETURN_NAMES = ("model_list", "model_count")
    FUNCTION = "list_models"
    CATEGORY = "utils/model"
    DESCRIPTION = "列出指定类型的所有模型文件"

    def list_models(self, model_type):
        """列出模型"""
        try:
            # 获取模型列表
            models = folder_paths.get_filename_list(model_type)
            model_count = len(models)

            # 生成列表字符串
            model_list = f"模型类型: {model_type}\n"
            model_list += f"模型总数: {model_count}\n\n"
            model_list += "模型列表:\n"
            model_list += "\n".join([f"{i+1}. {model}" for i, model in enumerate(models)])

        except Exception as e:
            model_list = f"读取模型列表时发生错误:\n{str(e)}"
            model_count = 0

        return (model_list, model_count)


# 节点映射
NODE_CLASS_MAPPINGS = {
    "ModelMetadataReader": ModelMetadataReader,
    "ModelListReader": ModelListReader,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "ModelMetadataReader": "Model Metadata Reader",
    "ModelListReader": "Model List Reader",
}