import os
import json
import gradio as gr
import json
from langchain_legal import run_pipeline

global data
data = {}


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


def get_ai_final_summary(text_input):
    global data
    data = run_pipeline(text_input)
    return data['ai_final_summary']


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


def get_inital_assesment() -> str:
    initial = data.get('ai_first_5', None)
    print(data)
    if initial is None:
        return 'No llm answer yet'
    return initial


def get_applicable_eu_regs() -> str:
    eval = data.get('ai_eur_lex_eval', None)
    if eval is None:
        return 'No llm answer yet'
    sources = data.get('eur_lex_lookup')
    res = f'EVALUATION:\n{eval}'
    if sources is not None:
        res = f'SOURCES:\n{sources}\n\n\n' + res
    return res


def get_applicable_national_regs() -> str:
    eval = data.get('ai_aus_eval', None)
    nat = 'aus'
    if eval is None:
        eval = data.get('ai_ger_eval', None)
        nat = 'ger'
        if eval is None:
            return 'No llm answer yet'
    if nat == 'aus':
        sources = data.get('aus_lookup')
    else:
        sources = data.get('ger_lookup')
    res = f'EVALUATION:\n{eval}'
    if sources is not None:
        res = f'SOURCES:\n{sources}\n\n\n' + res
    return res


def get_ai_challange() -> str:
    challange = data.get('ai_challange', None)
    print(data)
    if challange is None:
        return 'No llm answer yet'
    return challange


CALLBACKS = [
    get_inital_assesment,
    get_applicable_eu_regs,
    get_applicable_national_regs,
    get_ai_challange,
]


def main():
    with gr.Blocks() as demo:
        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=text_output
        )

        with gr.Tabs():
            for name, cb in zip(STEP_NAMES, CALLBACKS):
                with gr.Tab(f'{name}') as t:
                    tab_text = gr.Textbox(placeholder='No answer yet')
                    t.select(cb, outputs=tab_text)

    demo.launch()


if __name__ == '__main__':
    main()