import random
from pathlib import Path
from typing import Any, NamedTuple, cast

import pandas as pd
import streamlit as st
import streamlit.components.v1 as components
from jaxtyping import Float
from torch import Tensor

from deception_detection.experiment import Experiment, ExperimentConfig
from deception_detection.prompt_scorer import create_prompt_scorer
from deception_detection.results.token_html import get_token_scores_html
from deception_detection.results.utils import get_dataframe_of_scores_on_dataset
from deception_detection.scores import Scores
from deception_detection.tokenized_data import TokenizedDataset


def display_token_scores(
    str_tokens: list[str], scores: Float[Tensor, " seq"], extent: float = 6.0
) -> None:
    """Display colored token scores in Streamlit using custom HTML/CSS"""
    html = get_token_scores_html(str_tokens, scores, extent=extent)
    st.html("<style> .token { display: inline-block; padding: 2px; margin: 1px; } </style>")
    # wrap text in div to allow for background color
    div_style = "background-color: #ddd; overflow-wrap: break-word; padding: 20px;"
    full_html = f"<div style='{div_style}'>{html}</div>"
    components.html(full_html, height=1000, scrolling=True)


@st.cache_data
def load_cached_experiment(path: Path) -> Experiment:
    """Load experiment from path with caching"""
    experiment = Experiment(ExperimentConfig.from_path(path))
    return experiment


class DatasetInfo(NamedTuple):
    """Container for cached dataset information"""

    tokenized_dataset: TokenizedDataset
    normalized_scores: Scores
    dataframe: pd.DataFrame
    scorer: Any


@st.cache_data
def get_cached_dataset_data(experiment_path: Path, dataset_name: str) -> DatasetInfo:
    """Load and cache dataset-specific data"""
    experiment = load_cached_experiment(experiment_path)
    dataset = [d for d in experiment.eval_datasets if d.name == dataset_name][0]
    tok_dataset = TokenizedDataset.from_dataset(dataset, experiment.tokenizer)
    all_token_scores = experiment.get_eval_scores(all_tokens=True)[dataset_name]
    scorer = create_prompt_scorer("mean", experiment, [0.01], default_fpr=0.01)
    normalized_scores = scorer.normalize_scores(all_token_scores)
    df = get_dataframe_of_scores_on_dataset(experiment, dataset_name, scorer, normalize_scores=True)
    return DatasetInfo(
        tokenized_dataset=tok_dataset,
        normalized_scores=normalized_scores,
        dataframe=df,
        scorer=scorer,
    )


class SampleDetails(NamedTuple):
    idx: int
    label: str
    score: float
    score_percentile: float
    metadata: dict[str, Any]
    str_tokens: list[str]
    normalized_scores: Float[Tensor, " seq"]
    masked_normalized_scores: Float[Tensor, " seq"]


def get_sample_details(
    i: int,
    dataset_info: DatasetInfo,
) -> SampleDetails:
    """Create SampleDetails object for a sample"""
    # Apply mask to normalized scores
    normalized_scores = dataset_info.normalized_scores[i]
    detection_mask = dataset_info.tokenized_dataset.detection_mask
    assert detection_mask is not None
    masked_scores = normalized_scores * detection_mask[i][-len(normalized_scores) :]
    df_row = dataset_info.dataframe.iloc[i]

    return SampleDetails(
        idx=i,
        label=df_row["label"],  # type: ignore
        score=df_row["score"],  # type: ignore
        score_percentile=dataset_info.scorer.control_percentile(df_row["score"]),
        metadata={k: v for k, v in df_row.items() if k not in ["label", "score"]},  # type: ignore
        str_tokens=dataset_info.tokenized_dataset.str_tokens[i],
        normalized_scores=normalized_scores,
        masked_normalized_scores=masked_scores,
    )


