import gradio as gr
from PIL import Image
import base64
import io
from utils.tools import txt2string,base64_convert
from component.gradio_iframe import iFrame
from chatbots.initialization_bot.chatbot import Chatbot as InitialBot
from chatbots.improvement_bot.chatbot import Chatbot as ImproveBot
from chatbots.repolish_bot.chatbot import Chatbot as PolistBot


def process_inputs(session_state,image):
    session_state['screen_shot']=base64_convert(image)
    return session_state

def generate_code(session_state):
    base64_img=session_state['screen_shot']

    initial_bot=InitialBot()
    
    result=initial_bot.generate_code(base64_img)
    session_state['html_source']=result
    return result,result,session_state

def improve_code(session_state):
    base64_img=session_state['screen_shot']
    html_source=session_state['html_source']

    chatbot=ImproveBot()
    result=chatbot.generate_code(html_source,base64_img)
    session_state['html_source']=result
    return result,result,session_state

default_source=txt2string('./html_source/experiment1.html')

def initial_state():
    state={
        'name':'session state',
        'screen_shot':None,
        'html_source':default_source,
    }
    return state

def repolish(session_state,chat_input):
    html_source=session_state['html_source']
    reference_img=session_state['screen_shot']
    current_img=base64_convert(chat_input['files'][0])
    chatbot=PolistBot()
    result=chatbot.generate_code(chat_input['text'],html_source,[reference_img,current_img])
    return result,result


def render(session_state,source):
    session_state['html_source']=source
    return session_state,source


def interface():
    with gr.Blocks() as demo:
        session_state =gr.State(initial_state)
        with gr.Tab("Generate"):
            with gr.Row():
                with gr.Column():
                    gr.Markdown(
                    """
                    #  [A Large Language Model Converting UI to Codes] 🔍
                    """
                    )
                    gr.Markdown(
                    """
                     ## Automatically converting web screenshots to html codes !
                    """
                    )
            with gr.Row():
                with gr.Column():
                    screen_shot_upload=gr.Image(type="pil")
                    generate_button=gr.Button(value='Generate')
                    improve_button=gr.Button(value='Self-improve')
                    with gr.Group():
                        with gr.Row():

                            chat_input = gr.MultimodalTextbox(interactive=True,
                                                    file_count="multiple",
                                                    placeholder="Enter message or upload file...", show_label=False,submit_btn=False,scale=8)
                            #repolish_system_text=gr.Text()
                            repolish_button=gr.Button(value='Repolish',scale=2) 

        with gr.Tab("Source"):
            source_area=gr.TextArea(value=default_source,interactive=True,label='html source',lines=35)
            render_button=gr.Button(value='Rerender')

        with gr.Tab('Preview'):
            with gr.Row():
                html_area = iFrame(value=default_source,height=800)

        with gr.Tab('Debug'):
                chat_area = gr.Chatbot(
                    elem_id="chatbot",
                    bubble_full_width=False,
                    scale=1,
                )
        
                #chat_input = gr.MultimodalTextbox(interactive=True,
                #                                file_count="multiple",
                #                                placeholder="Enter message or upload file...", show_label='ergergergr')

                #chat_msg = chat_input.submit(add_message, [chatbot, chat_input], [chatbot, chat_input])
                #bot_msg = chat_msg.then(bot, chatbot, chatbot, api_name="bot_response")
        

        screen_shot_upload.change(process_inputs,inputs=[session_state,screen_shot_upload],outputs=[session_state])
        generate_button.click(generate_code,inputs=[session_state],outputs=[source_area,html_area,session_state])
        improve_button.click(improve_code,inputs=[session_state],outputs=[source_area,html_area,session_state])
        render_button.click(render,inputs=[session_state,source_area],outputs=[session_state,html_area])
        repolish_button.click(repolish,inputs=[session_state,chat_input],outputs=[html_area,source_area])
    return demo