import os
import gradio as gr
import torch
from diffusers import StableDiffusionXLPipeline
from huggingface_hub import HfApi, login
from huggingface_hub.utils import validate_repo_id, HfHubHTTPError
import tempfile
import re
import json
import glob
import gdown
import requests
import subprocess
from urllib.parse import urlparse, unquote
from pathlib import Path

# ---------------------- UTILITY FUNCTIONS ----------------------

def get_save_dtype(save_precision_as):
    """Determines the save dtype based on the user's choice."""
    if save_precision_as == "fp16":
        return torch.float16
    elif save_precision_as == "bf16":
        return torch.bfloat16
    elif save_precision_as == "float":
        return torch.float32  # Using float32 for "float" option
    else:
        return None

def determine_load_checkpoint(model_to_load):
    """Determines if the model to load is a checkpoint or a Diffusers model."""
    if model_to_load.endswith('.ckpt') or model_to_load.endswith('.safetensors'):
        return True
    elif os.path.isdir(model_to_load):
        required_folders = {"unet", "text_encoder", "text_encoder_2", "tokenizer", "tokenizer_2", "scheduler", "vae"}
        if required_folders.issubset(set(os.listdir(model_to_load))) and os.path.isfile(os.path.join(model_to_load, "model_index.json")):
            return False
    return None  # handle this case as required

def increment_filename(filename):
    """Increments the filename to avoid overwriting existing files."""
    base, ext = os.path.splitext(filename)
    counter = 1
    while os.path.exists(filename):
        filename = f"{base}({counter}){ext}"
        counter += 1
    return filename

def create_model_repo(api, user, orgs_name, model_name, make_private=False):
    """Creates a Hugging Face model repository if it doesn't exist."""
    if orgs_name == "":
        repo_id = user["name"] + "/" + model_name.strip()
    else:
        repo_id = orgs_name + "/" + model_name.strip()

    try:
        validate_repo_id(repo_id)
        api.create_repo(repo_id=repo_id, repo_type="model", private=make_private)
        print(f"Model repo '{repo_id}' didn't exist, creating repo")
    except HfHubHTTPError as e:
        print(f"Model repo '{repo_id}' exists, skipping create repo")

    print(f"Model repo '{repo_id}' link: https://huggingface.co/{repo_id}\n")

    return repo_id

def is_diffusers_model(model_path):
    """Checks if a given path is a valid Diffusers model directory."""
    required_folders = {"unet", "text_encoder", "text_encoder_2", "tokenizer", "tokenizer_2", "scheduler", "vae"}
    return required_folders.issubset(set(os.listdir(model_path))) and os.path.isfile(os.path.join(model_path, "model_index.json"))

# ---------------------- CONVERSION AND UPLOAD FUNCTIONS ----------------------

def load_sdxl_model(args, is_load_checkpoint, load_dtype, output_widget):
    """Loads the SDXL model from a checkpoint or Diffusers model."""
    model_load_message = "checkpoint" if is_load_checkpoint else "Diffusers" + (" as fp16" if args.fp16 else "")
    with output_widget:
        print(f"Loading {model_load_message}: {args.model_to_load}")

    if is_load_checkpoint:
        loaded_model_data = load_from_sdxl_checkpoint(args, output_widget)
    else:
        loaded_model_data = load_sdxl_from_diffusers(args, load_dtype)

    return loaded_model_data

def load_from_sdxl_checkpoint(args, output_widget):
    """Loads the SDXL model components from a checkpoint file (placeholder)."""
    # text_encoder1, text_encoder2, vae, unet, _, _ = sdxl_model_util.load_models_from_sdxl_checkpoint(
    #    "sdxl_base_v1-0", args.model_to_load, "cpu"
    # )

    # Implement Load model from ckpt or safetensors
    text_encoder1, text_encoder2, vae, unet = None, None, None, None

    with output_widget:
        print("Loading from Checkpoint not implemented, please implement based on your model needs.")

    return text_encoder1, text_encoder2, vae, unet

def load_sdxl_from_diffusers(args, load_dtype):
    """Loads an SDXL model from a Diffusers model directory."""
    pipeline = StableDiffusionXLPipeline.from_pretrained(
        args.model_to_load, torch_dtype=load_dtype, tokenizer=None, tokenizer_2=None, scheduler=None
    )
    text_encoder1 = pipeline.text_encoder
    text_encoder2 = pipeline.text_encoder_2
    vae = pipeline.vae
    unet = pipeline.unet

    return text_encoder1, text_encoder2, vae, unet

