import json
from pathlib import Path

import pandas as pd
import plotly.graph_objects as go
import streamlit as st
from datasets import load_dataset

DATASET_ID = "hoololi/llm-calculations"
LOCAL_PARQUET = Path("llm_calculations.parquet")
MODEL_INFO_PATH = Path("model_info_snapshot.json")
MODE_ORDER = ["raw", "prompted", "tool"]
MODE_COLORS = {"raw": "#6EA8FE", "prompted": "#FFB86B", "tool": "#69DB7C"}

st.set_page_config(page_title="LLM Calculations Explorer", page_icon="🧮", layout="wide")


@st.cache_data(show_spinner="Loading dataset...")
def load_data() -> pd.DataFrame:
    if LOCAL_PARQUET.exists():
        df = pd.read_parquet(LOCAL_PARQUET)
    else:
        ds = load_dataset(DATASET_ID, split="train")
        df = ds.to_pandas()

    if "operation" not in df.columns and "question" in df.columns:
        df["operation"] = df["question"]
    if "cost" in df.columns:
        df["cost"] = pd.to_numeric(df["cost"], errors="coerce")
    if "mode" in df.columns:
        df["mode"] = pd.Categorical(df["mode"], categories=MODE_ORDER, ordered=True)
    return df


@st.cache_data(show_spinner=False)
def load_model_info() -> dict:
    if not MODEL_INFO_PATH.exists():
        return {}
    return json.loads(MODEL_INFO_PATH.read_text(encoding="utf-8"))


def selected_models(model_a: str, model_b: str) -> list[str]:
    models = [model_a]
    if model_b and model_b != model_a:
        models.append(model_b)
    return models


def summarize(df: pd.DataFrame, model_order: list[str]) -> pd.DataFrame:
    if df.empty:
        return pd.DataFrame()
    summary = (
        df.groupby(["model", "mode"], observed=False)
        .agg(
            rows=("mode", "size"),
            accuracy_pct=("final_result_correct", lambda s: s.mean() * 100),
            run_success_pct=("run_success", lambda s: s.mean() * 100 if s.notna().any() else None),
            avg_total_tokens=("total_tokens", "mean"),
            avg_input_tokens=("input_tokens", "mean"),
            avg_output_tokens=("output_tokens", "mean"),
            avg_cost=("cost", "mean"),
            total_cost=("cost", "sum"),
            avg_latency_seconds=("latency_seconds", "mean"),
        )
        .reset_index()
    )
    summary["mode"] = pd.Categorical(summary["mode"], categories=MODE_ORDER, ordered=True)
    summary["model"] = pd.Categorical(summary["model"], categories=model_order, ordered=True)
    summary = summary.sort_values(["model", "mode"])
    summary["model"] = summary["model"].astype(str)
    summary["mode"] = summary["mode"].astype(str)
    return summary.round(4)


def metric_chart(summary: pd.DataFrame, metric: str, yaxis_title: str, percent=False) -> go.Figure:
    fig = go.Figure()
    for mode in MODE_ORDER:
        part = summary[summary["mode"] == mode]
        fig.add_bar(
            x=part["model"],
            y=part[metric],
            name=mode,
            marker_color=MODE_COLORS[mode],
            hovertemplate=f"{yaxis_title}: %{{y:.4g}}<extra></extra>",
        )
    fig.update_layout(
        template="plotly_dark",
        barmode="group",
        height=260,
        margin=dict(l=35, r=10, t=10, b=35),
        legend=dict(orientation="h", yanchor="top", y=0.99, xanchor="right", x=0.99),
        paper_bgcolor="rgba(0,0,0,0)",
        plot_bgcolor="rgba(255,255,255,0.035)",
        yaxis_title=yaxis_title,
        xaxis_title=None,
    )
    fig.update_xaxes(categoryorder="array", categoryarray=summary["model"].drop_duplicates().tolist(), tickangle=0)
    fig.update_yaxes(gridcolor="rgba(255,255,255,0.12)")
    if percent:
        fig.update_yaxes(range=[0, 100], tickvals=[0, 20, 40, 60, 80, 100])
    return fig


