import torch
import os
import logging

import folder_paths
import comfy.utils
import comfy.lora_convert
import comfy.model_patcher

# Import custom implementations
from ..module.load_lora import load_lora, model_lora_keys_unet, model_lora_keys_clip, standardize_lora_key_format

logger = logging.getLogger(__name__)

DTYPES = {
    "float16": torch.float16,
    "bfloat16": torch.bfloat16,
    "float32": torch.float32,
}


def load_lora_with_metadata(lora_path, dtype="float16"):
    """
    Load a LoRA/LyCORIS file and preserve its metadata through conversion.
    Returns:
        tuple: (converted_lora_dict, metadata_dict)
    """
    import json
    import safetensors.torch

    # Load the file
    lora_data = comfy.utils.load_torch_file(lora_path, safe_load=True)
    # convert the dtype
    for key, value in lora_data.items():
        if isinstance(value, torch.Tensor):
            lora_data[key] = value.to(DTYPES.get(dtype, torch.float16))
    metadata = {}

    # Extract metadata from safetensors
    if lora_path.endswith('.safetensors'):
        with safetensors.torch.safe_open(lora_path, framework="pt") as f:
            metadata = dict(f.metadata())

    # Parse network args if present
    network_args = None
    if 'ss_network_args' in metadata:
        try:
            network_args = json.loads(metadata['ss_network_args'])
            metadata['network_args_parsed'] = network_args
            logger.info(f"Detected algorithm: {network_args.get('algo', 'unknown')}")
        except json.JSONDecodeError:
            logger.warning("Failed to parse ss_network_args")

    # Convert lora format
    converted = comfy.lora_convert.convert_lora(lora_data)

    # Preserve metadata by attaching it to the converted dict
    class LoraWithMetadata(dict):
        def __init__(self, data, metadata):
            super().__init__(data)
            self.metadata = metadata

    return LoraWithMetadata(converted, metadata), metadata


class LycorisLoaderNode:
    """Load LyCORIS adapters (LoHA, LoKr, LoCon, etc.) into ComfyUI models"""

    CATEGORY = "loaders"
    RETURN_TYPES = ("MODEL", "CLIP")
    RETURN_NAMES = ("model", "clip")
    FUNCTION = "load_lycoris"
    OUTPUT_NODE = False

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("MODEL",),
                "clip": ("CLIP",),
                "lora_name": (folder_paths.get_filename_list("loras"),),
                "strength_model": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}),
                "strength_clip": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}),
            },
            "optional": {
                "lycoris_type": (["any", "LoHA", "LoKr", "LoCon", "LyCORIS"], {"default": "any",
                                                                       "tooltip": "Force a specific LyCORIS type. 'Any' will try to detect automatically."}),
                "dtype": (["float16", "bfloat16", "float32"], {"default": "float16",
                                                               "tooltip": "Data type for the loaded models. 'float16' is recommended for most GPUs."}),
            }
        }

    def load_lycoris(self,
                     model,
                     clip,
                     lora_name,
                     strength_model,
                     strength_clip,
                     lycoris_type="any",
                     dtype="float16"):
        if strength_model == 0 and strength_clip == 0:
            return (model, clip)

        lora_path = folder_paths.get_full_path("loras", lora_name)
        if not os.path.exists(lora_path):
            logger.error(f"LoRA file not found: {lora_path}")
            return (model, clip)

        try:
            # Use our new loader function
            lora_data, metadata = load_lora_with_metadata(lora_path, dtype=dtype)
            # Build key map
            key_map = {}
            if model is not None and strength_model != 0:
                key_map = model_lora_keys_unet(model.model, key_map)
            if clip is not None and strength_clip != 0:
                key_map = model_lora_keys_clip(clip.cond_stage_model, key_map)

            # Load patches - lora_data now has metadata attached
            loaded = load_lora(lora_data, key_map, lycoris_type)

            # Apply patches to models
            new_model = model
            new_clip = clip

            if strength_model != 0 and model is not None:
                new_model = model.clone()
                new_model.add_patches(loaded, strength_model)
            if strength_clip != 0 and clip is not None:
                new_clip = clip.clone()
                new_clip.add_patches(loaded, strength_clip)

            logger.info(f"Successfully loaded LyCORIS: {lora_name}")
            return (new_model, new_clip)

        except Exception as e:
            logger.error(f"Failed to load LyCORIS {lora_path}: {e}", exc_info=True)
            return (model, clip)