import json
import os

import gradio as gr
from dotenv import load_dotenv
from gradio_modal import Modal

from logger import logger
from model import get_guardian_response, get_prompt
from utils import (
    get_messages,
    get_result_description,
    load_command_line_args,
    to_snake_case,
    to_title_case,
)

load_command_line_args()
load_dotenv()

catalog = {}

toy_json = '{"name": "John"}'

with open("catalog.json") as f:
    logger.debug("Loading catalog from json.")
    catalog = json.load(f)


def update_selected_test_case(button_name, state: gr.State, event: gr.EventData):
    target_sub_catalog_name, target_test_case_name = event.target.elem_id.split("---")
    state["selected_sub_catalog"] = target_sub_catalog_name
    state["selected_criteria_name"] = target_test_case_name
    state["selected_test_case"] = [
        t
        for sub_catalog_name, sub_catalog in catalog.items()
        for t in sub_catalog
        if t["name"] == to_snake_case(button_name) and to_snake_case(sub_catalog_name) == target_sub_catalog_name
    ][0]
    return state


def on_test_case_click(state: gr.State):
    selected_sub_catalog = state["selected_sub_catalog"]
    selected_criteria_name = state["selected_criteria_name"]
    selected_test_case = state["selected_test_case"]

    logger.debug(f'Changing to test case "{selected_criteria_name}" from catalog "{selected_sub_catalog}".')

    is_context_editable = selected_criteria_name == "context_relevance"
    is_user_message_editable = selected_sub_catalog == "harmful_content_in_user_prompt"
    is_assistant_message_editable = (
        selected_sub_catalog == "harmful_content_in_assistant_response"
        or selected_criteria_name == "groundedness"
        or selected_criteria_name == "answer_relevance"
    )
    is_tools_present = "tools" in selected_test_case and selected_test_case["tools"] is not None

    test_case_name = f'<h2>{to_title_case(selected_test_case["name"])}</h2>'

    criteria = selected_test_case["criteria"]

    # update context field:
    if is_context_editable:
        context = gr.update(
            value=selected_test_case["context"], interactive=True, visible=True, elem_classes=["input-box"]
        )
    else:
        context = gr.update(
            visible=selected_test_case["context"] is not None,
            value=selected_test_case["context"],
            interactive=False,
            elem_classes=["read-only", "input-box"],
        )

    tools = gr.update(
        visible=is_tools_present,
        value=selected_test_case["tools"] if is_tools_present else toy_json,
        elem_classes=["read-only", "margin-bottom"],
    )

    # update user message field
    if is_user_message_editable:
        user_message = gr.update(
            value=selected_test_case["user_message"], visible=True, interactive=True, elem_classes=["input-box"]
        )
    else:
        user_message = gr.update(
            value=selected_test_case["user_message"], interactive=False, elem_classes=["read-only", "input-box"]
        )

    # update assistant message field
    if is_tools_present:
        assistant_message_json = gr.update(
            visible=True,
            value=selected_test_case["assistant_message"],
            elem_classes=["read-only", "margin-bottom"],
        )
        assistant_message_text = gr.update(visible=False)
    else:
        if is_assistant_message_editable:
            assistant_message_text = gr.update(
                value=selected_test_case["assistant_message"],
                visible=True,
                interactive=True,
                elem_classes=["input-box"],
            )
        else:
            assistant_message_text = gr.update(
                visible=selected_test_case["assistant_message"] is not None,
                value=selected_test_case["assistant_message"],
                interactive=False,
                elem_classes=["read-only", "input-box"],
            )
        assistant_message_json = gr.update(visible=False)

    result_text = gr.update(visible=False, value="")
    return (
        test_case_name,
        criteria,
        context,
        user_message,
        assistant_message_text,
        assistant_message_json,
        tools,
        result_text,
    )


def change_button_color(event: gr.EventData):
    return [
        (
            gr.update(elem_classes=["catalog-button", "selected"])
            if v.elem_id == event.target.elem_id
            else gr.update(elem_classes=["catalog-button"])
        )
        for c in catalog_buttons.values()
        for v in c.values()
    ]


