import os
import shutil
import gradio as gr

from uuid import uuid4
from requests import Session
import urllib.parse


API_URL = os.environ.get('API_URL')
API_TOKEN = os.environ.get('API_TOKEN')


class Client:
    def __init__(self):
        self.c = Session()
        self.c.headers.update({
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {API_TOKEN}'
        })

    def send_message(self, input):
        res = self.c.post(
            API_URL,
            json={'message': {'text': input}}
        )
        return res.json()

    def send_images(self, text, images):
        res = self.c.post(
            API_URL,
            json={'message': {'file': images, 'text': text}}
        )
        return res.json()


client = Client()


def send_message(input, files):
    print('sending message', files)
    res = client.send_images(input, files)
    print(res)
    return res['message']['data']

def upload_files(files):
    if not os.path.exists('uploads'):
        os.makedirs('uploads')

    print('uploading files', files)
    file_urls = []
    for file in files:
        file_urls.append(
            f'https://philipp-zettl-easybits-fal-lora.hf.space/gradio_api/file={urllib.parse.quote_plus(file)}'
        )
    print(file_urls)
    return file_urls


with gr.Blocks() as demo:
    gr.Markdown("## Demo")
    with gr.Row(equal_height=True):
        with gr.Column():
            gr.Markdown("### Input")
            inpt = gr.Textbox(label="LoRA Key-Word")
            upload_button = gr.UploadButton("Upload Images", file_count='multiple', file_types=['image'])
            hidden_files = gr.State([])
            upload_button.upload(upload_files, upload_button, hidden_files)
        with gr.Column():
            gr.Markdown("### Output")
            output = gr.Textbox(lines=5, label="Output Text")
    with gr.Row():
        send_button = gr.Button("Send")
        send_button.click(send_message, inputs=[inpt, hidden_files], outputs=[output])



demo.launch(share=True, allowed_paths=['uploads/*'])