import glob
import os
import random
import re
import subprocess
import sys

import spaces
import torch
import gradio as gr
from huggingface_hub import HfApi, ModelCard, hf_hub_download, login
from diffusers import Krea2Pipeline

# Optional HF token (needed for private repos or higher rate limits)
TOKEN = os.environ.get("HF_TOKEN")
if TOKEN:
    login(token=TOKEN)
api = HfApi(token=TOKEN)

DTYPE = torch.bfloat16
BASE_MODEL = "krea/Krea-2-Turbo"
MAX_SEED = 2**31 - 1

TARGET_RANK = 64
BLOCK_NAME = "Krea2TransformerBlock"
ARTIFACT_REPO = "multimodalart/krea2-aoti-kernels"
ARTIFACT_FILE = f"Krea2TransformerBlock-lora-r{TARGET_RANK}/package.pt2"

LORA_REPOS = [
    ("krea/Krea-2-LoRA-retroanime", "Retro Anime"),
    ("krea/Krea-2-LoRA-rainywindow", "Rainy Window"),
    ("krea/Krea-2-LoRA-vintagetarot", "Vintage Tarot"),
    ("krea/Krea-2-LoRA-sunsetblur", "Sunset Blur"),
    ("krea/Krea-2-LoRA-dotmatrix", "Dot Matrix"),
    ("krea/Krea-2-LoRA-neondrip", "Neon Drip"),
    ("krea/Krea-2-LoRA-darkbrush", "Dark Brush"),
    ("krea/Krea-2-LoRA-kidsdrawing", "Kids Drawing"),
    ("krea/Krea-2-LoRA-softwatercolor", "Soft Watercolor Art Deco"),
]
DEFAULT_SCALE = 1.0


def _read_trigger(repo: str) -> str:
    try:
        text = ModelCard.load(repo, token=TOKEN).text
        m = re.search(r"[Tt]rigger word[:\*\s]*`([^`]+)`", text)
        if m:
            return m.group(1).strip()
    except Exception:
        pass
    return ""


def resolve_custom_lora(repo: str):
    files = api.list_repo_files(repo)
    safes = [f for f in files if f.endswith(".safetensors")]
    if not safes:
        raise gr.Error(f"No .safetensors weights found in {repo}.")
    weight_name = max(safes, key=lambda f: ("/" not in f, "lora" in f.lower()))
    return weight_name, _read_trigger(repo)


pipe = Krea2Pipeline.from_pretrained(BASE_MODEL, torch_dtype=DTYPE)

LORAS = []
for repo, title in LORA_REPOS:
    key = repo.split("LoRA-")[-1]
    try:
        weight_name, trigger = resolve_custom_lora(repo)
    except Exception as e:
        print(f"[lora] resolve failed for {repo}: {e}")
        continue
    LORAS.append(
        {
            "key": key,
            "title": title,
            "repo": repo,
            "weight_name": weight_name,
            "trigger": trigger,
            "scale": DEFAULT_SCALE,
        }
    )

print(f"[lora] {len(LORAS)}/{len(LORA_REPOS)} styles ready: {[l['key'] for l in LORAS]}")

pipe.transformer.enable_lora_hotswap(target_rank=TARGET_RANK)
FIRST = LORAS[0]
pipe.transformer.load_lora_adapter(
    FIRST["repo"], weight_name=FIRST["weight_name"], adapter_name="style", token=TOKEN
)
pipe.transformer.set_adapters("style", weights=1.0)
pipe.to("cuda")

CURRENT = {
    "repo": FIRST["repo"],
    "weight_name": FIRST["weight_name"],
    "scale": None,
}

AOTI_MODEL = None


def _full_weights(block, scale: float):
    w = {}
    for n, p in block.named_parameters(remove_duplicate=False):
        if scale != 1.0 and ".lora_B." in n:
            w[n] = p * scale
        else:
            w[n] = p
    for n, b in block.named_buffers(remove_duplicate=False):
        w[n] = b
    return w


def _patch_all(aoti_model, scale: float):
    n = 0
    for block in pipe.transformer.modules():
        if block.__class__.__name__ == BLOCK_NAME:
            block.forward = aoti_model.with_weights(_full_weights(block, scale))
            n += 1
    return n


