import gradio as gr
import subprocess
import os
from huggingface_hub import whoami, HfApi, login
import random
import time

api = HfApi()
REPO_TYPES = ["model", "dataset", "space"]

def upload_model(source_url, dst_repo, token,civitai_api_key, new_name, dst_repo_path, repo_type):
    try:
        _ = whoami(token)  # Check token validity

        # Check if destination path ends with '/'
        if not dst_repo_path.endswith('/'):
            raise gr.Error("Your destination path must end with a /")

        # Check if source URL is provided
        if not source_url:
            raise gr.Error("You haven't provided a source URL for the model.")

        # Check if destination repository is provided
        if not dst_repo:
            raise gr.Error("You haven't provided a destination repository.")
            
        if not civitai_api_key:
           raise gr.Error("You haven't provided a Civitai Api Key.")
            
        # Login to Hugging Face Hub
        login(token=token)

        # Create a directory to store the downloaded model
        download_dir = "/home/user/apps/downloads/" + str(int(time.time())) + str(random.getrandbits(8)) + "/"
        subprocess.check_call(["mkdir", "-p", download_dir])
        # Download the model using wget
        subprocess.check_call(["wget", f"{source_url}&token={civitai_api_key}", "-P", download_dir])

        # List files in the download directory
        files = os.listdir(download_dir)

        # Set the destination path for the model
        if new_name:
            dst_repo_path = dst_repo_path.strip("/") + "/" + new_name.strip("/")
        else:
            dst_repo_path = dst_repo_path.strip("/") + "/" + files[0]

        # Upload the model file to the destination repository
        api.upload_file(
            path_or_fileobj=download_dir + files[0],
            path_in_repo=dst_repo_path,
            repo_id=dst_repo,
            repo_type=repo_type
        )

        # Clean up: remove downloaded file and directory
        os.remove(download_dir + files[0])
        os.rmdir(download_dir)

        # Generate the URL of the uploaded model
        if repo_type == "space":
            repo_url = f"https://huggingface.co/spaces/{dst_repo}"
        elif repo_type == "dataset":
            repo_url = f"https://huggingface.co/datasets/{dst_repo}"
        else:  # Assuming repo_type is "model"
            repo_url = f"https://huggingface.co/{dst_repo}"

        return (
            f'Your model has been successfully uploaded to <a href="{repo_url}" target="_blank" style="text-decoration:underline">your Hugging Face repository</a>.',
            "success.jpg",
        )
    except Exception as e:
        return str(e), "error.jpg"

# Interface setup
interface = gr.Interface(
    fn=upload_model,
    inputs=[
        gr.Textbox(label="Source URL", placeholder="Source URL Civitai (e.g. https://civitai.com/api/download/models/486156?type=Model&format=SafeTensor)"),
        gr.Textbox(label="Huggingface repository", placeholder="Destination repository HF (e.g. username/repo-name)"),
        gr.Textbox(label="Huggingface Token", placeholder="Write access token HF", type="password"),
        gr.Textbox(label="Civitai Api Key", placeholder="Civitai Api Key", type="password"),
        gr.Textbox(label="New name", placeholder="New name for the model file (optional)"),
        gr.Textbox(label="Huggingface repository Path", placeholder="Destination path within your repository (e.g. /models/Stable-diffusion/)"),
        gr.Dropdown(label="Repo Types", choices=REPO_TYPES, value="model"),
    ],
    outputs=[
        gr.Markdown(label="Output"),
        gr.Image(show_label=False),
    ],
    theme=gr.themes.Base(),
    title="Upload a Civitai Model to Hugging Face Repository",
    description="Upload a model file to your Hugging Face repository. Provide the source URL of the model, your repository information, and your write access token.",
    article="<p>Find your write token at <a href='https://huggingface.co/settings/tokens' target='_blank'>token settings</a></p>",
    allow_flagging="never",
    live=False,  # Prevents automatic re-runs on input change
)



interface.launch()