import os
import json
import copy
import time
import requests
import random
import logging
import numpy as np
import spaces
from typing import Any, Dict, List, Optional, Union
from civitai_utils import get_civitai_safetensors, LORA_CHECKPOINTS_CACHE

import torch
from PIL import Image
import gradio as gr

from diffusers import (
    DiffusionPipeline,
    AutoencoderKL,
    ZImagePipeline
)

from huggingface_hub import (
    hf_hub_download,
    HfFileSystem,
    ModelCard,
    snapshot_download)

from diffusers.utils import load_image
from typing import Iterable

# Load loras as list of dictionaries
loras = []
with open(os.path.join(os.getcwd(), "loras.json"), "r") as f:
    loras = json.load(f)

dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"
base_model = "Tongyi-MAI/Z-Image-Turbo"

print(f"Loading {base_model} pipeline...")

# Initialize Pipeline
pipe = ZImagePipeline.from_pretrained(
    base_model,
    torch_dtype=dtype,
    low_cpu_mem_usage=False,
).to(device)

# ======== AoTI compilation + FA3 ========
# As per reference for optimization
try:
    print("Applying AoTI compilation and FA3...")
    pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
    spaces.aoti_blocks_load(pipe.transformer.layers, "zerogpu-aoti/Z-Image", variant="fa3")
    print("Optimization applied successfully.")
except Exception as e:
    print(f"Optimization warning: {e}. Continuing with standard pipeline.")

MAX_SEED = np.iinfo(np.int32).max
SEED_RANDOM = -1

class calculateDuration:
    def __init__(self, activity_name=""):
        self.activity_name = activity_name

    def __enter__(self):
        self.start_time = time.time()
        return self
    
    def __exit__(self, exc_type, exc_value, traceback):
        self.end_time = time.time()
        self.elapsed_time = self.end_time - self.start_time
        if self.activity_name:
            print(f"Elapsed time for {self.activity_name}: {self.elapsed_time:.6f} seconds")
        else:
            print(f"Elapsed time: {self.elapsed_time:.6f} seconds")


def update_selection(evt: gr.SelectData, width, height):
    selected_lora = loras[evt.index]
    new_placeholder = f"Type a prompt for {selected_lora['title']}"
    lora_repo = selected_lora["repo"]
    updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✅"
    if "aspect" in selected_lora:
        if selected_lora["aspect"] == "portrait":
            width = 768
            height = 1024
        elif selected_lora["aspect"] == "landscape":
            width = 1024
            height = 768
        else:
            width = 1024
            height = 1024
    return (
        gr.update(placeholder=new_placeholder),
        updated_text,
        evt.index,
        width,
        height,
    )


def load_lora_from_hub(lora: dict, lora_scale: float):
    """Load LoRA weights from huggingface hub"""
    with calculateDuration(f"Loading LoRA weights for {lora.get('title')}"):
        try:
            pipe.load_lora_weights(
                lora.get("repo", ""), 
                weight_name=lora.get("weights", None), 
                adapter_name="default",
                low_cpu_mem_usage=True
            )
            # Set adapter scale
            pipe.set_adapters(["default"], adapter_weights=[lora_scale])
        except Exception as e:
            print(f"Error loading LoRA: {e}")
            gr.Warning("Failed to load LoRA weights. Generating with base model.")


def load_local_lora(lora: dict, lora_scale: float):
    """Load LoRA weights from local cache folder"""
    with calculateDuration(f"Loading LoRA weights for {lora.get('title')}"):
        try:
            pipe.load_lora_weights(
                LORA_CHECKPOINTS_CACHE,
                cache_dir=LORA_CHECKPOINTS_CACHE,
                adapter_name="local_lora",
                weight_name=lora.get("weights", None),
                local_files_only=True,
                low_cpu_mem_usage=True
            )
            # Set adapter scale
            pipe.set_adapters(["local_lora"], adapter_weights=[lora_scale])
        except Exception as e:
            print(f"Error loading LoRA: {e}")
            gr.Warning("Failed to load LoRA weights. Generating with base model.")


