"""
Streamlit app to view insider trading transcripts and their classifications
"""

from dataclasses import dataclass

import streamlit as st

from deception_detection.data.base import DialogueDataset
from deception_detection.repository import DatasetRepository
from deception_detection.types import Dialogue


@dataclass
class TranscriptSummary:
    index: int
    trade_classification: str
    report_classification: str
    label: str
    messages: Dialogue


def load_data() -> DialogueDataset:
    """Load the rollout data from the JSON file"""
    dataset = DatasetRepository().get("insider_trading__onpolicy", "llama-70b")
    return dataset


def get_transcript_summaries(data: DialogueDataset) -> list[TranscriptSummary]:
    """Extract summary information for each transcript"""
    summaries = []
    assert data.metadata is not None
    for i in range(len(data)):
        summary = TranscriptSummary(
            index=i,
            trade_classification=data.metadata["trade_classification"][i],
            report_classification=data.metadata["report_classification"][i],
            label=data.metadata["label"][i],
            messages=data.dialogues[i],
        )
        summaries.append(summary)
    return summaries


def color_classification(classification: str) -> str:
    """Get formatted markdown text with color corresponding to classification"""
    colors = {
        # Trade classifications
        "Yes": "blue",
        "No": "red",
        "Ambiguous": "orange",
        # Report classifications
        "Explicit": "blue",
        "Concealed": "red",
        "Implied": "orange",
        # Overall label
        "honest": "blue",
        "ambiguous": "orange",
        "deceptive": "red",
    }
    if classification in colors:
        return f":{colors[classification]}[{classification}]"
    else:
        return classification


def show_classifications(summary: TranscriptSummary):
    """Display the classifications for a single transcript"""
    cols = st.columns(3)
    cols[0].markdown(f"""**Trade** - {color_classification(summary.trade_classification)}""")
    cols[1].markdown(f"""**Report** - {color_classification(summary.report_classification)}""")
    cols[2].markdown(f"""**Overall Label** - {color_classification(summary.label)}""")


def main():
    st.title("Insider Trading Transcripts Viewer")

    # Load data
    data = load_data()
    summaries = get_transcript_summaries(data)

    # Create summary table
    st.subheader("Transcript Summaries")

    # add styling
    st.markdown("<style> .container { padding: 10px; } </style>")

    for summary in summaries:
        st.write(f"### Sample {summary.index}")
        show_classifications(summary)

        with st.expander("Conversation"):
            for msg in [summary.messages[-3], summary.messages[-1]]:
                with st.chat_message(msg.role):
                    st.markdown(msg.content.replace("\n", "\n\n"))


if __name__ == "__main__":
    main()
