"""Session Viewer — Streamlit app.

Run:  streamlit run notebooks/apps/session_viewer.py
"""

import sys
import os
from datetime import datetime, timezone

import pandas as pd
import streamlit as st

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from utils.connection import init, query_df


# ---------------------------------------------------------------------------
# Connection
# ---------------------------------------------------------------------------

@st.cache_resource
def get_connection():
    return init()


# ---------------------------------------------------------------------------
# Data loaders
# ---------------------------------------------------------------------------

@st.cache_data(ttl=120)
def load_users(_conn) -> list[str]:
    df = query_df(_conn, """
        SELECT   user_email, MAX(last_updated) AS last_active
        FROM     sessions
        WHERE    user_email IS NOT NULL
        GROUP BY user_email
        ORDER BY last_active DESC
    """)
    return df["user_email"].tolist()


@st.cache_data(ttl=120)
def load_sessions_for_user(_conn, email: str) -> pd.DataFrame:
    return query_df(_conn, """
        SELECT id, source, model, repo_name, git_branch, org,
               device_name, cwd, first_seen, last_updated
        FROM   sessions
        WHERE  user_email = %s
        ORDER  BY first_seen DESC
    """, (email,))


@st.cache_data(ttl=60)
def load_messages(_conn, session_id: str) -> pd.DataFrame:
    return query_df(_conn, """
        SELECT m.id, m.msg_type, m.timestamp, m.content, m.thinking, m.model, m.source
        FROM   messages m
        WHERE  m.session_id = %s
        ORDER  BY m.timestamp ASC
    """, (session_id,))


@st.cache_data(ttl=60)
def load_tokens(_conn, session_id: str) -> pd.DataFrame:
    return query_df(_conn, """
        SELECT t.message_id, t.input_tokens, t.output_tokens, t.cached_tokens, t.thinking_tokens
        FROM   token_usage t
        JOIN   messages m ON m.id = t.message_id
        WHERE  m.session_id = %s
    """, (session_id,))


@st.cache_data(ttl=60)
def load_tool_calls(_conn, session_id: str) -> pd.DataFrame:
    return query_df(_conn, """
        SELECT tc.message_id, tc.tool_id, tc.tool_name, tc.tool_input
        FROM   tool_calls tc
        JOIN   messages m ON m.id = tc.message_id
        WHERE  m.session_id = %s
    """, (session_id,))


@st.cache_data(ttl=60)
def load_tool_results(_conn, session_id: str) -> pd.DataFrame:
    return query_df(_conn, """
        SELECT tr.message_id, tr.call_id, tr.output, tr.status
        FROM   tool_results tr
        JOIN   messages m ON m.id = tr.message_id
        WHERE  m.session_id = %s
    """, (session_id,))


# ---------------------------------------------------------------------------
# Rendering
# ---------------------------------------------------------------------------

TYPE_LABELS = {
    "user":        "USER",
    "assistant":   "ASST",
    "tool_call":   "CALL",
    "tool_result": "RSLT",
    "system":      "SYS ",
}

TYPE_COLORS = {
    "user":        "#3b82f6",
    "assistant":   "#22c55e",
    "tool_call":   "#eab308",
    "tool_result": "#a855f7",
    "system":      "#6b7280",
}


def render_session_meta(row: pd.Series) -> None:
    parts = [
        f"repo: {row.get('repo_name') or '—'}",
        f"branch: {row.get('git_branch') or '—'}",
        f"model: {row.get('model') or '—'}",
        f"source: {row.get('source') or '—'}",
        f"device: {row.get('device_name') or '—'}",
        f"cwd: {row.get('cwd') or '—'}",
        f"first: {_fmt_ts(row.get('first_seen'))}",
        f"last: {_fmt_ts(row.get('last_updated'))}",
    ]
    st.code("\n".join(parts), language=None)


def render_token_summary(tokens_df: pd.DataFrame) -> None:
    if tokens_df.empty:
        return
    t = tokens_df[["input_tokens", "output_tokens", "cached_tokens", "thinking_tokens"]].sum()
    cols = st.columns(5)
    cols[0].metric("Input", f"{int(t['input_tokens']):,}")
    cols[1].metric("Output", f"{int(t['output_tokens']):,}")
    cols[2].metric("Cached", f"{int(t['cached_tokens']):,}")
    cols[3].metric("Thinking", f"{int(t['thinking_tokens']):,}")
    cols[4].metric("Billed", f"{int(t['input_tokens'] + t['output_tokens']):,}")