@spaces.GPU
def run_lora(prompt, image_strength, cfg_scale, steps, selected_index, seed, width, height, lora_scale):
    # Clean up previous LoRAs in both cases
    with calculateDuration("Unloading LoRA"):
        pipe.unload_lora_weights()
    
    prompt_mash = prompt
    
    # Check if a LoRA is selected
    if selected_index is not None and selected_index < len(loras):
        selected_lora = loras[selected_index]
        trigger_word = selected_lora["trigger_word"]
        
        # Prepare Prompt with Trigger Word
        if len(trigger_word):
            if "trigger_position" in selected_lora:
                if selected_lora["trigger_position"] == "prepend":
                    prompt_mash = f"{trigger_word} {prompt}"
                else:
                    prompt_mash = f"{prompt} {trigger_word}"
            else:
                prompt_mash = f"{trigger_word} {prompt}"

        # Special handling of lora loading if there's a civitai key
        if selected_lora.get("src") == "civitai":
            load_local_lora(selected_lora, lora_scale)
        else:
            load_lora_from_hub(selected_lora, lora_scale)
    else:
        # Base Model Case
        print("No LoRA selected. Running with Base Model.")
        prompt_mash = prompt

    seed_val = seed
    if seed_val == SEED_RANDOM:
        seed_val = random.randint(0, MAX_SEED)
    generator = torch.Generator(device=device).manual_seed(seed_val)

    with calculateDuration("Generating image"):
        # For Turbo models, guidance_scale is typically 0.0
        forced_guidance = 0.0 # Turbo mode
        
        final_image = pipe(
            prompt=prompt_mash,
            height=int(height),
            width=int(width),
            num_inference_steps=int(steps),
            guidance_scale=forced_guidance,
            generator=generator,
        ).images[0]
        
    yield final_image, seed_val