def on_submit(criteria, context, user_message, assistant_message_text, assistant_message_json, tools, state):
    criteria_name = state["selected_criteria_name"]
    if criteria_name == "function_calling_hallucination":
        assistant_message = assistant_message_json
    else:
        assistant_message = assistant_message_text
    test_case = {
        "name": criteria_name,
        "criteria": criteria,
        "context": context,
        "user_message": user_message,
        "assistant_message": assistant_message,
        "tools": tools,
    }

    messages = get_messages(test_case=test_case, sub_catalog_name=state["selected_sub_catalog"])

    logger.debug(
        f"Starting evaluation for subcatelog {state['selected_sub_catalog']} and criteria name {state['selected_criteria_name']}"
    )

    result_label = get_guardian_response(messages=messages, criteria_name=criteria_name)["assessment"]  # Yes or No

    html_str = f"<p>{get_result_description(state['selected_sub_catalog'], state['selected_criteria_name'])} <strong>{result_label}</strong></p>"
    # html_str = f"{get_result_description(state['selected_sub_catalog'], state['selected_criteria_name'])} {result_label}"
    return gr.update(value=html_str)


def on_show_prompt_click(criteria, context, user_message, assistant_message_text, assistant_message_json, tools, state):
    criteria_name = state["selected_criteria_name"]
    if criteria_name == "function_calling_hallucination":
        assistant_message = assistant_message_json
    else:
        assistant_message = assistant_message_text
    test_case = {
        "name": criteria_name,
        "criteria": criteria,
        "context": context,
        "user_message": user_message,
        "assistant_message": assistant_message,
        "tools": tools,
    }

    messages = get_messages(test_case=test_case, sub_catalog_name=state["selected_sub_catalog"])
    prompt = get_prompt(messages, criteria_name)
    prompt = prompt.replace("<", "&lt;").replace(">", "&gt;").replace("\\n", "<br>")
    return gr.Markdown(prompt)


ibm_blue = gr.themes.Color(
    name="ibm-blue",
    c50="#eff6ff",
    c100="#dbeafe",
    c200="#bfdbfe",
    c300="#93c5fd",
    c400="#60a5fa",
    c500="#0F62FE",
    c600="#2563eb",
    c700="#1d4ed8",
    c800="#1e40af",
    c900="#1e3a8a",
    c950="#1d3660",
)

head_style = """
<style>
    @media (min-width: 1536px)
    {
        .gradio-container {
            max-width: 1124px !important;
        }
    }
    @media (min-width: 1024px)
    {
        .gradio-container {
            max-width: 1124px !important;
        }
    }
</style>
"""

