import gradio as gr
import diff
import os
from PIL import Image
import uuid
import random
import bg
from rembg import remove as rm
import datetime 
from huggingface_hub import (HfApi)

proc1=diff
proc2=diff 

save_data='https://huggingface.co/datasets/Omnibus/game-data-1/raw/main/'
token_self = os.environ['HF_TOKEN']
api=HfApi(token=token_self)
#bg1=gr.Interface.load("spaces/Omnibus/fast-bg-remove-mod")
#proc2=gr.Interface.load("spaces/Omnibus/fast-diff")

sky = "https://huggingface.co/spaces/Omnibus/game-test/resolve/main/assets/sky.png"
platform = 'https://huggingface.co/spaces/Omnibus/game-test/resolve/main/assets/platform.png'
star = 'https://huggingface.co/spaces/Omnibus/game-test/resolve/main/assets/star.png'
bomb = 'https://huggingface.co/spaces/Omnibus/game-test/resolve/main/assets/bomb.png'
dude = 'https://huggingface.co/spaces/Omnibus/game-test/resolve/main/assets/dude.png'



def game_fn(sky=sky,star=star,bomb=bomb,platform=platform,dude=dude):
    html_mod=f"""
        <div id="my_game" style="height:600px">
            <iframe
            id="myIframe"
            src="https://omnibus-game-test-static.static.hf.space/index.html?sky={sky}&platform={platform}&star={star}&bomb={bomb}&dude={dude}"        
        	frameborder="0"
        	width="100%"
            height="100%"
            ></iframe>
        </div>"""
    return html_mod

def update_game(inp):
    
    return game_fn(sky=sky)

def sky_fn(inp):
    rand = random.randint(1,200)
    for i in range(rand):
        inp+=" "
    uid=uuid.uuid4()
    output=proc1.send_it(inp,5,1)
    print(output)
    outp=Image.open(output[0])
    width, height = outp.size
    rat = width/height
    if width > height:
        outp = outp.resize((600*rat,600))
    elif width < height:
        outp = outp.resize((800,800*rat))     
    else:
        outp = outp.resize((800,536))
    outp.save(f"{uid}_sky.png")
    out = os.path.abspath(f"{uid}_sky.png")
    #out = os.path.abspath(outp)
    out_url = f'https://omnibus-game-test.hf.space/file={out}'    
    return output[0],out_url

def star_fn(inp):
    rand = random.randint(1,200)
    for i in range(rand):
        inp+=" "
    uid=uuid.uuid4()
    output=proc2.send_it(inp,5,1)
    outp=Image.open(output[0])
    out=rm(outp)
    outrs = out.resize((36,36))
    outrs.save(f"{uid}_star.png")
    out_file = os.path.abspath(f"{uid}_star.png")
    out_url = f'https://omnibus-game-test.hf.space/file={out_file}'    
    return out,out_url

def enemy_fn(inp):
    rand = random.randint(1,200)
    for i in range(rand):
        inp+=" "
    uid=uuid.uuid4()
    output=proc2.send_it(inp,5,1)
    outp=Image.open(output[0])
    out=rm(outp)
    outrs = out.resize((24,24))
    outrs.save(f"{uid}_enemy.png")
    out_file = os.path.abspath(f"{uid}_enemy.png")
    out_url = f'https://omnibus-game-test.hf.space/file={out_file}'    
    return out,out_url

def save_game(background_img,star_img,enemy_img,platform_img,dude_img):
    #background_image.save(f"{uid}-background_img.png")
    #star_img.save(f"{uid}-star_img.png")
    #enemy_img.save(f"{uid}-enemy_img.png")
    #platform_img.save(f"{uid}-platform_img.png")
    #dude_img.save(f"{uid}-dude_img.png")
    user_repo=save_data.split('datasets/',1)[1].split('/raw',1)[0]
    try:
        r = requests.get(f'{save_data}/game_data.json') 
        lod = json.loads(r.text)
    except:
        lod=[]
        pass
    timestamp=str(datetime.datetime.now())
    print(timestamp)
    try: 
        api.upload_file(
        path_or_fileobj=background_img,
        path_in_repo=f"{save_data.split('main/',1)[1]}/images/background/{uid}-background_img.png",
        repo_id=save_data.split('datasets/',1)[1].split('/raw',1)[0],
        token=token_self,
        repo_type="dataset",
        )
        #os.remove("tmp1.json")

    except Exception as e:
        print (e)
        pass    
    block = {'index': len(lod) + 1,
            'timestamp': timestamp,
            'background_url': f'https://huggingface.co/datasets/{user_repo}/raw/main/images/background/{uid}-background_img.png',
            'star_url': f'https://huggingface.co/datasets/{user_repo}/raw/main/images/star/{star_img}',
            'enemy_url': f'https://huggingface.co/datasets/{user_repo}/raw/main/images/enemy/{enemy_img}',
            #'platform_url': f'https://huggingface.co/datasets/{repo}/{space}/raw/main/{platform_img}',
            #'dude_url': f'https://huggingface.co/datasets/{repo}/{space}/raw/main/{dude_img}',
            }
    print(block)
    lod.append(block)    
    
    json_object = json.dumps(lod, indent=4)
    with open("tmp1.json", "w") as outfile:
        outfile.write(json_object)
    try: 
        api.upload_file(
        path_or_fileobj="tmp1.json",
        path_in_repo=f"{save_data.split('main/',1)[1]}/game_data.json",
        repo_id=save_data.split('datasets/',1)[1].split('/raw',1)[0],
        token=token_self,
        repo_type="dataset",
        )
        os.remove("tmp1.json")
        print("success")
    except Exception as e:
        print (e)
        pass


with gr.Blocks() as app:
    with gr.Row():
        with gr.Column():
            prompt_sky=gr.Textbox(label="Background",value="beautiful landscape, real, 8k")
            btn_sky=gr.Button("Make")
            out_im_sky=gr.Image(type='filepath')
            out_sky_url=gr.Textbox(visible=False)
        with gr.Column():
            prompt_star=gr.Textbox(label="Star",value="Colorful Star, blank background")
            btn_star=gr.Button("Make")
            out_im_star=gr.Image(type='filepath')
            out_star_url=gr.Textbox(visible=False)
        with gr.Column():
            prompt_enemy=gr.Textbox(label="Enemy",value="Ball on fire, blank background")
            btn_enemy=gr.Button("Make")
            out_im_enemy=gr.Image(type='filepath')
            out_enemy_url=gr.Textbox(visible=False)
        gr.Column()
        gr.Column()
        
    with gr.Row():
        update_game=gr.Button("Make Game")
        #start_prompt=gr.Textbox(value="beautiful landscape, real, 8k",visible=False)
    html_game = gr.HTML()
    save_btn=gr.Button("Save (Experimental)")

    save_btn.click(save_game,[out_im_sky,out_im_star,out_im_enemy,out_im_enemy,out_im_enemy],None)
    update_game.click(game_fn,[out_sky_url,out_star_url,out_enemy_url],html_game)
    btn_sky.click(sky_fn,prompt_sky,[out_im_sky,out_sky_url])
    btn_star.click(star_fn,prompt_star,[out_im_star,out_star_url])
    btn_enemy.click(enemy_fn,prompt_enemy,[out_im_enemy,out_enemy_url])

    #app.load(sky_fn,prompt_sky,[out_im_sky,out_sky_url]).then(game_fn,[out_sky_url],html_game)
app.queue(concurrency_count=10).launch(show_api=False)