import gradio as gr
from models import models
from PIL import Image
import requests
import uuid
import io 
import base64
import random
from transforms import RGBTransform # from source code mentioned above


loaded_model=[]
for i,model in enumerate(models):
    try:
        loaded_model.append(gr.load(f'models/{model}'))
    except Exception as e:
        print(e)
        pass
print (loaded_model)



def run_dif_color(out_prompt,model_drop,cnt,color,tint):
    h=color.lstrip('#')
    #h = input('Enter hex: ').lstrip('#')
    #print('RGB =', tuple(int(h[i:i+2], 16) for i in (0, 2, 4)))
    color=tuple(int(h[i:i+2], 16) for i in (0, 2, 4))
    
    print (color)
    p_seed=""
    out_box=[]
    out_html=""
    #for i,ea in enumerate(loaded_model):
    
    
    
    for i in range(int(cnt)):
        rand=random.randint(1,500)
        for i in range(rand):
            p_seed+=" "
        try:
            #model=gr.load(f'models/{model[int(model_drop)]}')        
            model=loaded_model[int(model_drop)]
            out_img=model(out_prompt+p_seed)
            print(out_img)

            raw=Image.open(out_img)
            raw=raw.convert('RGB')

            colorize = RGBTransform().mix_with(color,factor=float(tint)).applied_to(raw)    
            
            out_box.append(colorize)           
        except Exception as e:
            print(e)
            out_html=str(e)
            pass
            
        yield out_box,out_html


def run_dif(out_prompt,model_drop,cnt):
    p_seed=""
    out_box=[]
    out_html=""
    #for i,ea in enumerate(loaded_model):
    for i in range(int(cnt)):
        p_seed+=" "
        try:
            model=loaded_model[int(model_drop)]
            out_img=model(out_prompt+p_seed)
            print(out_img)
            out_box.append(out_img)           
        except Exception as e:
            print(e)
            out_html=str(e)
            pass
        yield out_box,out_html

def run_dif_og(out_prompt,model_drop,cnt):
    out_box=[]
    out_html=""
    #for i,ea in enumerate(loaded_model):
    for i in range(cnt):
        try:
            #print (ea)
            model=loaded_model[int(model_drop)]
            out_img=model(out_prompt)
            print(out_img)
            url=f'https://omnibus-top-20.hf.space/file={out_img}'
            print(url)
            uid = uuid.uuid4()
            #urllib.request.urlretrieve(image, 'tmp.png')
            #out=Image.open('tmp.png')
            r = requests.get(url, stream=True)
            
            if r.status_code == 200:
                img_buffer = io.BytesIO(r.content)
                print (f'bytes:: {io.BytesIO(r.content)}')
                str_equivalent_image = base64.b64encode(img_buffer.getvalue()).decode()
                img_tag = "<img src='data:image/png;base64," + str_equivalent_image + "'/>"                
                out_html+=f"<div  class='img_class'><a href='https://huggingface.co/models/{models[i]}'>{models[i]}</a><br>"+img_tag+"</div>"
                out = Image.open(io.BytesIO(r.content))
                out_box.append(out)
            html_out = "<div class='grid_class'>"+out_html+"</div>"
            yield out_box,html_out
        except Exception as e:
            out_html+=str(e)
            html_out = "<div class='grid_class'>"+out_html+"</div>"
            
            yield out_box,html_out

def thread_dif(out_prompt,mod):
    out_box=[]
    out_html=""
    #for i,ea in enumerate(loaded_model):
    try:
        print (ea)
        model=loaded_model[int(mod)]
        out_img=model(out_prompt)
        print(out_img)
        url=f'https://omnibus-top-20.hf.space/file={out_img}'
        print(url)
        uid = uuid.uuid4()
        #urllib.request.urlretrieve(image, 'tmp.png')
        #out=Image.open('tmp.png')
        r = requests.get(url, stream=True)
        
        if r.status_code == 200:
            img_buffer = io.BytesIO(r.content)
            print (f'bytes:: {io.BytesIO(r.content)}')
            str_equivalent_image = base64.b64encode(img_buffer.getvalue()).decode()
            img_tag = "<img src='data:image/png;base64," + str_equivalent_image + "'/>"                
            
            #out_html+=f"<div  class='img_class'><a href='https://huggingface.co/models/{models[i]}'>{models[i]}</a><br>"+img_tag+"</div>"
            out = Image.open(io.BytesIO(r.content))
            out_box.append(out)
        else:
            out_html=r.status_code
        html_out = "<div class='grid_class'>"+out_html+"</div>"
        return out_box,html_out
    except Exception as e:
        out_html=str(e)
        #out_html+=str(e)
        html_out = "<div class='grid_class'>"+out_html+"</div>"
        
        return out_box,html_out


            
def start_threads(prompt):
    t1 = threading.Thread(target=thread_dif, args=(prompt,0))
    t2 = threading.Thread(target=thread_dif, args=(prompt,1))
    t1.start()
    t2.start()
    print (t1)
    print (t2)
    a1,a2=t1.result()
    b1,b2=t2.result()
    return a1,a2
            
css="""
.grid_class{
display:flex;
height:100%;
}
.img_class{
min-width:200px;
}

"""
with gr.Blocks(css=css) as app:
    with gr.Row():
        with gr.Column():
            inp=gr.Textbox(label="Prompt")
            btn=gr.Button()
        with gr.Column():
            col = gr.ColorPicker(label="Color Tint")
            tint = gr.Slider(label="Tint Strength", minimum=0, maximum=1, step=0.01, value=0.30)
        
    with gr.Row():
        model_drop=gr.Dropdown(label="Models", choices=models, type='index', value=models[0])
        cnt = gr.Number(value=1)
    out_html=gr.HTML()
    outp=gr.Gallery()
    btn.click(run_dif_color,[inp,model_drop,cnt,col,tint],[outp,out_html])
app.launch()