import mesop as me
from rag_app.chat_utils import State, _make_style_chat_bubble_wrapper, _ROLE_ASSISTANT, on_chat_input, _make_chat_bubble_style, \
    on_click_submit_chat_msg, _STYLE_CHAT_BUBBLE_NAME, handle_pdf_upload
from rag_app.rag_2 import check_if_exists, precompute_index


_COLOR_BACKGROUND = me.theme_var("background")

_STYLE_APP_CONTAINER = me.Style(
    background=_COLOR_BACKGROUND,
    display="flex",
    flex_direction="column",
    height="100%",
    margin=me.Margin.symmetric(vertical=0, horizontal="auto"),
    width="min(1024px, 100%)",
    box_shadow="0 3px 1px -2px #0003, 0 2px 2px #00000024, 0 1px 5px #0000001f",
    padding=me.Padding(top=20, left=20, right=20),
)


@me.page(security_policy=me.SecurityPolicy(
    allowed_iframe_parents=["https://huggingface.co"]
))
def app():
    if not check_if_exists():
        print("precomputing the index")
        precompute_index()
        
    state = me.state(State)
    with me.box(style=_STYLE_APP_CONTAINER):
        with me.box(style=me.Style(
                width="min(680%, 100%)",
                margin=me.Margin.symmetric(vertical=36, horizontal="auto"),
                flex_grow=1,
                overflow_y="auto",
                padding=me.Padding(left=20, right=20)
        )):
            me.text("""
            FinanceGPT - Powered by open source language models capable of document QnA on Annual
            Investor Reports of top companies.
            """,
                    style=me.Style(font_size=20, margin=me.Margin(bottom=24), text_align="center")
                    )
            me.text("ℹ️ Upload annual reports to start asking questions.",
                    style=me.Style(font_size=12, margin=me.Margin(bottom=24), text_align="center")
                    )
            for index, msg in enumerate(state.output):
                with me.box(style=_make_style_chat_bubble_wrapper(msg.role), key=f"msg-{index}"):
                    if msg.role == _ROLE_ASSISTANT:
                        me.text("assistant", style=_STYLE_CHAT_BUBBLE_NAME)
                    with me.box(style=_make_chat_bubble_style(msg.role)):
                        me.markdown(msg.content)
            if state.in_progress:
                me.progress_spinner()
                with me.box(key="scroll-to", style=me.Style(height=250)):
                    pass
        with me.box(style=me.Style(
                padding=me.Padding(top=30, left=20, right=20),
                display="flex",
                flex_direction="row"
        )):
            with me.content_uploader(
                    accepted_file_types=["application/pdf"],
                    on_upload=handle_pdf_upload,
                    type="icon",
                    style=me.Style(font_weight="bold", margin=me.Margin(right=8)),
            ):
                me.icon("attach_file")
            
            with me.box(style=me.Style(flex_grow=1)):
                me.input(
                    label="Enter your prompt",
                    key=f"input-{len(state.output)}",
                    on_input=on_chat_input,
                    on_enter=on_click_submit_chat_msg,
                    style=me.Style(width="100%")
                )
            with me.content_button(
                    color="primary",
                    type="flat",
                    disabled=state.in_progress,
                    on_click=on_click_submit_chat_msg,
                    style=me.Style(margin=me.Margin(top=8, left=8))
            ):
                me.icon("send" if not state.in_progress else "pending")