import gradio as gr
import base64
import requests
import os
import json

MAINTENANCE_NOTICE1 = 'Hint 1: If the app report "Something went wrong, connection error out", please turn off your proxy and retry.\nHint 2: If you upload a large size of image like 10MB, it may take some time to upload and process. Please be patient and wait.'
default_chatbox = []
NOTES = 'This app is adapted from <a href="https://modelscope.cn/models/Fengshenbang/Ziya-Visual-Lyrics-14B">https://modelscope.cn/models/Fengshenbang/Ziya-Visual-Lyrics-14B</a>. It would be recommended to check out the repo if you want to see the detail of our model. And most of the codes attach to this demo are modify from <a href="https://arxiv.org/abs/2312.05278">Lyrics</a>.'

def post(
        input_text,
        image_path,
        temperature,
        top_p,
        result_previous,
        hidden_image
        ):
    # 服务的URL
    with open(image_path,'rb') as f:
        image_prompt = base64.b64encode(f.read())

    service_url = os.getenv('URL_PATH')
    # 准备请求数据
    data = {
        'image_strs': image_prompt,
        'prompt': input_text
    }
        # 发送请求
    print("In the fn")
    response = requests.post(service_url, data=data)
    print(response.json())
    result_text = []
    result_text.append((input_text, response.json()))
    print(result_text)
    return "", result_text, None



def clear_fn(value):
    return "", default_chatbox, None

def clear_fn2(value):
    return default_chatbox

def io_fn(a, b, c):
    print(f"call io_fn")
    return a, b


gr.close_all()
examples = []
with open("./examples/example_inputs.jsonl") as f:
    for line in f:
        data = json.loads(line)
        examples.append(data)

def main():
    with gr.Blocks(css='style.css') as demo:

        with gr.Row():
            with gr.Column(scale=4.5):
                with gr.Group():
                    input_text = gr.Textbox(label='Input Text', placeholder='Please enter text prompt below and press ENTER.')
                    with gr.Row():
                        run_button = gr.Button('Generate')
                        clear_button = gr.Button('Clear')

                    img_path = gr.Image(type="filepath", label="Image Prompt", value=None)
                    
                with gr.Row():
                    temperature = gr.Slider(maximum=1, value=0.7, minimum=0, label='Temperature')
                    top_p = gr.Slider(maximum=1, value=0.1, minimum=0, label='Top P')
                with gr.Group():
                    with gr.Row():
                        with gr.Column(scale=7):    
                            maintenance_notice = gr.Markdown(MAINTENANCE_NOTICE1)
                        with gr.Column(scale=2):    
                            change_button = gr.Button('Change hint to English', visible=False)
            with gr.Column(scale=5.5):
                result_text = gr.components.Chatbot(label='Multi-round conversation History', value=[]).style(height=550)
                hidden_image_hash = gr.Textbox(visible=False)

        gr_examples = gr.Examples(examples=[[example["text"], example["image"]] for example in examples], 
                                    inputs=[input_text, img_path],
                                    label="Example Inputs (Click to insert an examplet into the input box)",
                                    examples_per_page=3)
        
        print(gr.__version__)
        run_button.click(fn=post,inputs=[input_text, img_path, temperature, top_p, result_text, hidden_image_hash],
                            outputs=[input_text, result_text, hidden_image_hash])
        input_text.submit(fn=post,inputs=[input_text, img_path, temperature, top_p, result_text, hidden_image_hash],
                            outputs=[input_text, result_text, hidden_image_hash])
        clear_button.click(fn=clear_fn, inputs=clear_button, outputs=[input_text, result_text, img_path])
        img_path.upload(fn=clear_fn2, inputs=clear_button, outputs=[result_text])
        img_path.clear(fn=clear_fn2, inputs=clear_button, outputs=[result_text])
        
    demo.queue(concurrency_count=10)
    demo.launch(server_name="0.0.0.0")

if __name__ == '__main__':
    print('start service')
    main()