import gradio as gr
import torch
import spaces
from diffusers import DiffusionPipeline
from pathlib import Path
import gc
import subprocess


subprocess.run('pip cache purge', shell=True)
device = "cuda" if torch.cuda.is_available() else "cpu"
torch.set_grad_enabled(False)


models = [
    "camenduru/FLUX.1-dev-diffusers",
    "black-forest-labs/FLUX.1-schnell",
    "sayakpaul/FLUX.1-merged",
    "John6666/blue-pencil-flux1-v001-fp8-flux",
    "John6666/copycat-flux-test-fp8-v11-fp8-flux",
    "John6666/nepotism-fuxdevschnell-v3aio-fp8-flux",
    "John6666/niji-style-flux-devfp8-fp8-flux",
    "John6666/fluxunchained-artfulnsfw-fut516xfp8e4m3fnv11-fp8-flux",
    "John6666/fastflux-unchained-t5f16-fp8-flux",
    "John6666/the-araminta-flux1a1-fp8-flux",
    "John6666/acorn-is-spinning-flux-v11-fp8-flux",
    "John6666/fluxescore-dev-v10fp16-fp8-flux",
    # "",
]


num_loras = 3


def is_repo_name(s):
    import re
    return re.fullmatch(r'^[^/,\s\"\']+/[^/,\s\"\']+$', s)


def is_repo_exists(repo_id):
    from huggingface_hub import HfApi
    api = HfApi()
    try:
        if api.repo_exists(repo_id=repo_id): return True
        else: return False
    except Exception as e:
        print(f"Error: Failed to connect {repo_id}. ")
        print(e)
        return True # for safe


def clear_cache():
    torch.cuda.empty_cache()
    gc.collect()


def get_repo_safetensors(repo_id: str):
    from huggingface_hub import HfApi
    api = HfApi()
    try:
        if not is_repo_name(repo_id) or not is_repo_exists(repo_id): return gr.update(value="", choices=[])
        files = api.list_repo_files(repo_id=repo_id)
    except Exception as e:
        print(f"Error: Failed to get {repo_id}'s info. ")
        print(e)
        return gr.update(choices=[])
    files = [f for f in files if f.endswith(".safetensors")]
    if len(files) == 0: return gr.update(value="", choices=[])
    else: return gr.update(value=files[0], choices=files)


# Initialize the base model
base_model = models[0]
pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
last_model = models[0]

def change_base_model(repo_id: str, progress=gr.Progress(track_tqdm=True)):
    global pipe
    global last_model
    try:
        if repo_id == last_model or not is_repo_name(repo_id) or not is_repo_exists(repo_id): return
        progress(0, desc=f"Loading model: {repo_id}")
        clear_cache()
        pipe = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
        last_model = repo_id
        progress(1, desc=f"Model loaded: {repo_id}")
    except Exception as e:
        print(e)
    return gr.update(visible=True)


def compose_lora_json(lorajson: list[dict], i: int, name: str, scale: float, filename: str, trigger: str):
    lorajson[i]["name"] = str(name) if name != "None" else ""
    lorajson[i]["scale"] = float(scale)
    lorajson[i]["filename"] = str(filename)
    lorajson[i]["trigger"] = str(trigger)
    return lorajson


def is_valid_lora(lorajson: list[dict]):
    valid = False
    for d in lorajson:
        if "name" in d.keys() and d["name"] and d["name"] != "None": valid = True
    return valid


def get_trigger_word(lorajson: list[dict]):
    trigger = ""
    for d in lorajson:
        if "name" in d.keys() and d["name"] and d["name"] != "None" and d["trigger"]:
            trigger += ", " + d["trigger"]
    return trigger

# https://huggingface.co/docs/diffusers/v0.23.1/en/api/loaders#diffusers.loaders.LoraLoaderMixin.fuse_lora
# https://github.com/huggingface/diffusers/issues/4919
def fuse_loras(pipe, lorajson: list[dict]):
    if not lorajson or not isinstance(lorajson, list): return
    a_list = []
    w_list = []
    for d in lorajson:
        if not d or not isinstance(d, dict) or not d["name"] or d["name"] == "None": continue
        k = d["name"]
        if is_repo_name(k) and is_repo_exists(k):
            a_name = Path(k).stem
            pipe.load_lora_weights(k, weight_name=d["filename"], adapter_name = a_name)
        elif not Path(k).exists():
            print(f"LoRA not found: {k}")
            continue
        else:
            w_name = Path(k).name
            a_name = Path(k).stem
            pipe.load_lora_weights(k, weight_name = w_name, adapter_name = a_name)
        a_list.append(a_name)
        w_list.append(d["scale"])
    if not a_list: return
    pipe.set_adapters(a_list, adapter_weights=w_list)
    pipe.fuse_lora(adapter_names=a_list, lora_scale=1.0)
    #pipe.unload_lora_weights()


change_base_model.zerogpu = True
fuse_loras.zerogpu = True


def description_ui():
    gr.Markdown(
        """
- Mod of [multimodalart/flux-lora-the-explorer](https://huggingface.co/spaces/multimodalart/flux-lora-the-explorer),
 [gokaygokay/FLUX-Prompt-Generator](https://huggingface.co/spaces/gokaygokay/FLUX-Prompt-Generator).
"""
    )