dashboard/reasoning.py

"""dashboard/reasoning.py — render the agent's full reasoning trace.

The trace is a list of step dicts, built incrementally during graph streaming
and rendered (live during generation, persistently afterwards) inside an
expander on the Verdict tab.

Step shapes:
    {"kind": "assistant_text", "text": str}
    {"kind": "tool_call",      "id": str, "name": str, "args": dict, "status": "pending"|"done"}
    {"kind": "tool_result",    "tool_call_id": str, "name": str, "snippet": str, "full": str}
    {"kind": "synthesis",      "status": "in_progress"|"done", "error": str | None}
"""
from __future__ import annotations

import json
from typing import Any

import streamlit as st

from dashboard.theme import (
    BG, BG_MUTED, BORDER, BORDER_STRONG,
    GREEN, RED, AMBER, BLUE,
    TEXT, TEXT_MUTED,
    INFO_BG, INFO_BORDER,
    WARN_BG, WARN_BORDER,
)
from dashboard.components import source_badge


SNIPPET_LEN = 200


# ──────────────────────────────────────────────────────────────────────────────
# Trace builder helpers — called from app.py during graph.stream
# ──────────────────────────────────────────────────────────────────────────────

def _stringify(content: Any) -> str:
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = []
        for block in content:
            if isinstance(block, dict) and block.get("type") == "text":
                parts.append(block.get("text", ""))
            elif isinstance(block, str):
                parts.append(block)
        return "\n".join(p for p in parts if p)
    return str(content)


def absorb_agent_messages(trace: list[dict], messages: list) -> None:
    """Append assistant_text + tool_call steps from a batch of agent messages."""
    for msg in messages:
        text = _stringify(getattr(msg, "content", "")).strip()
        if text:
            trace.append({"kind": "assistant_text", "text": text})
        for tc in getattr(msg, "tool_calls", None) or []:
            trace.append({
                "kind": "tool_call",
                "id": tc.get("id", ""),
                "name": tc.get("name", ""),
                "args": tc.get("args") or {},
                "status": "pending",
            })


def absorb_tool_messages(trace: list[dict], messages: list) -> None:
    """Append tool_result steps and mark matching tool_call steps done."""
    for msg in messages:
        tcid = getattr(msg, "tool_call_id", "")
        name = getattr(msg, "name", "") or ""
        full = _stringify(getattr(msg, "content", ""))
        snippet = full.strip().replace("\n", " ")
        if len(snippet) > SNIPPET_LEN:
            snippet = snippet[:SNIPPET_LEN] + "…"
        trace.append({
            "kind": "tool_result",
            "tool_call_id": tcid,
            "name": name,
            "snippet": snippet,
            "full": full,
        })
        for step in trace:
            if step["kind"] == "tool_call" and step["id"] == tcid:
                step["status"] = "done"


def mark_synthesis(trace: list[dict], status: str, error: str | None = None) -> None:
    """Add or update a synthesis step."""
    for step in trace:
        if step["kind"] == "synthesis":
            step["status"] = status
            step["error"] = error
            return
    trace.append({"kind": "synthesis", "status": status, "error": error})


# ──────────────────────────────────────────────────────────────────────────────
# Progress estimation — called from app.py live fragment
# ──────────────────────────────────────────────────────────────────────────────

def estimate_progress(trace: list[dict]) -> tuple[float, str]:
    """Return (0.0–1.0, status label) based on current trace state."""
    if not trace:
        return 0.02, "Starting…"
    for step in trace:
        if step["kind"] == "synthesis":
            if step["status"] == "done":
                return 1.0, "Brief ready — switching to Verdict tab"
            return 0.92, "Synthesizing brief…"
    completed = sum(1 for s in trace if s["kind"] == "tool_call" and s["status"] == "done")
    pending   = sum(1 for s in trace if s["kind"] == "tool_call" and s["status"] == "pending")
    if completed == 0 and pending > 0:
        return 0.06, "Running first queries…"
    value = 0.10 + min(completed / 10, 1.0) * 0.78
    return min(value, 0.88), f"Round {completed} — investigating…"


# ──────────────────────────────────────────────────────────────────────────────
# Pre-rendering: pair tool_call + tool_result into unified "tool_round" steps
# ──────────────────────────────────────────────────────────────────────────────

def _pair_steps(trace: list[dict]) -> list[dict]:
    """Walk the flat trace and merge each tool_call with its matching tool_result."""
    results_by_id: dict[str, dict] = {}
    for step in trace:
        if step["kind"] == "tool_result":
            results_by_id[step["tool_call_id"]] = step

    paired: list[dict] = []
    for step in trace:
        kind = step["kind"]
        if kind == "assistant_text":
            paired.append(step)
        elif kind == "tool_call":
            paired.append({
                "kind": "tool_round",
                "call": step,
                "result": results_by_id.get(step["id"]),
            })
        elif kind == "tool_result":
            pass  # consumed above
        elif kind == "synthesis":
            paired.append(step)
    return paired