def _load_artifact():
    global AOTI_MODEL
    from spaces.zero.torch.aoti import LazyAOTIModel

    pt2 = hf_hub_download(
        repo_id=ARTIFACT_REPO,
        filename=ARTIFACT_FILE,
        repo_type="dataset",
        token=TOKEN,
    )
    AOTI_MODEL = LazyAOTIModel(pt2)
    print(f"AoTI artifact loaded: {ARTIFACT_FILE}")


try:
    _load_artifact()
except Exception as e:
    print(f"AoTI artifact unavailable ({e}); running eager.")


def _ensure_adapter(repo: str, weight_name: str) -> bool:
    if CURRENT["repo"] == repo and CURRENT["weight_name"] == weight_name:
        return False
    pipe.transformer.load_lora_adapter(
        repo,
        weight_name=weight_name,
        adapter_name="style",
        hotswap=True,
        token=TOKEN,
    )
    pipe.transformer.set_adapters("style", weights=1.0)
    CURRENT["repo"] = repo
    CURRENT["weight_name"] = weight_name
    CURRENT["scale"] = None
    return True


gallery_items = []
for lora in LORAS:
    try:
        img = hf_hub_download(lora["repo"], "images/05_turbo.png")
    except Exception:
        img = None
    gallery_items.append((img, lora["title"]))


def _mash_prompt(prompt: str, trigger: str) -> str:
    prompt = (prompt or "").strip()
    if trigger and trigger.lower() not in prompt.lower():
        return f"{prompt}, {trigger}" if prompt else trigger
    return prompt


CHIP = (
    "background:#171717;border:1px solid #262626;border-radius:5px;padding:2px 8px;"
    "font-family:'JetBrains Mono',ui-monospace,monospace;font-size:12px;color:#d4d4d5;"
)


def _info_html(title: str, trigger: str, scale=None, thumb=None) -> str:
    parts = [
        '<div style="display:flex;flex-wrap:wrap;align-items:center;gap:8px;'
        'font-size:14px;color:#a3a3a3;">'
    ]
    if thumb:
        parts.append(
            f'<img src="{thumb}" style="width:40px;height:40px;border-radius:6px;'
            'object-fit:cover;border:1px solid #262626;" />'
        )
    parts.append(f'<span style="font-weight:600;color:#f5f5f5;">{title}</span>')
    if trigger:
        parts.append(f'<span>Trigger</span><span style="{CHIP}">{trigger}</span>')
    if scale is not None:
        parts.append(f'<span>weight</span><span style="{CHIP}">{scale}</span>')
    parts.append("</div>")
    return "".join(parts)


def _preview_path(repo: str):
    try:
        files = api.list_repo_files(repo)
    except Exception:
        return None
    imgs = sorted(
        f for f in files if f.lower().endswith((".jpg", ".jpeg", ".png", ".webp"))
    )
    return imgs[0] if imgs else None


_REPO_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*/[A-Za-z0-9][A-Za-z0-9._-]*$")
_NO_LORA_HTML = '<div style="color:#737373;font-size:14px;">No LoRA loaded</div>'


def preview_custom_lora(repo: str):
    repo = (repo or "").strip()
    if not repo:
        return (
            None,
            _NO_LORA_HTML,
            gr.update(placeholder="Select a LoRA, then type a prompt"),
        )
    if not _REPO_RE.match(repo):
        return gr.update(), gr.update(), gr.update()
    try:
        _weight, trigger = resolve_custom_lora(repo)
    except Exception:
        return None, _NO_LORA_HTML, gr.update()
    img = _preview_path(repo)
    thumb = f"https://huggingface.co/{repo}/resolve/main/{img}" if img else None
    info = _info_html(repo, trigger, thumb=thumb)
    placeholder = (
        f'Type a prompt. "{trigger}" is added automatically.'
        if trigger
        else "Type a prompt"
    )
    return None, info, gr.update(placeholder=placeholder)


def update_selection(evt: gr.SelectData):
    lora = LORAS[evt.index]
    chip = CHIP
    info = (
        '<div style="display:flex;flex-wrap:wrap;align-items:center;gap:8px;'
        'font-size:14px;color:#a3a3a3;">'
        f'<span style="font-weight:600;color:#f5f5f5;">{lora["title"]}</span>'
        f'<span>Trigger</span><span style="{chip}">{lora["trigger"]}</span>'
        f'<span>weight</span><span style="{chip}">{lora["scale"]}</span>'
        "</div>"
    )
    placeholder = f'Type a prompt. "{lora["trigger"]}" is added automatically.'
    return evt.index, info, gr.update(placeholder=placeholder), gr.update(
        value=lora["scale"]
    )


