import functools
import logging
import os.path
import re
from typing import TYPE_CHECKING, Optional

import network
import torch

if TYPE_CHECKING:
    from backend.patcher.clip import CLIP
    from backend.patcher.unet import UnetPatcher

from backend.args import dynamic_args
from backend.logging import setup_logger
from backend.patcher.lora import load_lora, model_lora_keys_clip, model_lora_keys_unet
from backend.utils import load_torch_file
from modules import errors, scripts, sd_models, shared

logger = logging.getLogger("lora")
setup_logger(logger)


load_lora_state_dict = functools.partial(load_torch_file, safe_load=True)


def process_anima(lora: dict[str, torch.Tensor]):
    # LLMAdapter was moved from transformer to text_encoder

    keys = list(lora.keys())
    for k in keys:
        if k.startswith("diffusion_model.llm_adapter"):
            lora[k.replace("diffusion_model", "text_encoders.qwen3_06b")] = lora.pop(k)
        elif k.startswith("lora_unet_llm_adapter"):
            lora[k.replace("lora_unet_llm_adapter", "lora_te_llm_adapter")] = lora.pop(k)


def load_lora_for_models(model: "UnetPatcher", clip: "CLIP", lora: dict[str, torch.Tensor], strength_model: float, strength_clip: float, filename: str = "default", online_mode: bool = False):
    if dynamic_args.nunchaku:
        model.model.diffusion_model.loras.append((filename, strength_model))
        return model, clip

    model_flag: str = type(model.model).__name__ if model is not None else "default"

    unet_keys = model_lora_keys_unet(model.model) if model is not None else {}
    clip_keys = model_lora_keys_clip(clip.cond_stage_model) if clip is not None else {}

    if model.model.diffusion_model.__class__.__name__ == "Anima":
        process_anima(lora)

    lora_unmatch = lora
    lora_unet, lora_unmatch = load_lora(lora_unmatch, unet_keys)
    lora_clip, lora_unmatch = load_lora(lora_unmatch, clip_keys)

    _unmatches = len(lora_unmatch)

    if _unmatches / len(lora) > 0.5:
        logger.warning(f"[LORA] LoRA mismatch for {model_flag}: {filename}")
        return model, clip

    if _unmatches > 0:
        logger.info(f"[LORA] Loading {os.path.basename(filename)} for {model_flag} with {_unmatches} unmatched keys")

    if model is not None and len(lora_unet) > 0:
        new_model = model.clone()
        loaded_keys = new_model.add_patches(filename=filename, patches=lora_unet, strength_patch=strength_model, online_mode=online_mode)
        skipped_keys = [item for item in lora_unet if item not in loaded_keys]
        if len(skipped_keys) / len(lora_unet) > 0.25:
            logger.warning(f"[LORA] Mismatch {filename} for {model_flag}-UNet with {len(skipped_keys)} keys mismatched in {len(loaded_keys)} keys")
        else:
            logger.info(f"[LORA] Loaded {os.path.basename(filename)} for {model_flag}-UNet with {len(loaded_keys)} keys at weight {strength_model} (skipped {len(skipped_keys)} keys) with on_the_fly = {online_mode}")
            model = new_model

    if clip is not None and len(lora_clip) > 0:
        new_clip = clip.clone()
        loaded_keys = new_clip.add_patches(filename=filename, patches=lora_clip, strength_patch=strength_clip, online_mode=online_mode)
        skipped_keys = [item for item in lora_clip if item not in loaded_keys]
        if len(skipped_keys) / len(lora_clip) > 0.25:
            logger.warning(f"[LORA] Mismatch {filename} for {model_flag}-CLIP with {len(skipped_keys)} keys mismatched in {len(loaded_keys)} keys")
        else:
            logger.info(f"[LORA] Loaded {os.path.basename(filename)} for {model_flag}-CLIP with {len(loaded_keys)} keys at weight {strength_clip} (skipped {len(skipped_keys)} keys) with on_the_fly = {online_mode}")
            clip = new_clip

    return model, clip


def load_network(name: str, network_on_disk: network.NetworkOnDisk):
    net = network.Network(name, network_on_disk)
    net.mtime = os.path.getmtime(network_on_disk.filename)
    return net


