"""Streamlit UI to view AI audit rollouts.

Could probably be adapted to generic rollout viewer."""

import html
import json
from pathlib import Path
from typing import Any

import streamlit as st

from deception_detection.rollouts.rollout_dialogue import RolloutDialogue


def load_json(path: Path) -> dict[str, Any]:
    """Load JSON data from file."""
    with open(path) as f:
        return json.load(f)


def display_rollout(rollout: RolloutDialogue) -> None:
    """Display a single rollout's details."""
    assert rollout.metadata is not None
    st.subheader(f"Experiment: {rollout.metadata['experiment']}")

    # Display deception stats
    col1, col2 = st.columns(2)
    with col1:
        deception_percentage = rollout.deception_percentage or float("nan")
        st.metric("Deception Percentage", f"{deception_percentage:.0%}")
    with col2:
        st.metric("Graded Responses", len(rollout.graded_deception))

    # Show input messages
    with st.expander("View Setup", expanded=False):
        for msg in rollout.input_messages[:-2]:
            st.subheader(f"{msg.role}")
            with st.container(border=True):
                st.markdown(
                    f"<div style='white-space: pre-wrap;'>{html.escape(msg.content)}</div>",
                    unsafe_allow_html=True,
                )

    with st.expander("View Last Message", expanded=True):
        msg = rollout.input_messages[-2]
        st.subheader(f"{msg.role}")
        with st.container(border=True):
            st.markdown(
                f"<div style='white-space: pre-wrap;'>{html.escape(msg.content)}</div>",
                unsafe_allow_html=True,
            )

    st.header("Responses")
    zipped = zip(
        rollout.output_messages, rollout.graded_deception, rollout.grader_reasoning, strict=True
    )
    for i, (out, grade, reasoning) in enumerate(zipped):
        with st.expander(f"{i}: Graded {grade}", expanded=False):
            st.subheader("Response")
            with st.container(border=True):
                st.markdown(
                    f"<div style='white-space: pre-wrap;'>{html.escape(out.content)}</div>",
                    unsafe_allow_html=True,
                )
            st.subheader("Grader Reasoning")
            with st.container(border=True):
                st.markdown(
                    f"<div style='white-space: pre-wrap;'>{html.escape(reasoning)}</div>",
                    unsafe_allow_html=True,
                )


def main() -> None:
    st.title("AI Audit Rollouts Analysis")

    # Load data
    data = load_json(Path("data/rollouts/ai_audit__no_reasoning__llama-70b.json"))

    # Display dataset info
    st.write(f"**Dataset Class:** {data['dataset_class']}")
    st.write(f"**Base Name:** {data['base_name']}")
    st.write(f"**Variant:** {data['variant']}")
    st.write(f"**Model:** {data['model']}")

    rollouts = [RolloutDialogue(**r) for r in data["rollouts"]]

    # Add sidebar for rollout selection
    st.sidebar.header("Rollout Selection")

    def get_rollout_info(i: int) -> str:
        metadata = rollouts[i].metadata
        assert metadata is not None
        expected_label = rollouts[i].expected_label.value
        return f"Rollout {i} ({metadata['experiment_variation']}) - {expected_label}"

    selected_index: int = st.sidebar.selectbox(
        "Choose Rollout", options=range(len(rollouts)), format_func=get_rollout_info
    )

    # Display selected rollout
    display_rollout(rollouts[selected_index])


if __name__ == "__main__":
    main()