@spaces.GPU(duration=120, size="large")
def _generate(
    repo,
    weight_name,
    full_prompt,
    scale,
    steps,
    guidance,
    width,
    height,
    seed,
):
    scale = float(scale)
    swapped = _ensure_adapter(repo, weight_name)
    if AOTI_MODEL is not None:
        if swapped or CURRENT["scale"] != scale:
            _patch_all(AOTI_MODEL, scale)
            CURRENT["scale"] = scale
    else:
        pipe.transformer.set_adapters("style", weights=scale)
    generator = torch.Generator("cuda").manual_seed(int(seed))
    return pipe(
        prompt=full_prompt,
        num_inference_steps=int(steps),
        guidance_scale=float(guidance),
        width=int(width),
        height=int(height),
        generator=generator,
    ).images[0]


def run_lora(
    prompt,
    custom_lora,
    selected_index,
    lora_scale,
    steps,
    guidance,
    width,
    height,
    seed,
    randomize,
    progress=gr.Progress(track_tqdm=True),
):
    if custom_lora and custom_lora.strip():
        repo = custom_lora.strip()
        weight_name, trigger = resolve_custom_lora(repo)
    elif selected_index is not None:
        lora = LORAS[selected_index]
        repo, weight_name, trigger = (
            lora["repo"],
            lora["weight_name"],
            lora["trigger"],
        )
    else:
        raise gr.Error(
            "Select a LoRA from the gallery or enter a custom LoRA path."
        )

    if randomize:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)
    full_prompt = _mash_prompt(prompt, trigger)

    try:
        image = _generate(
            repo,
            weight_name,
            full_prompt,
            lora_scale,
            steps,
            guidance,
            width,
            height,
            seed,
        )
    except Exception as e:
        msg = str(e)
        if (
            "hotswap" in msg.lower()
            or "rank" in msg.lower()
            or "target" in msg.lower()
        ):
            raise gr.Error(
                f"This LoRA isn't hotswap-compatible (needs rank <= {TARGET_RANK} "
                f"and the same target layers as the Krea-2 built-ins). Details: {msg}"
            )
        raise
    return image, seed


KREA_ACCENT = "#000000"

theme = gr.themes.Base(
    
    primary_hue=gr.themes.colors.gray,
    neutral_hue=gr.themes.colors.gray,
    font=[
        gr.themes.GoogleFont("Inter"),
        "ui-sans-serif",
        "system-ui",
        "sans-serif",
    ],
    font_mono=[
        gr.themes.GoogleFont("JetBrains Mono"),
        "ui-monospace",
        "monospace",
    ],
).set(
    
    body_background_fill="#000000",
    body_background_fill_dark="#000000",
    body_text_color="#d7d7d7",

    
    background_fill_primary="#000000",
    background_fill_secondary="#000000",
    block_background_fill="#000000",
    block_border_color="#121212",
    block_border_width="1px",
    block_label_background_fill="#000000",
    block_label_text_color="#d7d7d7",
    block_title_text_color="#d7d7d7",
    border_color_primary="#121212",

    
    input_background_fill="#000000",
    input_border_color="#121212",
    input_border_color_focus="#121212",

    
    button_primary_background_fill="#000000",
    button_primary_background_fill_hover="#121212",
    button_primary_text_color="#d7d7d7",
    button_primary_border_color="#121212",

    
    button_secondary_background_fill="#000000",
    button_secondary_background_fill_hover="#121212",
    button_secondary_text_color="#d7d7d7",
    button_secondary_border_color="#121212",

    
    slider_color="#000000",
)