def convert_and_save_sdxl_model(args, is_save_checkpoint, loaded_model_data, save_dtype, output_widget):
    """Converts and saves the SDXL model as either a checkpoint or a Diffusers model."""
    text_encoder1, text_encoder2, vae, unet = loaded_model_data
    model_save_message = "checkpoint" + ("" if save_dtype is None else f" in {save_dtype}") if is_save_checkpoint else "Diffusers"

    with output_widget:
        print(f"Converting and saving as {model_save_message}: {args.model_to_save}")

    if is_save_checkpoint:
        save_sdxl_as_checkpoint(args, text_encoder1, text_encoder2, vae, unet, save_dtype, output_widget)
    else:
        save_sdxl_as_diffusers(args, text_encoder1, text_encoder2, vae, unet, save_dtype, output_widget)

def save_sdxl_as_checkpoint(args, text_encoder1, text_encoder2, vae, unet, save_dtype, output_widget):
    """Saves the SDXL model components as a checkpoint file (placeholder)."""
    # logit_scale = None
    # ckpt_info = None

    # key_count = sdxl_model_util.save_stable_diffusion_checkpoint(
    #    args.model_to_save, text_encoder1, text_encoder2, unet, args.epoch, args.global_step, ckpt_info, vae, logit_scale, save_dtype
    # )

    with output_widget:
        print("Saving as Checkpoint not implemented, please implement based on your model needs.")
        # print(f"Model saved. Total converted state_dict keys: {key_count}")

def save_sdxl_as_diffusers(args, text_encoder1, text_encoder2, vae, unet, save_dtype, output_widget):
    """Saves the SDXL model as a Diffusers model."""
    with output_widget:
        reference_model_message = args.reference_model if args.reference_model is not None else 'default model'
        print(f"Copying scheduler/tokenizer config from: {reference_model_message}")

    # Save diffusers pipeline
    pipeline = StableDiffusionXLPipeline(
        vae=vae,
        text_encoder=text_encoder1,
        text_encoder_2=text_encoder2,
        unet=unet,
        scheduler=None,  # Replace None if there is a scheduler
        tokenizer=None,  # Replace None if there is a tokenizer
        tokenizer_2=None  # Replace None if there is a tokenizer_2
    )

    pipeline.save_pretrained(args.model_to_save)

    with output_widget:
        print(f"Model saved as {save_dtype}.")

def convert_model(model_to_load, save_precision_as, epoch, global_step, reference_model, fp16, output_widget):
    """Main conversion function."""
    class Args:  # Defining Args locally within convert_model
        def __init__(self, model_to_load, save_precision_as, epoch, global_step, reference_model, output_path, fp16):
            self.model_to_load = model_to_load
            self.save_precision_as = save_precision_as
            self.epoch = epoch
            self.global_step = global_step
            self.reference_model = reference_model
            self.output_path = output_path #Using output_path even if hardcoded
            self.fp16 = fp16

    # Create a temporary directory for output
    with tempfile.TemporaryDirectory() as tmpdirname:
        args = Args(model_to_load, save_precision_as, epoch, global_step, reference_model, tmpdirname, fp16)
        args.model_to_save = increment_filename(os.path.splitext(args.model_to_load)[0] + ".safetensors")

        try:
            load_dtype = torch.float16 if fp16 else None
            save_dtype = get_save_dtype(save_precision_as)

            is_load_checkpoint = determine_load_checkpoint(model_to_load)
            is_save_checkpoint = not is_load_checkpoint  # reverse of load model

            loaded_model_data = load_sdxl_model(args, is_load_checkpoint, load_dtype, output_widget)
            convert_and_save_sdxl_model(args, is_save_checkpoint, loaded_model_data, save_dtype, output_widget)

            with output_widget:
                return f"Conversion complete. Model saved to {args.model_to_save}"

        except Exception as e:
            with output_widget:
                return f"Conversion failed: {e}"