# ──────────────────────────────────────────────────────────────────────────────
# Helpers
# ──────────────────────────────────────────────────────────────────────────────

_TOOL_CATEGORIES: dict[str, str] = {
    "search_filing": "filing",
    "get_financial_metrics": "filing",
    "get_analyst": "filing",
    "search_transcript": "transcript",
    "get_news": "news",
    "search_news": "news",
}


def _tool_category(name: str) -> str:
    for key, cat in _TOOL_CATEGORIES.items():
        if key in name:
            return cat
    return "tool"


def _fmt_args(args: dict, max_len: int = 120) -> str:
    if not args:
        return ""
    try:
        s = json.dumps(args, ensure_ascii=False, separators=(", ", ": "), default=str).strip("{}")
    except Exception:
        s = str(args)
    return s if len(s) <= max_len else s[: max_len - 1] + "…"


def _trunc(text: str, n: int = 75) -> str:
    text = text.strip().replace("\n", " ")
    return text if len(text) <= n else text[:n - 1] + "…"


def _details_block(full: str) -> str:
    """HTML <details> collapsible for full tool result. Empty if content fits in snippet."""
    if not full or len(full.strip()) <= SNIPPET_LEN:
        return ""
    escaped = full.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
    return (
        f'<details style="margin-top:7px;">'
        f'<summary style="cursor:pointer;font-size:0.72rem;color:{TEXT_MUTED};font-weight:500;'
        f'list-style:none;display:inline-flex;align-items:center;gap:4px;user-select:none;">'
        f'<span style="font-size:0.58rem;">▶</span> Show full result'
        f'</summary>'
        f'<div style="margin-top:6px;padding:8px 10px;background:{BG_MUTED};'
        f'border:1px solid {BORDER};border-radius:6px;'
        f'font-family:ui-monospace,SFMono-Regular,Menlo,monospace;'
        f'font-size:0.72rem;line-height:1.5;color:{TEXT_MUTED};'
        f'white-space:pre-wrap;word-break:break-all;'
        f'max-height:300px;overflow-y:auto;">{escaped}</div>'
        f'</details>'
    )


def _details_wrap(summary_html: str, body_html: str, is_open: bool) -> str:
    """Wrap body_html in a <details> block with the given summary."""
    open_attr = " open" if is_open else ""
    return (
        f'<details{open_attr} style="margin:4px 0;">'
        f'<summary style="list-style:none;cursor:pointer;user-select:none;">'
        f'{summary_html}'
        f'</summary>'
        f'{body_html}'
        f'</details>'
    )


# ──────────────────────────────────────────────────────────────────────────────
# Step renderers
# ──────────────────────────────────────────────────────────────────────────────

def _render_assistant_text(step: dict, is_recent: bool = True) -> None:
    preview = _trunc(step["text"])
    summary = (
        f'<div style="background:{INFO_BG};border:1px solid {INFO_BORDER};'
        f'border-left:4px solid {BLUE};border-radius:0 10px 10px 0;'
        f'padding:7px 14px;font-size:0.8rem;color:{TEXT_MUTED};">'
        f'💭 <em style="color:{BLUE};font-weight:600;">Thinking</em>'
        f'<span style="margin-left:8px;color:{TEXT_MUTED};">{preview}</span>'
        f'</div>'
    )
    body = (
        f'<div style="background:{INFO_BG};border:1px solid {INFO_BORDER};'
        f'border-left:4px solid {BLUE};border-radius:0 10px 10px 0;'
        f'padding:10px 14px;margin-top:2px;">'
        f'<div style="font-size:0.82rem;line-height:1.55;color:{TEXT};white-space:pre-wrap;">'
        f'{step["text"]}</div>'
        f'</div>'
    )
    st.markdown(_details_wrap(summary, body, is_recent), unsafe_allow_html=True)