def price_per_million(value) -> str:
    if value in (None, ""):
        return "n/a"
    try:
        return f"${float(value) * 1_000_000:.2f}/1M tokens"
    except (TypeError, ValueError):
        return str(value)


def model_card(model: str, model_info: dict) -> None:
    info = model_info.get(model, {})
    if not info:
        st.caption(f"No OpenRouter snapshot information available for `{model}`.")
        return
    pricing = info.get("pricing", {}) if isinstance(info.get("pricing", {}), dict) else {}
    architecture = info.get("architecture", {}) if isinstance(info.get("architecture", {}), dict) else {}
    provider = info.get("top_provider", {}) if isinstance(info.get("top_provider", {}), dict) else {}
    supported = info.get("supported_parameters", []) or []
    description = (info.get("description") or "").strip().replace("\n", " ")
    if len(description) > 320:
        description = description[:320].rstrip() + "…"

    st.markdown(f"### {info.get('name') or model}")
    st.code(model, language=None)
    st.markdown(
        f"""
- **Context length:** {info.get('context_length') or provider.get('context_length') or 'n/a'}
- **Max completion tokens:** {provider.get('max_completion_tokens') or 'n/a'}
- **Modality:** {architecture.get('modality') or 'n/a'}
- **Tokenizer:** {architecture.get('tokenizer') or 'n/a'}
- **Input price:** {price_per_million(pricing.get('prompt') or pricing.get('input'))}
- **Output price:** {price_per_million(pricing.get('completion') or pricing.get('output'))}
- **Supported parameters:** {', '.join(supported[:8]) if supported else 'n/a'}{'…' if len(supported) > 8 else ''}
"""
    )
    if description:
        st.caption(description)


df = load_data()
model_info = load_model_info()
models = sorted(df["model"].dropna().unique().tolist())
operation_categories = sorted(df["operation_category"].dropna().unique().tolist()) if "operation_category" in df.columns else []

st.title("🧮 LLM Calculations Explorer")
st.markdown("Simple explorer for the `hoololi/llm-calculations` dataset. Compare raw, prompted, and tool modes.")

with st.container(border=True):
    c1, c2, c3, c4 = st.columns([1.2, 1.2, 1.4, 1])
    with c1:
        model_a = st.selectbox("Model A", models, index=0)
    with c2:
        model_b = st.selectbox("Model B optional", [""] + models, index=0)
    with c3:
        selected_categories = st.multiselect(
            "Operation categories",
            operation_categories,
            default=operation_categories,
        )
    with c4:
        run_error_policy = st.selectbox(
            "Run errors",
            ["Count as failures", "Exclude from metrics"],
            index=0,
        )

chosen = selected_models(model_a, model_b)
view = df[df["model"].isin(chosen)].copy()
if selected_categories and "operation_category" in view.columns:
    view = view[view["operation_category"].isin(selected_categories)]
if run_error_policy == "Exclude from metrics" and "run_success" in view.columns:
    view = view[view["run_success"] == True]

summary = summarize(view, chosen)

st.subheader("Metrics summary")
st.dataframe(summary, width="stretch", hide_index=True)

r1c1, r1c2 = st.columns(2)
with r1c1:
    st.markdown("**Accuracy**")
    st.plotly_chart(metric_chart(summary, "accuracy_pct", "Accuracy (%)", True), width="stretch")
with r1c2:
    st.markdown("**Tokens**")
    st.plotly_chart(metric_chart(summary, "avg_total_tokens", "Average total tokens"), width="stretch")

r2c1, r2c2 = st.columns(2)
with r2c1:
    st.markdown("**Cost**")
    st.plotly_chart(metric_chart(summary, "avg_cost", "Average cost, USD"), width="stretch")
with r2c2:
    st.markdown("**Latency**")
    st.plotly_chart(metric_chart(summary, "avg_latency_seconds", "Seconds"), width="stretch")

with st.expander("Selected model details from OpenRouter snapshot", expanded=False):
    info_cols = st.columns(len(chosen))
    for col, model in zip(info_cols, chosen):
        with col:
            model_card(model, model_info)

st.divider()
st.subheader("Browse dataset records")
st.caption("Inspect the individual rows behind the aggregate metrics.")

