import gradio as gr
import interface
import utils
import treegraph as tg
import logger as lg
import os
system_prompt = """
You are a critical AI radiology assistant.
You are helping a radiologist correctly fill out a radiology report.
The report is regarding a Knee MRI.
"""
username = os.environ.get("user")
passwd = os.environ.get("pass")

graph, nodes_dict = tg.build_tree_from_file("knee_template.json")
report_interface = interface.ReportChecklistInterface(
    llm=utils.LLM(model="gpt-3.5-turbo"),
    system_prompt=system_prompt,
    graph=graph,
    nodes_dict=nodes_dict,
)
logger = None

if report_interface.prime_model() is False:
    print("Model priming failed. Please try again.")
    exit()
else:
    print("Model priming successful.")

running = True


def check_report(report, name):
    global logger, report_interface
    if len(name.strip()) < 3:
        return "Please enter a name."
    else:
        logger = lg.Logger(log_file="log.txt")
        report_interface.logger = logger
        report_interface.username = name
    if running:
        results = report_interface.process_input(report)
        if results == "quit":
            quit_fn()
        elif results == "help":
            return report_interface.help_message
        elif results == "exception":
            return "An exception occurred. Please try again."
        else:
            return results
    else:
        return "Model has been stopped."


def quit_fn():
    global running
    running = False
    return "Model has been stopped."


with gr.Blocks(theme="soft") as demo:
    gr.Markdown("## Radiology Report Assistant")
    gr.Markdown(report_interface.help_message)
    name_textbox = gr.Textbox(label="Name")

    report_textbox = gr.TextArea(label="Report", lines=20, max_lines=50)
    with gr.Row():
        check_btn = gr.Button(
            value="Check Report",
        )
        clear_btn = gr.ClearButton(
            value="Clear Messages",
        )
        quit_btn = gr.Button(
            value="Quit",
        )
    results_textbox = gr.TextArea(label="Results", lines=20, max_lines=50)
    clear_btn.add([results_textbox, report_textbox])

    check_btn.click(
        fn=check_report,
        inputs=[report_textbox, name_textbox],
        outputs=[results_textbox],
    )
    quit_btn.click(fn=quit_fn, outputs=[results_textbox])

demo.launch(auth=(username, passwd))