def _render_tool_round(step: dict, is_recent: bool = True) -> None:
    call = step["call"]
    result = step["result"]
    pending = result is None

    if pending:
        bg, border_color, accent, icon = WARN_BG, WARN_BORDER, AMBER, "🔄"
    else:
        bg, border_color, accent, icon = BG_MUTED, BORDER, GREEN, "✅"

    name = call["name"]
    category = _tool_category(name)
    badge_html = source_badge(category)
    args_str = _fmt_args(call.get("args") or {})
    snippet_preview = _trunc(result["snippet"] if result else "") if result else ""

    # Summary row (shown when collapsed)
    status_text = f'<em style="color:{AMBER};">running…</em>' if pending else (
        f'<span style="color:{TEXT_MUTED};font-size:0.75rem;">{snippet_preview}</span>' if snippet_preview else ""
    )
    summary = (
        f'<div style="background:{bg};border:1px solid {border_color};'
        f'border-left:4px solid {accent};border-radius:0 10px 10px 0;'
        f'padding:7px 14px;display:flex;align-items:center;justify-content:space-between;gap:8px;">'
        f'<div style="display:flex;align-items:center;gap:6px;font-size:0.8rem;color:{TEXT};min-width:0;">'
        f'{icon} '
        f'<code style="background:{BG};border:1px solid {BORDER_STRONG};'
        f'padding:1px 6px;border-radius:4px;font-size:0.75rem;">{name}</code>'
        f'<span style="color:{TEXT_MUTED};font-size:0.75rem;overflow:hidden;text-overflow:ellipsis;white-space:nowrap;">'
        f'{status_text}</span>'
        f'</div>'
        f'{badge_html}'
        f'</div>'
    )

    # Full body (shown when open)
    args_html = (
        f'<div style="margin-top:6px;padding:4px 9px;background:{BG};'
        f'border:1px solid {border_color};border-radius:5px;'
        f'font-family:ui-monospace,SFMono-Regular,Menlo,monospace;'
        f'font-size:0.71rem;color:{TEXT_MUTED};word-break:break-all;">'
        f'{args_str}</div>'
    ) if args_str else ""

    if result:
        snippet = result.get("snippet") or "(empty)"
        details = _details_block(result.get("full") or "")
        result_html = (
            f'<div style="margin-top:9px;padding-top:8px;border-top:1px solid {border_color};">'
            f'<div style="font-size:0.6rem;font-weight:700;text-transform:uppercase;'
            f'letter-spacing:0.08em;color:{GREEN};margin-bottom:4px;">↩ Result</div>'
            f'<div style="font-size:0.78rem;color:{TEXT_MUTED};line-height:1.45;">{snippet}</div>'
            f'{details}'
            f'</div>'
        )
    else:
        result_html = (
            f'<div style="margin-top:8px;padding-top:8px;border-top:1px solid {border_color};'
            f'font-size:0.75rem;color:{AMBER};font-style:italic;">Waiting for result…</div>'
        )

    body = (
        f'<div style="background:{bg};border:1px solid {border_color};'
        f'border-left:4px solid {accent};border-radius:0 10px 10px 0;'
        f'padding:10px 14px;margin-top:2px;">'
        f'<div style="display:flex;align-items:center;justify-content:space-between;gap:8px;">'
        f'<div style="font-size:0.82rem;color:{TEXT};display:flex;align-items:center;gap:7px;">'
        f'{icon}'
        f'<code style="background:{BG};border:1px solid {BORDER_STRONG};'
        f'padding:1px 7px;border-radius:5px;font-size:0.77rem;">{name}</code>'
        f'</div>'
        f'{badge_html}'
        f'</div>'
        f'{args_html}'
        f'{result_html}'
        f'</div>'
    )

    st.markdown(_details_wrap(summary, body, is_recent), unsafe_allow_html=True)


def _render_synthesis(step: dict) -> None:
    if step.get("error"):
        icon, label, color = "❌", f"Synthesis failed: {step['error']}", RED
    elif step["status"] == "in_progress":
        icon, label, color = "🔄", "Synthesizing brief…", AMBER
    else:
        icon, label, color = "✅", "Brief synthesized", GREEN
    st.markdown(
        f'<div style="background:{BG_MUTED};border:1px solid {BORDER};'
        f'border-left:4px solid {color};border-radius:0 10px 10px 0;'
        f'padding:10px 14px;margin:6px 0;">'
        f'<div style="font-size:0.85rem;font-weight:600;color:{TEXT};">'
        f'{icon} {label}</div></div>',
        unsafe_allow_html=True,
    )


# ──────────────────────────────────────────────────────────────────────────────
# Public API
# ──────────────────────────────────────────────────────────────────────────────

def render_trace_body(trace: list[dict], recent_n: int = 2) -> None:
    """Render every step of the trace. Older steps are collapsed, last recent_n are expanded."""
    if not trace:
        st.caption("No reasoning steps yet.")
        return
    paired = _pair_steps(trace)
    n = len(paired)
    for i, step in enumerate(paired):
        is_recent = i >= n - recent_n
        kind = step["kind"]
        if kind == "assistant_text":
            _render_assistant_text(step, is_recent=is_recent)
        elif kind == "tool_round":
            _render_tool_round(step, is_recent=is_recent)
        elif kind == "synthesis":
            _render_synthesis(step)


def render(trace: list[dict]) -> None:
    """Render the reasoning expander on the Verdict tab (collapsed by default)."""
    if not trace:
        return
    n_rounds = sum(1 for s in trace if s["kind"] == "tool_call")
    label = f"🧠 Reasoning trace — {n_rounds} tool round{'s' if n_rounds != 1 else ''}, {len(trace)} steps"
    with st.expander(label, expanded=False):
        render_trace_body(trace)