def get_huggingface_safetensors(link) -> dict:
    split_link = link.split("/")
    if(len(split_link) == 2):
            model_card = ModelCard.load(link)
            base_model_list = model_card.data.get("base_model")

            # Relaxed check to allow Z-Image or Flux or others, assuming user knows what they are doing
            # or specifically check for Z-Image-Turbo
            if base_model_list[0] not in ["Tongyi-MAI/Z-Image-Turbo", "black-forest-labs/FLUX.1-dev"]:
                # Just a warning instead of error to allow experimentation
                print("Warning: Base model might not match.")
                
            image_path = model_card.data.get("widget", [{}])[0].get("output", {}).get("url", None)
            trigger_word = model_card.data.get("instance_prompt", "")
            image_url = f"https://huggingface.co/{link}/resolve/main/{image_path}" if image_path else None
            fs = HfFileSystem()
            try:
                list_of_files = fs.ls(link, detail=False)
                for file in list_of_files:
                    if(file.endswith(".safetensors")):
                        safetensors_name = file.split("/")[-1]
                    if (not image_url and file.lower().endswith((".jpg", ".jpeg", ".png", ".webp"))):
                        image_elements = file.split("/")
                        image_url = f"https://huggingface.co/{link}/resolve/main/{image_elements[-1]}"
            except Exception as e:
                print(e)
                gr.Warning(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
                raise Exception(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")
            
            lora_info = {
                "image": image_url,
                "title": split_link[1],
                "repo": link,
                "weights": safetensors_name,
                "trigger_word": trigger_word
            }
            return lora_info


def check_custom_model(link) -> dict:
    if(link.startswith("https://")):
        if(link.startswith("https://huggingface.co") or link.startswith("https://www.huggingface.co")):
            link_split = link.split("huggingface.co/")
            return get_huggingface_safetensors(link_split[1])
        elif "civitai" in link:
            return get_civitai_safetensors(link)
    else: 
        return {}


def seed_radio_handler(selected_option, used_seed_value):
    if selected_option == "Reuse":
        return gr.update(value=used_seed_value, interactive=True)
    else:  # "Randomize"
        return gr.update(value=SEED_RANDOM, interactive=False)


def add_custom_lora(custom_lora):
    global loras
    if(custom_lora):
        try:
            lora_info = check_custom_model(custom_lora)
            repo = lora_info.get("repo")
            image = lora_info.get("image")
            trigger_word = lora_info.get("trigger_word")
            path = lora_info.get("weights")
            title = lora_info.get("title")
            src = lora_info.get("src")

            repo = "civitai" if src == "civitai" else lora_info.get("repo")
            print(f"Loaded custom LoRA: {repo}")
            card = f'''
            <div class="custom_lora_card">
              <span>Loaded custom LoRA:</span>
              <div class="card_internal">
                <img src="{image}" />
                <div>
                    <h3>{title}</h3>
                    <small>{"Using: <code><b>"+trigger_word+"</code></b> as the trigger word" if trigger_word else "No trigger word found. If there's a trigger word, include it in your prompt"}<br></small>
                </div>
              </div>
            </div>
            '''
            existing_item_index = next((index for (index, item) in enumerate(loras) if item['title'] == title), None)
            if not existing_item_index:
                print(lora_info)
                existing_item_index = len(loras)
                loras.append(lora_info)
        
            return gr.update(visible=True, value=card), gr.update(visible=True), gr.Gallery(selected_index=None), f"Custom: {path}", existing_item_index, trigger_word
        except Exception as e:
            print(f"add_custom_lora() Exception: {e}")
            gr.Warning(f"Invalid LoRA: either you entered an invalid link, or a non-supported LoRA")
            return gr.update(visible=True, value=f"Invalid LoRA: either you entered an invalid link, a non-supported LoRA"), gr.update(visible=False), gr.update(), "", None, ""
    else:
        return gr.update(visible=False), gr.update(visible=False), gr.update(), "", None, ""


def remove_custom_lora():
    return gr.update(visible=False), gr.update(visible=False), gr.update(), "", None, ""


with gr.Blocks(delete_cache=(60, 60)) as demo:
    title = gr.HTML(
        """<h1>Z Image Turbo LoRA DLC 🧪</h1>""",
        elem_id="title",
    )
    selected_index = gr.State(None)
    with gr.Row():
        with gr.Column(scale=2):
            prompt = gr.Textbox(label="Enter Prompt", lines=2, placeholder="✦︎ Choose the LoRA and type the prompt (LoRA = None → Base Model = Active)")
            generate_button = gr.Button("Generate", variant="primary", elem_id="gen_btn")
            selected_info = gr.Markdown("### No LoRA Selected (Base Model)")
            with gr.Group():
                custom_lora = gr.Textbox(label="Enter Custom LoRA", placeholder="Paste the LoRA url & press Enter (e.g. https://huggingface.co/tarn59/pixel_art_style_lora_z_image_turbo )")
                gr.Markdown("[Check the list of Z-Image LoRA's](https://huggingface.co/models?other=base_model:adapter:Tongyi-MAI/Z-Image-Turbo)", elem_id="lora_list")
            custom_lora_info = gr.HTML(visible=False)
            custom_lora_button = gr.Button("Remove custom LoRA", visible=False)
            with gr.Row():
                with gr.Accordion("Advanced Settings", open=False):
                    with gr.Row():
                        image_strength = gr.Slider(label="Denoise Strength", info="Ignored for Z-Image-Turbo", minimum=0.1, maximum=1.0, step=0.01, value=0.75, visible=False)
                    with gr.Column():
                        with gr.Row():
                            cfg_scale = gr.Slider(label="CFG Scale", info="Forced to 0.0 for Turbo", minimum=0, maximum=20, step=0.5, value=0.0, interactive=False)
                            steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=9)
                        
                        with gr.Row():
                            width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=1024)
                            height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=1024)
                        
                        with gr.Row():
                            seed_radio = gr.Radio(["Reuse", "Randomize"], label="Seed Options", info="How the seed is generated", value="Randomize", interactive=True)
                            seed_input = gr.Number(value=SEED_RANDOM, label="Seed", precision=0, interactive=False)

                        lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=3, step=0.01, value=0.95)

        with gr.Column(scale=3):
            # Create a scrolling horizontal gallery
            with gr.Row(variant="compact"):
                gallery = gr.Gallery(
                    [(item["image"], item["title"]) for item in loras],
                    label="Z-Image LoRAs",
                    columns=len(loras),
                    object_fit="cover",
                    height=132,
                    elem_id="gallery",
                    allow_preview=False
                )
            result = gr.Image(label="Generated Image", format="jpg", height=1024)
            used_seed = gr.Number(label="Seed used", interactive=False) 

    seed_radio.change(seed_radio_handler, show_progress="hidden", inputs=[seed_radio, used_seed], outputs=seed_input)

    gallery.select(
        update_selection,
        inputs=[width, height],
        outputs=[prompt, selected_info, selected_index, width, height]
    )
    custom_lora.input(
        add_custom_lora,
        inputs=[custom_lora],
        outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, prompt]
    )
    custom_lora_button.click(
        remove_custom_lora,
        outputs=[custom_lora_info, custom_lora_button, gallery, selected_info, selected_index, custom_lora]
    )
    gr.on(
        triggers=generate_button.click,
        fn=run_lora,
        inputs=[prompt, image_strength, cfg_scale, steps, selected_index, seed_input, width, height, lora_scale],
        outputs=[result, used_seed]
    )

demo.queue()
demo.launch(mcp_server=True, ssr_mode=False, show_error=True)