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

api=HfApi()
FILE_TYPES = ["Embedding", "Model", "Lora", "Lycoris"]

def duplicate(source_url, preview_url, model_id, dst_repo, token, dst_repo_path, repo_type):
    try:
        _ = whoami(token)
        # ^ this will throw if token is invalid

        # make sure the user fills out the other required paths.
        if not dst_repo_path[len(dst_repo_path)-1] == '/':
                raise Exception("Your destination path *must* end with a /")
        if not source_url:
                raise Exception("You haven't chosen a file to download!")
        if not dst_repo:
                raise Exception("You haven't chosen a repo to download to")
        login(token=token)

        # keep things separate, partly in case people download different files with same name (`download.zip`). Especially, it also allows saving filename to work
        dir="/home/user/apps/downloads/"+str(int(time.time()))+str(random.getrandbits(8))+"/"
        subprocess.check_call([r"mkdir","-p",dir])
        subprocess.check_call([r"aria2c","-x16","--split=16",source_url,"--dir="+dir])

        filename=os.listdir(dir)[0].rsplit('.',1)[0]
        subprocess.check_call([r"curl","https://civitai.com/api/v1/model-versions/"+model_id,source_url,"-O",dir+filename+".civitai.info"])
        subprocess.check_call([r"curl",preview_url,"-O",filename+".preview.png"])

	match repo_type:
		case "Embedding":
			repopath="/embeddings"
		case "Model":
			repopath="/models/Stable-diffusion"
		case "Lora":
			repopath="/models/Lora/Lora"
		case "Lycoris":
			repopath="/models/Lora/Lycoris"

	if nsfw==True:
		repopath=repopath+"/nsfw"

        api.upload_folder(
                path_or_fileobj=dir,
                path_in_repo=repopath,
                repo_id=dst_repo,
                repo_type=repo_type
        )

        # now clean up
        files=os.listdir(dir)
        os.remove(dir+files[0])
        os.remove(dir+files[1])
        os.remove(dir+files[2])
        os.rmdir(dir)

        return (
            f'Find your commit at the top of <a href=\'https://hf.co/spaces/{dst_repo}/commits/main{repopath}\' target="_blank" style="text-decoration:underline">your commits</a>',
            "sp.jpg",
        )

    except Exception as e:
        blames=["grandma","my boss","your boss","God","you","you. It's *all* your fault.","the pope"]
        blameweights=(1,1,1,1,4,2,1)
        excuses=["I blame it all on "+random.choices(blames,weights=blameweights)[0],"It's my fault, sorry.","I did it on purpose.","That file doesn't want to be downloaded.","You nincompoop!"]
        excusesweights=(12,1,1,2,3)
        excuse=random.choices(excuses,weights=excusesweights)[0]
        return (
            f"""
        ### Error 😢😢😢
        
        {e}
        
        
        
        <i>""" + excuse+"</i>",
            None,
        )


interface = gr.Interface(
    fn=duplicate,
    inputs=[
        gr.Textbox(placeholder="Source URL for model (e.g. civitai.com/api/download/models/4324322534)"),
        gr.Textbox(placeholder="Preview URL for model (e.g. https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/df8e24d7-4094-42cf-0d8f-a50afd28f800/width=450,optimized=true/00143-3805918203.jpeg)"),
        gr.Textbox(placeholder="model ID (e.g. 4324322534)"),
        gr.Textbox(placeholder="Destination repository (e.g. osanseviero/dst)",value="spaces/mirroring/civitai_mirror/"),
        gr.Textbox(placeholder="Write access token", type="password"),
        gr.Textbox(placeholder="Destination for your file within your repo. Don't include the filename, end path with a / (e.g. /models/Stable-diffusion/)"),
        gr.Dropdown(choices=FILE_TYPES, value="Model"),
    ],
    outputs=[
        gr.Markdown(label="output"),
        gr.Image(show_label=False),
    ],
    title="Download a civitai model to your repo!",
    description="Download a model from Civitai to your repo on HF. Hopefully will soon parse the modelid from the download url, then get the preview automatically from the api response (lets be realistic, this will never happen); in the meantime you need to provide the model id number from the download url , and the location of the preview file. This Space is a an experimental demo.",
    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, # since i keep wondering, this prevents it from running again automatically when an input changes
)
interface.launch(enable_queue=True)