# lora_handling.py
import torch
from typing import Any, Dict, List, Optional, Union
import gradio as gr
from huggingface_hub import ModelCard, HfFileSystem
from flux_app.utilities import calculate_shift, retrieve_timesteps, calculateDuration  # Absolute import
import numpy as np
from PIL import Image
import copy
from flux_app.lora import loras
# FLUX pipeline (continued from previous response)
@torch.inference_mode()
def flux_pipe_call_that_returns_an_iterable_of_images(
    self,
    prompt: Union[str, List[str]] = None,
    prompt_2: Optional[Union[str, List[str]]] = None,
    height: Optional[int] = None,
    width: Optional[int] = None,
    num_inference_steps: int = 28,
    timesteps: List[int] = None,
    guidance_scale: float = 3.5,
    num_images_per_prompt: Optional[int] = 1,
    generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
    latents: Optional[torch.FloatTensor] = None,
    prompt_embeds: Optional[torch.FloatTensor] = None,
    pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
    output_type: Optional[str] = "pil",
    return_dict: bool = True,
    joint_attention_kwargs: Optional[Dict[str, Any]] = None,
    max_sequence_length: int = 512,
    good_vae: Optional[Any] = None,
):
    height = height or self.default_sample_size * self.vae_scale_factor
    width = width or self.default_sample_size * self.vae_scale_factor
    
    self.check_inputs(
        prompt,
        prompt_2,
        height,
        width,
        prompt_embeds=prompt_embeds,
        pooled_prompt_embeds=pooled_prompt_embeds,
        max_sequence_length=max_sequence_length,
    )

    self._guidance_scale = guidance_scale
    self._joint_attention_kwargs = joint_attention_kwargs
    self._interrupt = False

    batch_size = 1 if isinstance(prompt, str) else len(prompt)
    device = self._execution_device

    lora_scale = joint_attention_kwargs.get("scale", None) if joint_attention_kwargs is not None else None
    prompt_embeds, pooled_prompt_embeds, text_ids = self.encode_prompt(
        prompt=prompt,
        prompt_2=prompt_2,
        prompt_embeds=prompt_embeds,
        pooled_prompt_embeds=pooled_prompt_embeds,
        device=device,
        num_images_per_prompt=num_images_per_prompt,
        max_sequence_length=max_sequence_length,
        lora_scale=lora_scale,
    )
    
    num_channels_latents = self.transformer.config.in_channels // 4
    latents, latent_image_ids = self.prepare_latents(
        batch_size * num_images_per_prompt,
        num_channels_latents,
        height,
        width,
        prompt_embeds.dtype,
        device,
        generator,
        latents,
    )
    
    sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
    image_seq_len = latents.shape[1]
    mu = calculate_shift(
        image_seq_len,
        self.scheduler.config.base_image_seq_len,
        self.scheduler.config.max_image_seq_len,
        self.scheduler.config.base_shift,
        self.scheduler.config.max_shift,
    )
    timesteps, num_inference_steps = retrieve_timesteps(
        self.scheduler,
        num_inference_steps,
        device,
        timesteps,
        sigmas,
        mu=mu,
    )
    self._num_timesteps = len(timesteps)

    guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32).expand(latents.shape[0]) if self.transformer.config.guidance_embeds else None

    for i, t in enumerate(timesteps):
        if self.interrupt:
            continue

        timestep = t.expand(latents.shape[0]).to(latents.dtype)

        noise_pred = self.transformer(
            hidden_states=latents,
            timestep=timestep / 1000,
            guidance=guidance,
            pooled_projections=pooled_prompt_embeds,
            encoder_hidden_states=prompt_embeds,
            txt_ids=text_ids,
            img_ids=latent_image_ids,
            joint_attention_kwargs=self.joint_attention_kwargs,
            return_dict=False,
        )[0]

        latents_for_image = self._unpack_latents(latents, height, width, self.vae_scale_factor)
        latents_for_image = (latents_for_image / self.vae.config.scaling_factor) + self.vae.config.shift_factor
        image = self.vae.decode(latents_for_image, return_dict=False)[0]
        yield self.image_processor.postprocess(image, output_type=output_type)[0]
        latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
        torch.cuda.empty_cache()
        
    latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
    latents = (latents / good_vae.config.scaling_factor) + good_vae.config.shift_factor
    image = good_vae.decode(latents, return_dict=False)[0]
    self.maybe_free_model_hooks()
    torch.cuda.empty_cache()
    yield self.image_processor.postprocess(image, output_type=output_type)[0]


