#!/usr/bin/env python
"""Streamlit log viewer for OpenHands agent completion logs."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import streamlit as st

LOG_ROOTS = [
    Path("output/Agent/logs"),
]


def load_runs(root: Path) -> list[Path]:
    if not root.exists():
        return []
    return sorted([p for p in root.iterdir() if p.is_dir()])


def load_log_files(run_dir: Path) -> list[Path]:
    return sorted(run_dir.glob("*.json"))


def load_json(path: Path) -> dict[str, Any]:
    try:
        return json.loads(path.read_text())
    except Exception as exc:  # pragma: no cover - UI guard
        st.error(f"Failed to parse {path}: {exc}")
        return {}


def render_message(msg: dict[str, Any]) -> None:
    msg_type = msg.get("type") or msg.get("role")
    if msg_type == "message":
        role = msg.get("role", "user")
        st.markdown(f"**{role}**")
        for chunk in msg.get("content", []):
            if isinstance(chunk, dict) and chunk.get("text"):
                st.write(chunk["text"])
    elif msg_type == "function_call":
        args = msg.get("arguments", "")
        preview = (args[:80] + "...") if len(args) > 80 else args
        st.markdown(f"**Tool Call:** `{msg.get('name')}` — {preview}")
        st.code(msg.get("arguments"), language="json")
    elif msg_type == "function_call_output":
        st.markdown("**Tool Output**")
        st.code(msg.get("output", ""), language="text")
    elif msg_type == "reasoning":
        st.markdown("**Reasoning**")
        if msg.get("summary"):
            st.write(msg["summary"])
        elif msg.get("encrypted_content"):
            st.text("(encrypted content)")
    else:
        st.write(msg)


def render_response(resp: dict[str, Any]) -> None:
    st.subheader("Response")
    message = resp.get("message", {})
    if message:
        st.markdown(f"**role:** {message.get('role')}")
        for chunk in message.get("content", []):
            if isinstance(chunk, dict) and chunk.get("text"):
                st.write(chunk["text"])
    tool_calls = resp.get("tool_calls") or []
    for tc in tool_calls:
        with st.expander(f"Tool call: {tc.get('function', {}).get('name')}"):
            st.code(json.dumps(tc, indent=2), language="json")


def main() -> None:
    st.set_page_config(page_title="OpenHands Log Viewer", layout="wide")
    st.title("OpenHands Log Viewer")

    root_options = [p for p in LOG_ROOTS if p.exists()]
    if not root_options:
        st.warning("No log roots found under output/*Agent/logs.")
        return

    selected_root = st.sidebar.selectbox(
        "Log root", root_options, format_func=lambda p: str(p)
    )
    runs = load_runs(selected_root)
    if not runs:
        st.info("No runs found.")
        return
    run = st.sidebar.selectbox("Run", runs, format_func=lambda p: p.name)
    log_files = load_log_files(run)
    if not log_files:
        st.info("No log files in this run.")
        return
    log_file = st.sidebar.selectbox("Log file", log_files, format_func=lambda p: p.name)

    data = load_json(log_file)
    if not data:
        return

    st.subheader("Metadata")
    cols = st.columns(4)
    cols[0].metric("Model", data.get("llm_path", ""))
    cols[1].metric("Latency (s)", f"{data.get('latency_sec', 0):.2f}")
    cols[2].metric("Cost", data.get("cost", ""))
    cols[3].metric("Timestamp", data.get("timestamp", ""))

    st.subheader("Input")
    for msg in data.get("input", []):
        msg_type = msg.get("type", msg.get("role", "message"))
        label = msg_type
        if msg_type == "function_call":
            name = msg.get("name", "")
            label = f"{msg_type} • {name}".strip()
        if msg_type == "function_call_output":
            label = msg_type
        with st.expander(label, expanded=False):
            render_message(msg)

    if data.get("response"):
        render_response(data["response"])

    if usage := data.get("usage_summary"):
        with st.expander("Usage summary"):
            st.json(usage)

    st.caption(f"File: {log_file}")


if __name__ == "__main__":
    main()