def upload_to_huggingface(model_path, hf_token, orgs_name, model_name, make_private, output_widget):
    """Uploads a model to the Hugging Face Hub."""
    try:
        login(hf_token, add_to_git_credential=True)
        api = HfApi()
        user = api.whoami(hf_token)
        model_repo = create_model_repo(api, user, orgs_name, model_name, make_private)

        # Determine upload parameters (adjust as needed)
        path_in_repo = ""
        trained_model = os.path.basename(model_path)

        path_in_repo_local = path_in_repo if path_in_repo and not is_diffusers_model(model_path) else ""

        notification = f"Uploading {trained_model} from {model_path} to https://huggingface.co/{model_repo}"
        with output_widget:
            print(notification)

        if os.path.isdir(model_path):
            if is_diffusers_model(model_path):
                commit_message = f"Upload diffusers format: {trained_model}"
                print("Detected diffusers model. Adjusting upload parameters.")
            else:
                commit_message = f"Upload checkpoint: {trained_model}"
                print("Detected regular model. Adjusting upload parameters.")

            api.upload_folder(
                folder_path=model_path,
                path_in_repo=path_in_repo_local,
                repo_id=model_repo,
                commit_message=commit_message,
                ignore_patterns=".ipynb_checkpoints",
            )
        else:
            commit_message = f"Upload file: {trained_model}"
            api.upload_file(
                path_or_fileobj=model_path,
                path_in_repo=path_in_repo_local,
                repo_id=model_repo,
                commit_message=commit_message,
            )
        with output_widget:
            return f"Model upload complete! Check it out at https://huggingface.co/{model_repo}/tree/main"

    except Exception as e:
        with output_widget:
            return f"Upload failed: {e}"

# ---------------------- GRADIO INTERFACE ----------------------

def main(model_to_load, save_precision_as, epoch, global_step, reference_model, fp16, hf_token, orgs_name, model_name, make_private):
  """Main function orchestrating the entire process."""
  output = gr.Markdown()

  # Hardcode output_path
  #output_path = "./converted_model" ##This is incorrect! This will save to current working directory, which isnt ideal.
  # Create tempdir, will only be there for the function
  with tempfile.TemporaryDirectory() as output_path:
    conversion_output = convert_model(model_to_load, save_precision_as, epoch, global_step, reference_model, fp16, output)

    upload_output = upload_to_huggingface(output_path, hf_token, orgs_name, model_name, make_private, output)

    # Return a combined output
    return f"{conversion_output}\n\n{upload_output}"

with gr.Blocks() as demo:

    # Add initial warnings (only once)
    gr.Markdown(f"""
        ## **⚠️ IMPORTANT WARNINGS ⚠️**
        This App is Coded by an LLM partially, and for more information please go here: [Ktiseos Nyx](https://github.com/Ktiseos-Nyx/Sdxl-to-diffusers). The colab edition of this may indeed break AUP. This space is running on CPU and in theory SHOULD work, but may be slow. Earth and Dusk/ Ktiseos Nyx does not have the enterprise budget for ZERO GPU or any gpu sadly! Thank you to the community, John6666 especially for coming to aid when gemini would NOT fix the requirements. Support Ktiseos Nyx & Myself on Ko-fi: [![Ko-fi](https://img.shields.io/badge/Support%20me%20on%20Ko--fi-F16061?logo=ko-fi&logoColor=white&style=flat)](https://ko-fi.com/Z8Z8L4EO)
    """)

    model_to_load = gr.Textbox(label="Model to Load (Checkpoint or Diffusers)", placeholder="Path to model")
    with gr.Row():
        save_precision_as = gr.Dropdown(
            choices=["fp16", "bf16", "float"], value="fp16", label="Save Precision As"
        )
        fp16 = gr.Checkbox(label="Load as fp16 (Diffusers only)")
    with gr.Row():
        epoch = gr.Number(value=0, label="Epoch to Write (Checkpoint)")
        global_step = gr.Number(value=0, label="Global Step to Write (Checkpoint)")

    reference_model = gr.Textbox(label="Reference Diffusers Model",
                                 placeholder="e.g., stabilityai/stable-diffusion-xl-base-1.0")
    #output_path = gr.Textbox(label="Output Path", value="./converted_model") #Remove text box - using temp file approach

    gr.Markdown("## Hugging Face Hub Configuration")
    hf_token = gr.Textbox(label="Hugging Face Token", placeholder="Your Hugging Face write token")
    with gr.Row():
        orgs_name = gr.Textbox(label="Organization Name (Optional)", placeholder="Your organization name")
        model_name = gr.Textbox(label="Model Name", placeholder="The name of your model on Hugging Face")
    make_private = gr.Checkbox(label="Make Repository Private", value=False)

    convert_button = gr.Button("Convert and Upload")
    output = gr.Markdown()

    convert_button.click(fn=main,
                       inputs=[model_to_load, save_precision_as, epoch, global_step, reference_model,
                               fp16, hf_token, orgs_name, model_name, make_private],
                       outputs=output)

demo.launch()