browse_df = view.copy()

b1, b2, b3 = st.columns([1, 1, 1.4])
with b1:
    browse_modes = st.multiselect("Modes", MODE_ORDER, default=MODE_ORDER)
with b2:
    correctness_filter = st.selectbox(
        "Final answer",
        ["All", "Correct only", "Incorrect only", "Run errors only", "Tool issues only"],
        index=0,
    )
with b3:
    search_text = st.text_input("Search operation or answer", value="")

if browse_modes:
    browse_df = browse_df[browse_df["mode"].astype(str).isin(browse_modes)]

if correctness_filter == "Correct only":
    browse_df = browse_df[browse_df["final_result_correct"] == True]
elif correctness_filter == "Incorrect only":
    browse_df = browse_df[browse_df["final_result_correct"] == False]
elif correctness_filter == "Run errors only":
    browse_df = browse_df[browse_df["run_success"] == False]
elif correctness_filter == "Tool issues only":
    browse_df = browse_df[
        (browse_df["mode"].astype(str) == "tool")
        & (
            (browse_df.get("tool_error_type") != "none")
            | (browse_df.get("tool_result_correct") == False)
            | (browse_df.get("tool_expression_scope").isin(["partial_expression", "unrelated_expression", "unknown"]))
        )
    ]

if search_text.strip():
    needle = search_text.strip().lower()
    browse_df = browse_df[
        browse_df["operation"].fillna("").str.lower().str.contains(needle, regex=False)
        | browse_df["answer_text"].fillna("").str.lower().str.contains(needle, regex=False)
    ]

browse_df = browse_df.sort_values(["operation_id", "model", "mode"])

display_cols = [
    "operation_id",
    "operation_category",
    "model",
    "mode",
    "operation",
    "correct_result_str",
    "extracted_answer_str",
    "final_result_correct",
    "run_success",
    "tool_expression_scope",
    "tool_error_type",
]
display_cols = [col for col in display_cols if col in browse_df.columns]

st.dataframe(browse_df[display_cols], width="stretch", hide_index=True, height=260)

if not browse_df.empty:
    browse_df = browse_df.reset_index(drop=True)
    row_labels = browse_df.apply(
        lambda r: f"{r['operation_id']} | {r['model']} | {r['mode']} | final_correct={r['final_result_correct']}",
        axis=1,
    ).tolist()
    selected_label = st.selectbox("Record details", row_labels, index=0)
    selected_row = browse_df.iloc[row_labels.index(selected_label)]

    d1, d2 = st.columns([1, 1])
    with d1:
        st.markdown("#### Operation")
        st.write(selected_row.get("operation"))
        st.caption(f"Category: {selected_row.get('operation_category')}")
        st.markdown("#### Expected / extracted")
        st.json(
            {
                "correct_result": selected_row.get("correct_result_str", selected_row.get("correct_result")),
                "extracted_answer": selected_row.get("extracted_answer_str", selected_row.get("extracted_answer")),
                "final_result_correct": bool(selected_row.get("final_result_correct")),
                "extraction_error_type": selected_row.get("extraction_error_type"),
                "answer_empty": bool(selected_row.get("answer_empty")),
            }
        )
    with d2:
        st.markdown("#### Diagnostics")
        st.json(
            {
                "run_success": bool(selected_row.get("run_success")),
                "run_error": selected_row.get("run_error"),
                "tool_called": selected_row.get("tool_called"),
                "tool_expression_scope": selected_row.get("tool_expression_scope"),
                "tool_error_type": selected_row.get("tool_error_type"),
                "tool_result_correct": selected_row.get("tool_result_correct"),
            }
        )

    st.markdown("#### Model answer")
    st.code(selected_row.get("answer_text") or "", language=None)

    if str(selected_row.get("mode")) == "tool":
        st.markdown("#### Tool details")
        st.json(
            {
                "tool_expression": selected_row.get("tool_expression"),
                "tool_result": selected_row.get("tool_result_str", selected_row.get("tool_result")),
                "tool_error": selected_row.get("tool_error"),
                "tool_calls": selected_row.get("tool_calls"),
            }
        )
else:
    st.info("No records match the current filters.")