"""
LoRA Loader for WAN 2.2 - references files from lkzd7/WAN2.2_LoraSet_NSFW
"""
import urllib.parse
import re
from huggingface_hub import hf_hub_download

LORA_REPO = "lkzd7/WAN2.2_LoraSet_NSFW"
HF_TOKEN = None

LORA_FILES = [
    "Blink_Squatting_Cowgirl_Position_I2V_HIGH.safetensors",
    "Blink_Squatting_Cowgirl_Position_I2V_LOW.safetensors",
    "PENISLORA_22_i2v_HIGH_e320.safetensors",
    "PENISLORA_22_i2v_LOW_e496.safetensors",
    "Pornmaster_wan 2.2_14b_I2V_bukkake_v1.4_high_noise.safetensors",
    "Pornmaster_wan 2.2_14b_I2V_bukkake_v1.4_low_noise.safetensors",
    "W22_Multiscene_Photoshoot_Softcore_i2v_HN.safetensors",
    "W22_Multiscene_Photoshoot_Softcore_i2v_LN.safetensors",
    "WAN-2.2-I2V-Double-Blowjob-HIGH-v1.safetensors",
    "WAN-2.2-I2V-Double-Blowjob-LOW-v1.safetensors",
    "WAN-2.2-I2V-HandjobBlowjobCombo-HIGH-v1.safetensors",
    "WAN-2.2-I2V-HandjobBlowjobCombo-LOW-v1.safetensors",
    "WAN-2.2-I2V-SensualTeasingBlowjob-HIGH-v1.safetensors",
    "WAN-2.2-I2V-SensualTeasingBlowjob-LOW-v1.safetensors",
    "iGOON_Blink_Blowjob_I2V_HIGH.safetensors",
    "iGOON_Blink_Blowjob_I2V_LOW.safetensors",
    "iGoon - Blink_Front_Doggystyle_I2V_HIGH.safetensors",
    "iGoon - Blink_Front_Doggystyle_I2V_LOW.safetensors",
    "iGoon - Blink_Missionary_I2V_HIGH.safetensors",
    "iGoon - Blink_Missionary_I2V_LOW v2.safetensors",
    "iGoon - Blink_Missionary_I2V_LOW.safetensors",
    "iGoon%20-%20Blink_Back_Doggystyle_HIGH.safetensors",
    "iGoon%20-%20Blink_Back_Doggystyle_LOW.safetensors",
    "iGoon%20-%20Blink_Facial_I2V_HIGH.safetensors",
    "iGoon%20-%20Blink_Facial_I2V_LOW.safetensors",
    "iGoon_Blink_Missionary_I2V_HIGH v2.safetensors",
    "iGoon_Blink_Titjob_I2V_HIGH.safetensors",
    "iGoon_Blink_Titjob_I2V_LOW.safetensors",
    "lips-bj_high_noise.safetensors",
    "lips-bj_low_noise.safetensors",
    "mql_casting_sex_doggy_kneel_diagonally_behind_vagina_wan22_i2v_v1_high_noise.safetensors",
    "mql_casting_sex_doggy_kneel_diagonally_behind_vagina_wan22_i2v_v1_low_noise.safetensors",
    "mql_casting_sex_reverse_cowgirl_lie_front_vagina_wan22_i2v_v1_high_noise.safetensors",
    "mql_casting_sex_reverse_cowgirl_lie_front_vagina_wan22_i2v_v1_low_noise.safetensors",
    "mql_casting_sex_spoon_wan22_i2v_v1_high_noise.safetensors",
    "mql_casting_sex_spoon_wan22_i2v_v1_low_noise.safetensors",
    "mql_doggy_a_wan22_t2v_v1_high_noise .safetensors",
    "mql_doggy_a_wan22_t2v_v1_low_noise.safetensors",
    "mql_massage_tits_wan22_i2v_v1_high_noise.safetensors",
    "mql_massage_tits_wan22_i2v_v1_low_noise.safetensors",
    "mql_panties_aside_wan22_i2v_v1_high_noise.safetensors",
    "mql_panties_aside_wan22_i2v_v1_low_noise.safetensors",
    "mqlspn_a_wan22_t2v_v1_high_noise.safetensors",
    "mqlspn_a_wan22_t2v_v1_low_noise.safetensors",
    "sfbehind_v2.1_high_noise.safetensors",
    "sfbehind_v2.1_low_noise.safetensors",
    "sid3l3g_transition_v2.0_H.safetensors",
    "sid3l3g_transition_v2.0_L.safetensors",
    "wan2.2_i2v_high_ulitmate_pussy_asshole.safetensors",
    "wan2.2_i2v_low_ulitmate_pussy_asshole.safetensors",
    "wan22-mouthfull-140epoc-high-k3nk.safetensors",
    "wan22-mouthfull-152epoc-low-k3nk.safetensors",
]

LORA_PAIRS = {}
for f in LORA_FILES:
    name = urllib.parse.unquote(f).replace(".safetensors", "")
    is_high = bool(re.search(r'(high|HN|_H\b)', name, re.IGNORECASE))
    is_low = bool(re.search(r'(low|LN|_L\b)', name, re.IGNORECASE))
    group = re.sub(r'[\s_-]*(high|low|noise|HN|LN)([\s_-]*noise)?[\s_-]*(v?\d+(\.\d+)?)?\s*$', '', name, flags=re.IGNORECASE).strip()
    group = re.sub(r'[\s_]+$', '', group)
    if group not in LORA_PAIRS:
        LORA_PAIRS[group] = {"HIGH": None, "LOW": None}
    if is_high:
        LORA_PAIRS[group]["HIGH"] = f
    elif is_low:
        LORA_PAIRS[group]["LOW"] = f


def get_lora_choices():
    choices = []
    for group in sorted(LORA_PAIRS.keys()):
        p = LORA_PAIRS[group]
        if p["HIGH"] and p["LOW"]:
            choices.append(group)
        elif p["HIGH"]:
            choices.append(f"{group} (HIGH only)")
        elif p["LOW"]:
            choices.append(f"{group} (LOW only)")
    return choices


def download_lora(group_name):
    if not group_name:
        return None, None
    clean_name = re.sub(r'\s*\(HIGH only\)|\s*\(LOW only\)', '', group_name)
    if clean_name not in LORA_PAIRS:
        return None, None
    pair = LORA_PAIRS[clean_name]
    high_path, low_path = None, None
    if pair["HIGH"]:
        high_path = hf_hub_download(LORA_REPO, pair["HIGH"], token=HF_TOKEN)
    if pair["LOW"]:
        low_path = hf_hub_download(LORA_REPO, pair["LOW"], token=HF_TOKEN)
    return high_path, low_path


def load_lora_to_pipe(pipe, group_name, adapter_name="lora"):
    high_path, low_path = download_lora(group_name)
    if high_path and low_path:
        pipe.load_lora_weights(high_path, adapter_name=f"{adapter_name}_high")
        pipe.load_lora_weights(low_path, adapter_name=f"{adapter_name}_low")
        print(f"Loaded LoRA pair: {group_name}")
        return True
    elif high_path:
        pipe.load_lora_weights(high_path, adapter_name=adapter_name)
        print(f"Loaded LoRA: {group_name}")
        return True
    return False


def unload_lora(pipe):
    try:
        pipe.unload_lora_weights()
    except:
        pass