import arrow
import gradio as gr
import os
import re
import pandas as pd
from pathlib import Path
from time import sleep
from tqdm import tqdm

from api_calls import *

ROOT_DIR = Path(__file__).resolve().parents[0]

    
def export_to_txt(output):
    today_dt_str = arrow.now(tz="Asia/Taipei").format("YYYYMMDDTHHmmss")
    with open(f"esg_report_summary-{today_dt_str}.txt", "w") as f:
        f.write(output)
    return f"esg_report_summary-{today_dt_str}.txt"

def print_like_dislike(x: gr.LikeData):
    print(x.index, x.value, x.liked)

def add_text(history, text):
    history = history + [(text, None)]
    return history, gr.Textbox(value="", interactive=False)

def esgsumm_exe(openai_model_name, year, company_name, tone):
    query = "根據您提供的相關資訊和偏好語氣，以繁體中文生成一份符合GRI標準的報告草稿。報告將包括每個GRI披露項目的標題、相關公司行為的概要，以及公司的具體措施和效果。"
    response = api_rag_summ_chain_demo(openai_model_name, query, year, company_name, tone)
    full_anwser = ""
    for chunk in response.iter_content(chunk_size=32):
        if chunk:
            try:
                _c = chunk.decode('utf-8')
            except UnicodeDecodeError:
                _c = " "
            full_anwser += _c
            yield full_anwser
    # for character in response:
    #     full_text += character
    #     yield full_text

def esgqabot(history, openai_model_name, year, company_name):
    query = history[-1][0]
    response = api_rag_qa_chain_demo(openai_model_name, query, year, company_name)
    history[-1][1] = ""
    for chunk in response.iter_content(chunk_size=32):
        if chunk:
            try:
                _c = chunk.decode('utf-8')
            except UnicodeDecodeError:
                _c = " "
            history[-1][1] += _c
            yield history
    # for character in response:
    #     history[-1][1] += character
    #     yield history


css = """
#center {text-align: center}
footer {visibility: hidden}
a {color: rgb(255, 206, 10) !important}
"""
with gr.Blocks(css=css, theme=gr.themes.Monochrome(neutral_hue="lime")) as demo:

    gr.HTML("<h1>ESG RAG Playground</h1>", elem_id="center")
    gr.Markdown("Made by `Abao`", elem_id="center")
    gr.Markdown("---")

    # esgsumm
    with gr.Tab("ESG Report Summarization"):
        gr.HTML("<h2>Report Summarization</h2><p>Summarize report with tone & schema.</p>", elem_id="center")
        with gr.Row():
            with gr.Group():
                gr.Markdown("### Configuration", elem_id="center")
                esgsumm_report_tone = gr.Dropdown(
                    label="Tone", 
                    choices=["富有創意", "中庸", "精確"])
                esgsumm_openai_model_name = gr.Dropdown(
                    label="OpenAI Model", 
                    choices=["gpt-4-turbo-preview", "gpt-3.5-turbo"])
                esgsumm_year = gr.Dropdown(
                    label="Year",
                    choices=["111", "110", "109"]
                )
                esgsumm_company_name = gr.Dropdown(
                    label="Company Name",
                    choices=["台泥", "聯電", "裕融", "大同", "台積電", "鴻海", "中鋼", "中華電信"]
                )
                esgsumm_report_gen_button = gr.Button("Generate Report")

        with gr.Column():
                gr.Markdown("## Generate ESG Summarization", elem_id="center")
                with gr.Accordion("Revise Your Prompt", open=False):
                    esgsumm_checkbox_replace = gr.Checkbox(label="Replace with new prompt")
                    esgsumm_prompt_tmpl = gr.Textbox(
                        label="希望用於本次問答的prompt",
                        info="必須使用到的變數：{filtered_data}、{query}",
                        value=prompt_dict["qa"],
                        interactive=True,
                    )
                esgsumm_report_output = gr.Textbox(
                    label="Report Output",
                    interactive=False,
                    scale=4,
                )
                esgsumm_download_btn = gr.Button("Export Summary")
                esgsumm_download_file = gr.File(
                    label="Download Summary Text", file_types=[".txt"]
                )


    # esgqa
    with gr.Tab("ESG QA"):
        gr.HTML("<h2>ParallelQA (GPT-4 like)</h2><p>Test multiple LLMs at once.</p>", elem_id="center")
        with gr.Row():
            with gr.Group():
                gr.Markdown("### Configuration", elem_id="center")
                esgqa_openai_model_name = gr.Dropdown(
                    label="OpenAI Model", 
                    choices=["gpt-4-turbo-preview", "gpt-3.5-turbo"])
                esgqa_year = gr.Dropdown(
                    label="Year",
                    choices=["111", "110", "109"]
                )
                esgqa_company_name = gr.Dropdown(
                    label="Company Name",
                    choices=["台泥", "聯電", "裕融", "大同", "台積電", "鴻海", "中鋼", "中華電信"]
                )

            with gr.Column():
                gr.Markdown("## Chat with ESGQABot", elem_id="center")
                with gr.Accordion("Revise Your Prompt", open=False):
                    esgqa_checkbox_replace = gr.Checkbox(label="Replace with new prompt")
                    esgqa_prompt_tmpl = gr.Textbox(
                        label="希望用於本次問答的prompt",
                        info="必須使用到的變數：{filtered_data}、{query}",
                        value=prompt_dict["qa"],
                        interactive=True,
                    )
                esgqa_chatbot = gr.Chatbot(
                    [(None, "我是 ESGQABot\n有什麼能為您服務的嗎？")],
                    elem_id="chatbot",
                    scale=1,
                    height=700,
                    bubble_full_width=False
                )
                with gr.Row():
                    esgqa_chatbot_input = gr.Textbox(
                        scale=4,
                        show_label=False,
                        placeholder="Enter text and press enter, or upload an image",
                        container=False,
                    )
                    esgqa_chat_btn = gr.Button("💬")


    # esgsumm
    esgsumm_report_gen_button.click(
        esgsumm_exe, [esgsumm_openai_model_name, esgsumm_year, esgsumm_company_name, esgsumm_report_tone], esgsumm_report_output
    )
    esgsumm_download_btn.click(
        fn=export_to_txt,
        inputs=[esgsumm_report_output],
        outputs=esgsumm_download_file,
    )
    
    # esgqa
    esgqa_chatbot_input.submit(
        add_text, [esgqa_chatbot, esgqa_chatbot_input], [esgqa_chatbot, esgqa_chatbot_input], queue=False
    ).then(
        esgqabot, [esgqa_chatbot, esgqa_openai_model_name, esgqa_year, esgqa_company_name], esgqa_chatbot, api_name="esgqa_response"
    ).then(
        lambda: gr.Textbox(interactive=True), None, [esgqa_chatbot_input], queue=False
    )
    esgqa_chat_btn.click(
        add_text, [esgqa_chatbot, esgqa_chatbot_input], [esgqa_chatbot, esgqa_chatbot_input], queue=False
    ).then(
        esgqabot, [esgqa_chatbot, esgqa_openai_model_name, esgqa_year, esgqa_company_name], esgqa_chatbot, api_name="esgqa_response"
    ).then(
        lambda: gr.Textbox(interactive=True), None, [esgqa_chatbot_input], queue=False
    )
    esgqa_chatbot.like(print_like_dislike, None, None)


if __name__ == "__main__":
    demo.queue().launch(max_threads=10)