def get_huggingface_safetensors(link: str) -> tuple[str, str, str, str, str]:
    """
    Extracts LoRA information from a Hugging Face model card.

    Args:
        link: The Hugging Face model repository URL or ID (e.g., "user/repo" or
              "https://huggingface.co/user/repo").

    Returns:
        A tuple containing:
          - title (str): The repository name.
          - repo (str):  The full repository ID ("user/repo").
          - path (str): The filename of the .safetensors file.
          - trigger_word (str): The instance prompt (trigger word) from the model card.
          - image_url (str): URL of a preview image, if found.

    Raises:
        Exception: If the provided link is not a valid FLUX LoRA repository.
    """
    split_link = link.split("/")
    if len(split_link) == 2:
        model_card = ModelCard.load(link)
        base_model = model_card.data.get("base_model")
        print(base_model)

        # Allows Both FLUX.1-dev and FLUX.1-schnell
        if base_model not in ("black-forest-labs/FLUX.1-dev", "black-forest-labs/FLUX.1-schnell"):
            raise Exception("Flux LoRA Not Found!")

        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]}"
            return split_link[1], link, safetensors_name, trigger_word, image_url  # Return as soon as .safetensors is found
        except Exception as e:
            print(e)
            raise Exception(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA") # More concise exception
    else: #if the links is not complete
        raise Exception(f"You didn't include a link neither a valid Hugging Face repository with a *.safetensors LoRA")

def check_custom_model(link: str) -> tuple[str, str, str, str, str]:
    """
    Checks if the provided link is a Hugging Face URL and extracts LoRA info.

    Args:
        link: The URL or repository ID.

    Returns:
        The same tuple as `get_huggingface_safetensors`.
    """
    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])
    return get_huggingface_safetensors(link)



def create_lora_card(title: str, repo: str, trigger_word: str, image: str) -> str:
    """
    Generates HTML for a LoRA card in the Gradio UI.
    """
    trigger_word_info = (
        f"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"
    )
    return f'''
    <div class="custom_lora_card">
        <span>Loaded custom LoRA:</span>
        <div class="card_internal">
            <img src="{image}" />
            <div>
                <h3>{title}</h3>
                <small>{trigger_word_info}<br></small>
            </div>
        </div>
    </div>
    '''

def add_custom_lora(custom_lora: str, loras: list) -> tuple:
    """Adds a custom LoRA to the list of available LoRAs."""
    if custom_lora:
        try:
            title, repo, path, trigger_word, image = check_custom_model(custom_lora)
            print(f"Loaded custom LoRA: {repo}")
            card = create_lora_card(title, repo, trigger_word, image)

            # Check if the repo is already in the list
            existing_item_index = next((index for (index, item) in enumerate(loras) if item['repo'] == repo), None)
            if existing_item_index is None:  # Use 'is None' for comparison
                new_item = {
                    "image": image,
                    "title": title,
                    "repo": repo,
                    "weights": path,
                    "trigger_word": trigger_word
                }
                print(new_item)
                loras.append(new_item)  # Append to the passed-in loras list
                existing_item_index = len(loras) -1 #the index of new appended item


            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"Error loading LoRA: {e}")  # Debugging
            return gr.update(visible=True, value="Invalid 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() -> tuple:
    """Removes the custom LoRA from the UI."""
    return gr.update(visible=False), gr.update(visible=False), gr.update(), "", None, ""

def prepare_prompt(prompt: str, selected_index: Optional[int], loras: List[Dict]) -> str:
    """Combines the user prompt with the LoRA trigger word."""
    if selected_index is None:
        raise gr.Error("You must select a LoRA before proceeding.🧨")

    selected_lora = loras[selected_index]
    trigger_word = selected_lora.get("trigger_word")  # Use get()

    if trigger_word:
        trigger_position = selected_lora.get("trigger_position", "append")
        if trigger_position == "prepend":
            prompt_mash = f"{trigger_word} {prompt}"
        else:
            prompt_mash = f"{prompt} {trigger_word}"
    else:
        prompt_mash = prompt
    return prompt_mash

def unload_lora_weights(pipe, pipe_i2i):
    """Unloads LoRA weights from both pipelines."""
    if pipe is not None:
        pipe.unload_lora_weights()
    if pipe_i2i is not None:
        pipe_i2i.unload_lora_weights()


def load_lora_weights_into_pipeline(pipe_to_use, lora_path: str, weight_name: Optional[str]):
    """Loads LoRA weights into the specified pipeline."""
    pipe_to_use.load_lora_weights(
        lora_path,
        weight_name=weight_name,
        low_cpu_mem_usage=True
    )


def update_selection(evt: gr.SelectData, width, height, loras):
    """Updates the UI when a LoRA is selected from the gallery."""
    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,
    )