def load_networks(names: list[str], te_multipliers: list[float] = None, unet_multipliers: list[float] = None):
    current_sd = sd_models.model_data.get_sd_model()
    if current_sd is None:
        return

    loaded_networks.clear()

    unavailable_networks = []
    for name in names:
        if name.lower() in forbidden_network_aliases and available_networks.get(name) is None:
            unavailable_networks.append(name)
        elif available_network_aliases.get(name) is None:
            unavailable_networks.append(name)

    if unavailable_networks:
        update_available_networks_by_names(unavailable_networks)

    networks_on_disk = [available_networks.get(name, None) if name.lower() in forbidden_network_aliases else available_network_aliases.get(name, None) for name in names]
    if any(x is None for x in networks_on_disk):
        list_available_networks()
        networks_on_disk = [available_networks.get(name, None) if name.lower() in forbidden_network_aliases else available_network_aliases.get(name, None) for name in names]

    for network_on_disk, name in zip(networks_on_disk, names):
        try:
            net = load_network(name, network_on_disk)
            net.mentioned_name = name
            network_on_disk.read_hash()
            loaded_networks.append(net)
        except Exception:
            logger.error(f'Failed to load LoRA: "{name}"')
            continue

    online_mode = dynamic_args.online_lora or False

    if dynamic_args.ops.startswith("Mixed") or dynamic_args.ops.endswith("FP8"):
        online_mode = False

    compiled_lora_targets = []
    for n, u, t in zip(networks_on_disk, unet_multipliers, te_multipliers):
        if n is None:
            continue
        compiled_lora_targets.append([n.filename, u, t, online_mode])

    compiled_lora_targets_hash = str(compiled_lora_targets)
    if current_sd.current_lora_hash == compiled_lora_targets_hash:
        return

    current_sd.current_lora_hash = compiled_lora_targets_hash
    current_sd.forge_objects.unet = current_sd.forge_objects_original.unet
    current_sd.forge_objects.clip = current_sd.forge_objects_original.clip

    if dynamic_args.nunchaku:
        current_sd.forge_objects.unet.model.diffusion_model.loras.clear()

    for filename, strength_model, strength_clip, online_mode in compiled_lora_targets:
        lora_sd = load_lora_state_dict(filename)
        current_sd.forge_objects.unet, current_sd.forge_objects.clip = load_lora_for_models(current_sd.forge_objects.unet, current_sd.forge_objects.clip, lora_sd, strength_model, strength_clip, filename=filename, online_mode=online_mode)

    current_sd.forge_objects_after_applying_lora = current_sd.forge_objects.shallow_copy()


def process_network_files(names: Optional[list[str]] = None):
    candidates = []

    for _dir in [shared.cmd_opts.lora_dir, *shared.cmd_opts.lora_dirs]:
        candidates.extend(shared.walk_files(_dir, allowed_extensions=[".pt", ".ckpt", ".safetensors"]))

    for filename in candidates:
        if os.path.isdir(filename):
            continue
        name = os.path.splitext(os.path.basename(filename))[0]
        # if names is provided, only load networks with names in the list
        if names and name not in names:
            continue
        try:
            entry = network.NetworkOnDisk(name, filename)
        except OSError:  # should catch FileNotFoundError and PermissionError etc.
            errors.report(f"Failed to load network {name} from {filename}", exc_info=True)
            continue

        available_networks[name] = entry

        if entry.alias in available_network_aliases:
            forbidden_network_aliases.add(entry.alias.lower())

        available_network_aliases[name] = entry
        available_network_aliases[entry.alias] = entry


def update_available_networks_by_names(names: list[str]):
    process_network_files(names)


def list_available_networks():
    available_networks.clear()
    available_network_aliases.clear()
    available_network_hash_lookup.clear()
    forbidden_network_aliases.clear()
    forbidden_network_aliases.update(["none", "Addams"])

    os.makedirs(shared.cmd_opts.lora_dir, exist_ok=True)

    process_network_files()


re_network_name = re.compile(r"(.*)\s*\([0-9a-fA-F]+\)")


def infotext_pasted(infotext, params: dict):
    if "AddNet Module 1" in [x[1] for x in scripts.scripts_txt2img.infotext_fields]:
        return  # if the other extension is active, it will handle those fields, no need to do anything

    added = []

    for k in params:
        if not k.startswith("AddNet Model "):
            continue

        num = k[13:]

        if params.get("AddNet Module " + num) != "LoRA":
            continue

        name = params.get("AddNet Model " + num)
        if name is None:
            continue

        if m := re_network_name.match(name):
            name = m.group(1)

        multiplier = params.get("AddNet Weight A " + num, "1.0")

        added.append(f"<lora:{name}:{multiplier}>")

    if added:
        params["Prompt"] += "\n" + "".join(added)


extra_network_lora = None

available_networks: dict[str, "network.NetworkOnDisk"] = {}
available_network_aliases: dict[str, "network.NetworkOnDisk"] = {}
available_network_hash_lookup: dict[bytes, "network.NetworkOnDisk"] = {}
forbidden_network_aliases: set[str] = set()
loaded_networks: list["network.Network"] = []

list_available_networks()
