import json
import gradio as gr

from langchain_legal import run_pipeline


def load_json(path: str) -> dict:
    with open(path, 'r') as f:
        return json.load(f)


def get_ai_final_summary(text_input):
    data = run_pipeline(text_input)
    return data, data.get('ai_final_summary', 'No answer yet')


STEP_NAMES = [
    'Initial assessment',
    'Applicable EU regs',
    'Applicable national regs',
    'Complementary assessment',
]


def get_final_summary(data):
    summ = data.get('ai_final_summary')
    if summ is None:
        return gr.Textbox('No answer yet', label='Final summary')
    return gr.Textbox(summ, label='Final summary', show_copy_button=True)


def get_inital_assesment(data):
    initial = data.get('ai_first_5', None)
    if initial is None:
        return gr.Textbox('No llm answer yet', show_label=False)
    return gr.Textbox(initial, show_label=False, lines=25)


def get_eu_regs_eval(data):
    eval = data.get('ai_eur_lex_eval', None)
    if eval is None:
        return gr.Textbox('No llm answer yet', show_label=False)
    return gr.Textbox(eval, show_label=False, lines=25)


def get_eu_regs_sources(data):
    lookup = data.get('eur_lex_lookup', None)
    if lookup is None:
        return gr.Json({}, show_label=False)
    return gr.Json(lookup)


def get_national_regs_eval(data):
    eval = data.get('ai_aus_eval', None)
    if eval is None:
        eval = data.get('ai_ger_eval')
        if eval is None:
            return gr.Textbox('No llm answer yet', show_label=False)
    return gr.Textbox(eval, show_label=False, lines=25)


def get_national_regs_sources(data):
    lookup = data.get('aus_lookup', None)
    if lookup is None:
        lookup = data.get('ger_lookup')
        if lookup is None:
            return gr.Json({}, show_label=False)
    return gr.Json(lookup)


def get_ai_challange(data):
    challange = data.get('ai_challange', None)
    if challange is None:
        return gr.Textbox('No llm answer yet', show_label=False)
    return gr.Textbox(challange, show_label=False, lines=50)


def update_ui(data):
    return (
        get_inital_assesment(data),
        get_eu_regs_eval(data),
        get_eu_regs_sources(data),
        get_national_regs_eval(data),
        get_national_regs_sources(data),
        get_ai_challange(data),
    )


def main():
    with gr.Blocks() as demo:
        data = gr.State({})
        gr.Markdown("# Lawgarithm's LLM legal research tool")

        with gr.Column():
            text_input = gr.Textbox(
                label='Case description',
                placeholder='Write your legal case here',
            )
            submit_btn = gr.Button('Submit')
            text_output = gr.Textbox(label='Final Summary')
            submit_btn.click(
                get_ai_final_summary,
                inputs=[text_input],
                outputs=[data, text_output],
                show_progress='full',
                queue=True,
            )

        with gr.Tabs():
            with gr.Tab(STEP_NAMES[0]):
                inital_tab_text = gr.Textbox(
                    placeholder='No answer yet', show_label=False
                )
            with gr.Tab(STEP_NAMES[1]):
                eu_regs_tab_text = gr.Textbox(
                    placeholder='No answer yet', show_label=False
                )
                eu_regs_tab_sources = gr.Json({})
            with gr.Tab(STEP_NAMES[2]):
                nat_regs_tab_text = gr.Textbox(
                    placeholder='No answer yet', show_label=False
                )
                nat_regs_tab_sources = gr.Json({})
            with gr.Tab(STEP_NAMES[3]):
                comp_assesment_tab_text = gr.Textbox(
                    placeholder='No answer yet', show_label=False
                )
        data.change(
            update_ui,
            inputs=data,
            outputs=[
                inital_tab_text,
                eu_regs_tab_text,
                eu_regs_tab_sources,
                nat_regs_tab_text,
                nat_regs_tab_sources,
                comp_assesment_tab_text,
            ],
        )

    demo.launch()


if __name__ == '__main__':
    main()