def display_sample_details(sample: SampleDetails, dataset_info: DatasetInfo) -> None:
    """Display details for a single sample"""
    st.title(f"Sample {sample.idx} Details")

    # Navigation buttons in columns
    col1, col2 = st.columns([1, 4])
    with col1:
        if st.button("← Back to Overview"):
            st.query_params["sample"] = ""
            st.rerun()
    with col2:
        # Get all indices with same label
        same_label_indices = dataset_info.dataframe[
            dataset_info.dataframe["label"] == sample.label
        ].index.tolist()
        same_label_indices.remove(sample.idx)  # Remove current sample

        if same_label_indices and st.button(f"Next Random {sample.label.title()} Sample"):
            next_sample: int = random.choice(same_label_indices)  # type: ignore
            st.query_params["sample"] = str(next_sample)  # type: ignore
            st.rerun()

    # Display key information in columns
    col1, col2, col3 = st.columns(3)
    with col1:
        st.metric("Label", sample.label)
    with col2:
        st.metric("Score", f"{sample.score:.3f}")
    with col3:
        st.metric("Score Percentile", f"{sample.score_percentile:.1%}")

    # Toggle for showing scores on unmasked tokens
    st.subheader("Token Scores")
    show_scores_unmasked = st.toggle("Show scores on non-evaluation tokens", value=False)
    scores = sample.normalized_scores if show_scores_unmasked else sample.masked_normalized_scores
    display_token_scores(sample.str_tokens, scores)

    # Show additional metadata if present
    if sample.metadata:
        with st.expander("Additional Metadata"):
            for key, value in sample.metadata.items():
                # Format key with title case and replace underscores with spaces
                formatted_key = key.replace("_", " ").title()
                st.markdown(f"**:blue-background[{formatted_key}:]**\n{value}")


def display_overview_table(name: str, df: pd.DataFrame) -> None:
    """Display the main overview table with clickable links, split by label"""
    st.title(f"{name} Results")

    def get_link(i: int) -> str:
        return f'<a href="?dataset={name}&sample={i}" target="_self">Sample {i}</a>'

    # Create clickable links and prepare dataframe
    df_with_links = df.copy()
    df_with_links["link"] = [get_link(i) for i in range(len(df))]  # type: ignore
    cols = ["label", "link", "score"]
    if "experiment_variation" in df_with_links.columns:
        cols.append("experiment_variation")
    df_with_links = df_with_links[cols]
    df_with_links = df_with_links.rename(columns={"experiment_variation": "var"})  # type: ignore

    # Split into three columns
    cols = st.columns(3)

    # Define labels and their display names
    label_groups = [("honest", "Honest"), ("ambiguous", "Ambiguous"), ("deceptive", "Deceptive")]

    # Display each group in its own column
    for (label, display_name), col in zip(label_groups, cols, strict=True):
        with col:
            st.subheader(display_name)
            df_filtered = cast(pd.DataFrame, df_with_links[df_with_links["label"] == label])  # type: ignore

            # Add random sample button
            if not df_filtered.empty:
                random_idx: list[int] = df_filtered.index.to_list()
                if st.button(f"Random {display_name} Sample", key=f"random_{label}"):
                    random_sample = random.choice(random_idx)
                    st.query_params["sample"] = str(random_sample)
                    st.rerun()

            # Sort by score, highest first for deceptive, lowest first for honest
            df_filtered = df_filtered.sort_values("score", ascending=False)  # type: ignore
            # round score to 3 decimal places
            df_filtered["score"] = df_filtered["score"].round(3)
            table_body = cast(str, df_filtered.to_html(escape=False, index=False))  # type: ignore
            st.markdown(table_body, unsafe_allow_html=True)


def main(experiment_path: Path) -> None:
    st.set_page_config(page_title="Deception Detection Results", layout="wide")
    # Load experiment with caching
    experiment = load_cached_experiment(experiment_path)

    # Always show dataset selection in sidebar
    with st.sidebar:
        st.title("Dataset Selection")
        dataset_names = [d.name for d in experiment.eval_datasets]
        current_dataset = st.query_params.get("dataset")

        # Default to first dataset if none selected
        if current_dataset not in dataset_names:
            current_dataset = dataset_names[0]

        selected_dataset = st.selectbox(
            "Select Dataset", dataset_names, index=dataset_names.index(current_dataset)
        )

        if selected_dataset != current_dataset:
            st.query_params["dataset"] = selected_dataset
            st.query_params["sample"] = ""
            st.rerun()

    # Load cached data for selected dataset
    dataset_info = get_cached_dataset_data(experiment_path, selected_dataset)

    # Get the sample parameter from the URL
    sample_idx = st.query_params.get("sample")

    if sample_idx and sample_idx != "":
        # Display details page for the selected sample
        i = int(sample_idx)
        sample = get_sample_details(i, dataset_info)
        display_sample_details(sample, dataset_info)
    else:
        # Display the main overview table
        display_overview_table(selected_dataset, dataset_info.dataframe)


if __name__ == "__main__":
    import sys

    if len(sys.argv) != 2:
        print("Usage: streamlit run streamlit.py -- <experiment_path>")
        sys.exit(1)
    main(Path(sys.argv[1]))
