import gradio as gr
from response_db import StResponseDb
db = StResponseDb()
a = gr.Number(0)

def get_next_question(history):
    if len(history)==2:
        question = "What is the man doing?"
    elif len(history)==4:
        question = "How many apples are there?"
    else:
        question = "What color is the cat?"
    return question

def ask_a_question(input, taskid, history=[]):
    history.append(input)
    db.add(int(a.value), taskid, len(history)//2-1, history[-2], history[-1])
    history.append(get_next_question(history))
    
    # write some HTML
    html = "<div class='chatbot'>"
    for m, msg in enumerate(history):
        cls = "bot" if m%2 == 0 else "user"
        html += "<div class='msg {}'> {}</div>".format(cls, msg)
    html += "</div>"
    return html, history


css = """
.chatbox {display:flex;flex-direction:column}
.msg {padding:4px;margin-bottom:4px;border-radius:4px;width:80%}
.msg.user {background-color:cornflowerblue;color:white}
.msg.bot {background-color:lightgray;align-self:self-end}
.footer {display:none !important}
"""

def set_images(taskid):
    id1 = f"images/img_{int(10*(taskid-1)+1)}.jpg"
    id2 = f"images/img_{int(10*(taskid-1)+2)}.jpg"
    id3 = f"images/img_{int(10*(taskid-1)+3)}.jpg"
    id4 = f"images/img_{int(10*(taskid-1)+4)}.jpg"
    id5 = f"images/img_{int(10*(taskid-1)+5)}.jpg"
    id6 = f"images/img_{int(10*(taskid-1)+6)}.jpg"
    id7 = f"images/img_{int(10*(taskid-1)+7)}.jpg"
    id8 = f"images/img_{int(10*(taskid-1)+8)}.jpg"
    id9 = f"images/img_{int(10*(taskid-1)+9)}.jpg"
    id10 = f"images/img_{int(10*(taskid-1)+10)}.jpg"
    first_question = "How many dogs are there?"
    first_question_html = f"<div class='chatbot'><div class='msg bot'>{first_question}</div></div>"
    a.value = a.value+1
    return id1, id2, id3, id4, id5, id6, id7, id8, id9, id10, [first_question], first_question_html

with gr.Blocks(css=css) as demo:

    with gr.Column() as img_block:
        with gr.Row():
            img1 = gr.Image()
            img2 = gr.Image()
            img3 = gr.Image()
            img4 = gr.Image()
            img5 = gr.Image()
        with gr.Row():
            img6 = gr.Image()
            img7 = gr.Image()
            img8 = gr.Image()
            img9 = gr.Image()
            img10 = gr.Image()
    conversation = gr.HTML()
    history_log = gr.State([])

    with gr.Column():
        with gr.Row():
            taskid = gr.Number(label="Task ID (Choose from [1,2,3])")
            btn1 = gr.Button("Enter")
        btn1.click(set_images, inputs=taskid, outputs=[img1, img2, img3, img4, img5, img6, img7, img8, img9, img10, history_log, conversation])
    answer = gr.inputs.Textbox(placeholder="Insert answer here.", label="Answer the given question.")
    submit = gr.Button("Submit")
    submit.click(fn=ask_a_question, inputs=[answer, taskid, history_log], outputs=[conversation, history_log])

demo.launch()