def render_message(
    msg: pd.Series,
    tokens_map: dict,
    tc_map: dict,
    tr_map: dict,
    collapsed_types: set[str],
) -> None:
    mid = msg["id"]
    msg_type = msg["msg_type"] or "system"
    color = TYPE_COLORS.get(msg_type, "#6b7280")
    label = TYPE_LABELS.get(msg_type, "????")
    ts = _fmt_ts(msg["timestamp"])
    is_collapsed = msg_type in collapsed_types

    # header: always visible
    token_str = ""
    if mid in tokens_map:
        tk = tokens_map[mid]
        token_str = f"  [in:{tk['input_tokens']} out:{tk['output_tokens']} cache:{tk['cached_tokens']} think:{tk['thinking_tokens']}]"

    tool_str = ""
    if mid in tc_map:
        tool_str = f"  tool={tc_map[mid]['tool_name']}"
    if mid in tr_map:
        tool_str = f"  status={tr_map[mid].get('status', '?')}"

    header = f"[{label}] {ts}  {msg.get('model') or ''}{tool_str}{token_str}"

    st.markdown(
        f"<pre style='margin:0;padding:4px 8px;background:transparent;"
        f"border-left:3px solid {color};color:{color};font-size:13px;'>{_esc(header)}</pre>",
        unsafe_allow_html=True,
    )

    if is_collapsed:
        return

    # body parts as plain text blocks
    parts = []

    if msg.get("thinking"):
        parts.append(f"[thinking]\n{msg['thinking']}")

    if mid in tc_map:
        tc = tc_map[mid]
        inp = tc.get("tool_input")
        if inp:
            inp_str = str(inp) if not isinstance(inp, str) else inp
            if len(inp_str) > 2000:
                inp_str = inp_str[:2000] + "\n… [truncated]"
            parts.append(f"[tool input]\n{inp_str}")

    if mid in tr_map:
        tr = tr_map[mid]
        out = tr.get("output")
        if out:
            out_str = str(out)
            if len(out_str) > 3000:
                out_str = out_str[:3000] + "\n… [truncated]"
            parts.append(f"[tool output]\n{out_str}")

    if msg.get("content"):
        parts.append(str(msg["content"]))

    if parts:
        st.code("\n".join(parts), language=None)


def _fmt_ts(val) -> str:
    if val is None or pd.isna(val):
        return "—"
    if isinstance(val, datetime):
        return val.strftime("%Y-%m-%d %H:%M:%S")
    return str(val)[:19]


def _esc(s: str) -> str:
    return s.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")


def _build_map(df: pd.DataFrame, key_col: str) -> dict:
    result = {}
    for _, row in df.iterrows():
        result[row[key_col]] = row.to_dict()
    return result


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main() -> None:
    st.set_page_config(page_title="session viewer", layout="wide")
    st.title("Session Viewer")

    conn = get_connection()

    # --- sidebar ---
    with st.sidebar:
        st.header("Filters")

        users = load_users(conn)
        selected_user = st.selectbox("User", options=[""] + users, format_func=lambda x: x or "— select —")

        if not selected_user:
            st.info("Pick a user to start.")
            return

        sessions_df = load_sessions_for_user(conn, selected_user)
        if sessions_df.empty:
            st.warning("No sessions found.")
            return

        session_options = []
        for _, s in sessions_df.iterrows():
            date = _fmt_ts(s["first_seen"])[:16]
            repo = s["repo_name"] or "no-repo"
            label = f"{date} · {repo} · {s['id'][:12]}…"
            session_options.append((s["id"], label))

        selected_idx = st.selectbox(
            "Session",
            options=range(len(session_options)),
            format_func=lambda i: session_options[i][1],
        )
        selected_session_id = session_options[selected_idx][0]

        st.divider()

        type_filter = st.multiselect(
            "Show types",
            options=["user", "assistant", "tool_call", "tool_result", "system"],
            default=["user", "assistant", "tool_call", "tool_result", "system"],
        )

        search_term = st.text_input("Search")

        st.divider()
        st.subheader("Collapse body")
        collapse_user = st.checkbox("Collapse USER", value=False)
        collapse_asst = st.checkbox("Collapse ASSISTANT", value=False)
        collapse_call = st.checkbox("Collapse TOOL_CALL", value=True)
        collapse_rslt = st.checkbox("Collapse TOOL_RESULT", value=True)
        collapse_sys = st.checkbox("Collapse SYSTEM", value=True)

        st.divider()
        if st.button("Clear cache & reload"):
            st.cache_data.clear()
            st.rerun()

    # --- build collapsed set ---
    collapsed_types: set[str] = set()
    if collapse_user: collapsed_types.add("user")
    if collapse_asst: collapsed_types.add("assistant")
    if collapse_call: collapsed_types.add("tool_call")
    if collapse_rslt: collapsed_types.add("tool_result")
    if collapse_sys:  collapsed_types.add("system")

    # --- main area ---
    session_row = sessions_df[sessions_df["id"] == selected_session_id].iloc[0]
    render_session_meta(session_row)

    messages_df = load_messages(conn, selected_session_id)
    tokens_df = load_tokens(conn, selected_session_id)
    tc_df = load_tool_calls(conn, selected_session_id)
    tr_df = load_tool_results(conn, selected_session_id)

    render_token_summary(tokens_df)
    st.divider()

    tokens_map = _build_map(tokens_df, "message_id")
    tc_map = _build_map(tc_df, "message_id")
    tr_map = _build_map(tr_df, "message_id")

    # filter
    filtered = messages_df.copy()
    if type_filter:
        filtered = filtered[filtered["msg_type"].isin(type_filter)]
    if search_term:
        term = search_term.lower()
        mask = filtered.apply(
            lambda row: term in str(row.get("content", "") or "").lower()
            or term in str(row.get("thinking", "") or "").lower()
            or (row["id"] in tc_map and term in str(tc_map[row["id"]].get("tool_name", "")).lower())
            or (row["id"] in tc_map and term in str(tc_map[row["id"]].get("tool_input", "")).lower())
            or (row["id"] in tr_map and term in str(tr_map[row["id"]].get("output", "")).lower()),
            axis=1,
        )
        filtered = filtered[mask]

    st.caption(f"{len(filtered)} / {len(messages_df)} messages")

    for _, msg in filtered.iterrows():
        render_message(msg, tokens_map, tc_map, tr_map, collapsed_types)


main()