CSS = """
.gradio-container { background: #000 !important; }
#page { max-width: 1120px; margin: 0 auto; padding: 4px 8px 32px; }

#krea-header {
  padding: 32px 6px 22px;
  border-bottom: 1px solid #1a1a1a;
  margin-bottom: 22px;
}
#krea-header .eyebrow {
  font-family: 'JetBrains Mono', ui-monospace, monospace;
  font-size: 11px;
  letter-spacing: 0.24em;
  text-transform: uppercase;
  color: #737373;
}
#krea-header h1 {
  font-size: 42px;
  font-weight: 600;
  letter-spacing: -0.025em;
  line-height: 1.05;
  margin: 10px 0 6px;
  color: #fff;
}
#krea-header .subtitle {
  font-size: 15px;
  line-height: 1.5;
  color: #a3a3a3;
  margin: 0;
  max-width: 60ch;
}
#krea-header .meta {
  margin-top: 18px;
  display: flex;
  justify-content: space-between;
  align-items: center;
  flex-wrap: wrap;
  gap: 12px;
}
#krea-header .badges { display: flex; gap: 8px; }
#krea-header .badge {
  font-family: 'JetBrains Mono', ui-monospace, monospace;
  font-size: 10px;
  letter-spacing: 0.12em;
  text-transform: uppercase;
  color: #d4d4d5;
  border: 1px solid #262626;
  border-radius: 999px;
  padding: 4px 10px;
}
#krea-header .links { display: flex; gap: 16px; }
#krea-header .links a {
  font-family: 'JetBrains Mono', ui-monospace, monospace;
  font-size: 11px;
  letter-spacing: 0.08em;
  text-transform: uppercase;
  color: #737373;
  text-decoration: none;
  transition: color 0.15s ease;
}
#krea-header .links a:hover { color: #f5f5f5; }

#generate-btn { font-weight: 600; letter-spacing: 0.01em; }

#result-image { min-height: 420px; border-radius: 10px; overflow: hidden; }

.gradio-container code,
.gradio-container .prose code {
  background: #171717 !important;
  color: #d4d4d5 !important;
  border: 1px solid #262626 !important;
  border-radius: 5px !important;
  padding: 2px 7px !important;
  font-family: 'JetBrains Mono', ui-monospace, monospace !important;
  font-size: 0.85em !important;
}

footer { display: none !important; }
.gradio-container .prose a { color: """ + KREA_ACCENT + """; }
"""

HEADER = """ """

with gr.Blocks(title="Krea 2 LoRA Explorer", css=CSS, theme=theme) as demo:
    with gr.Column(elem_id="page"):
        gr.HTML(HEADER)
        selected_index = gr.State(None)

        with gr.Column():
            result = gr.Image(label="", format="png", elem_id="result-image", height=690)

            with gr.Column():
                generate_button = gr.Button(
                    "", variant="primary", elem_id="generate-btn"
                )
                prompt = gr.Textbox(
                    label="",
                    lines=40,
                    placeholder="",
                )
            
            with gr.Accordion(open=True):
                lora_scale = gr.Slider(0.0, 2.0, value=0.4, step=0.1, label="")
                steps = gr.Slider(1, 30, value=8, step=1, label="")
                guidance = gr.Slider(0.0, 10.0, value=0.4, step=0.1, label="")
                width = gr.Slider(1024, 2048, value=1024, step=512, label="")
                height = gr.Slider(1024, 2048, value=1024, step=512, label="")
                seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="")
                randomize = gr.Checkbox(value=True, label="")

            with gr.Column():
                selected_info = gr.HTML("")
                gallery = gr.Gallery(
                    value=gallery_items,
                    label="",
                    columns=3,
                    height="auto",
                    object_fit="cover",
                    allow_preview=False,
                    elem_id="lora-gallery",
                )

                with gr.Column():
                    custom_lora = gr.Textbox(
                        label="",
                        info="",
                        placeholder="",
                    )
                    custom_info = gr.HTML(_NO_LORA_HTML)
                    gr.Markdown("")

        gallery.select(
            update_selection,
            outputs=[selected_index, selected_info, prompt, lora_scale],
        )
        custom_lora.change(
            preview_custom_lora,
            custom_lora,
            [selected_index, custom_info, prompt],
        )
        custom_lora.submit(
            preview_custom_lora,
            custom_lora,
            [selected_index, custom_info, prompt],
        )

        inputs = [
            prompt,
            custom_lora,
            selected_index,
            lora_scale,
            steps,
            guidance,
            width,
            height,
            seed,
            randomize,
        ]
        gr.on(
            [generate_button.click, prompt.submit],
            run_lora,
            inputs,
            [result, seed],
        )

if __name__ == "__main__":
    demo.launch()