with gr.Blocks(
    title="Granite Guardian",
    theme=gr.themes.Soft(
        primary_hue=ibm_blue,
        font=[gr.themes.GoogleFont("IBM Plex Sans"), gr.themes.GoogleFont("Source Sans 3")],
    ),
    head=head_style,
    fill_width=False,
    css=os.path.join(os.path.dirname(os.path.abspath(__file__)), "./styles.css"),
) as demo:

    state = gr.State(
        value={"selected_sub_catalog": "harmful_content_in_user_prompt", "selected_criteria_name": "general_harm"}
    )

    starting_test_case = [
        t
        for sub_catalog_name, sub_catalog in catalog.items()
        for t in sub_catalog
        if t["name"] == state.value["selected_criteria_name"]
        and sub_catalog_name == state.value["selected_sub_catalog"]
    ][0]

    with gr.Row(elem_classes="header-row"):
        with gr.Column(scale=4):
            gr.HTML("<h2>IBM Granite Guardian 3.1</h2>", elem_classes="title")
            gr.HTML(
                elem_classes="system-description",
                value="<p>Granite Guardian models are specialized language models in the Granite family that can detect harms and risks in generative AI systems. They can be used with any large language model to make interactions with generative AI systems safe. Select an example in the left panel to see how the Granite Guardian model evaluates harms and risks in user prompts, assistant responses, and for hallucinations in retrival-augmented generation and function calling. In this demo, we use granite-guardian-3.1-8b.</p>",
            )
    with gr.Row(elem_classes="column-gap"):
        with gr.Column(scale=0, elem_classes="no-gap"):
            title_display_left = gr.HTML("<h2>Example Risks</h2>", elem_classes=["subtitle", "subtitle-harms"])
            accordions = []
            catalog_buttons: dict[str, dict[str, gr.Button]] = {}
            for i, (sub_catalog_name, sub_catalog) in enumerate(catalog.items()):
                with gr.Accordion(
                    to_title_case(sub_catalog_name), open=(i == 0), elem_classes="accordion"
                ) as accordion:
                    for test_case in sub_catalog:
                        elem_classes = ["catalog-button"]
                        elem_id = f"{sub_catalog_name}---{test_case['name']}"
                        if starting_test_case == test_case:
                            elem_classes.append("selected")

                        if sub_catalog_name not in catalog_buttons:
                            catalog_buttons[sub_catalog_name] = {}

                        catalog_buttons[sub_catalog_name][test_case["name"]] = gr.Button(
                            to_title_case(test_case["name"]),
                            elem_classes=elem_classes,
                            variant="secondary",
                            size="sm",
                            elem_id=elem_id,
                        )

                    accordions.append(accordion)

        with gr.Column(visible=True, scale=1) as test_case_content:
            with gr.Row(elem_classes="no-stretch"):
                test_case_name = gr.HTML(
                    f'<h2>{to_title_case(starting_test_case["name"])}</h2>', elem_classes="subtitle"
                )
                show_propt_button = gr.Button("Show prompt", size="sm", scale=0, min_width=110)

            criteria = gr.Textbox(
                label="Evaluation Criteria",
                lines=3,
                interactive=False,
                value=starting_test_case["criteria"],
                elem_classes=["read-only", "input-box", "margin-bottom"],
            )
            gr.HTML(elem_classes=["block", "content-gap"])
            context = gr.Textbox(
                label="Context",
                lines=3,
                interactive=True,
                value=starting_test_case["context"],
                visible=False,
                elem_classes=["input-box"],
            )

            tools = gr.Code(label="API Definition (Tools)", visible=False, language="json")

            user_message = gr.Textbox(
                label="User Prompt",
                lines=3,
                interactive=True,
                value=starting_test_case["user_message"],
                elem_classes=["input-box"],
            )

            assistant_message_text = gr.Textbox(
                label="Assistant Response",
                lines=3,
                interactive=True,
                visible=False,
                value=starting_test_case["assistant_message"],
                elem_classes=["input-box"],
            )

            assistant_message_json = gr.Code(
                label="Assistant Response",
                visible=False,
                language="json",
                value=None,
                elem_classes=["input-box"],
            )

            submit_button = gr.Button(
                "Evaluate",
                variant="primary",
                icon=os.path.join(os.path.dirname(os.path.abspath(__file__)), "send-white.png"),
                elem_classes="submit-button",
            )

            # result_text = gr.HTML(label='Result', elem_classes=['result-text', 'read-only', 'input-box'], visible=False, value='')
            result_text = gr.HTML(
                label="Result", elem_classes=["result-root"], show_label=True, visible=False, value=""
            )

            with Modal(visible=False, elem_classes="modal") as modal:
                prompt = gr.Markdown("")

    # events

    show_propt_button.click(
        on_show_prompt_click,
        inputs=[criteria, context, user_message, assistant_message_text, assistant_message_json, tools, state],
        outputs=prompt,
    ).then(lambda: gr.update(visible=True), None, modal)

    submit_button.click(lambda: gr.update(visible=True, value=""), None, result_text).then(
        on_submit,
        inputs=[criteria, context, user_message, assistant_message_text, assistant_message_json, tools, state],
        outputs=[result_text],
        scroll_to_output=True,
    )

    for button in [
        t for sub_catalog_name, sub_catalog_buttons in catalog_buttons.items() for t in sub_catalog_buttons.values()
    ]:
        button.click(
            change_button_color, inputs=None, outputs=[v for c in catalog_buttons.values() for v in c.values()]
        ).then(update_selected_test_case, inputs=[button, state], outputs=[state]).then(
            on_test_case_click,
            inputs=state,
            outputs=[
                test_case_name,
                criteria,
                context,
                user_message,
                assistant_message_text,
                assistant_message_json,
                tools,
                result_text,
            ],
        )

demo.launch(server_name="0.0.0.0")