import gradio as gr
import diff
import os
from PIL import Image
import uuid
import random

proc1=gr.Interface.load("spaces/Omnibus/fast-diff")
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="demo" 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+=" "
    #dif = gr.Interface.load("models/stabilityai/stable-diffusion-2-1-base", live=True),
    #dif = models[0]
    uid=uuid.uuid4()
    output=proc1(inp,5,1)
    in_url=f'https://omnibus-fast-diff.hf.space/file={output[0]}'
    #output=dif(inp)
    print(output)
    outp=Image.open(in_url)
    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+=" "
    #dif = gr.Interface.load("models/stabilityai/stable-diffusion-2-1-base", live=True),
    #dif = models[0]
    uid=uuid.uuid4()
    output=proc2(inp,5,1)
    #output=proc2.send_it(inp,5,1)
    #output=dif(inp)
    in_url=f'https://omnibus-fast-diff.hf.space/file={output[0]}'
    
    print(output)
    outp=Image.open(in_url)
    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((24,24))
    outp.save(f"{uid}_star.png")
    out = os.path.abspath(f"{uid}_star.png")
    #out = os.path.abspath(outp)
    out_url = f'https://omnibus-game-test.hf.space/file={out}'    
    return output[0],out_url

with gr.Blocks() as app:
    with gr.Row():
        with gr.Column():
            prompt=gr.Textbox()
            with gr.Row():
                btn=gr.Button("Make Image")
                update_game=gr.Button("Use Image")
        out_im_sky=gr.Image(type='filepath',width=50)
        out_im_star=gr.Image(type='filepath',width=50)
        out_sky_url=gr.Textbox(visible=False)
        out_star_url=gr.Textbox(visible=False)
        start_prompt=gr.Textbox(value="beautiful landscape, real, 8k",visible=False)
    html_game = gr.HTML()
    
    update_game.click(game_fn,[out_sky_url,out_star_url],html_game)
    btn.click(sky_fn,prompt,[out_im_sky,out_sky_url])
    btn.click(star_fn,prompt,[out_im_star,out_star_url])
    app.load(sky_fn,start_prompt,[out_im_sky,out_sky_url]).then(game_fn,[out_sky_url],html_game)
app.launch()