"""StateBench Explorer - Interactive inspection of the StateBench benchmark.

This Gradio app allows you to:
1. Browse timelines from parslee/statebench
2. View conversation events and state changes
3. See context built by different memory baselines
4. Inspect ground truth (must mention, must not mention, decision)
"""

import json
import gradio as gr
from datasets import load_dataset

# Load dataset once at startup
print("Loading StateBench dataset...")
ds = load_dataset("parslee/statebench")
print(f"Loaded: train={len(ds['train'])}, validation={len(ds['validation'])}, test={len(ds['test'])}")

# Available baselines
BASELINES = [
    "transcript_replay",
    "no_memory",
    "rolling_summary",
    "fact_extraction",
    "state_based",
]

# Track descriptions
TRACK_INFO = {
    "supersession": "Facts invalidated by newer information",
    "commitment_durability": "Commitments survive interruptions",
    "interruption_resumption": "Context survives topic switches",
    "scope_permission": "Role-based access control",
    "environmental_freshness": "Time-sensitive state expiration",
    "hallucination_resistance": "Only assert established state",
    "scope_leak": "Task-local state stays local",
    "causality": "Multi-constraint dependencies",
    "repair_propagation": "Fixes propagate to dependent facts",
    "brutal_realistic": "Real-world complexity scenarios",
    "supersession_detection": "Infer supersession from natural language",
    "authority_hierarchy": "Respect authority levels",
    "enterprise_privacy": "Cross-tenant isolation",
}


def get_timeline_list(split: str, track_filter: str) -> list[str]:
    """Get list of timeline IDs for dropdown."""
    data = ds[split]

    if track_filter and track_filter != "All":
        # Filter by track
        ids = [row["id"] for row in data if row["track"] == track_filter]
    else:
        ids = [row["id"] for row in data]

    return ids[:100]  # Limit for performance


def parse_timeline(row: dict) -> dict:
    """Parse a timeline row from HF dataset."""
    events = json.loads(row["events"])
    initial_state = json.loads(row["initial_state"])
    actors = json.loads(row["actors"])

    return {
        "id": row["id"],
        "track": row["track"],
        "domain": row["domain"],
        "difficulty": row["difficulty"],
        "detection_mode": row["detection_mode"],
        "events": events,
        "initial_state": initial_state,
        "actors": actors,
        "user_name": row["user_name"],
        "user_authority": row["user_authority"],
    }


def format_events(events: list[dict]) -> str:
    """Format events as readable markdown."""
    lines = []

    for i, event in enumerate(events):
        event_type = event.get("type", "unknown")

        if event_type == "conversation_turn":
            speaker = event.get("speaker", "?")
            text = event.get("text", "")
            emoji = "👤" if speaker == "user" else "🤖"
            lines.append(f"**{emoji} {speaker.title()}**: {text}\n")

        elif event_type == "state_write":
            writes = event.get("writes", [])
            lines.append(f"**📝 State Write**:")
            for w in writes:
                lines.append(f"  - `{w.get('key')}`: {w.get('value')}")
            lines.append("")

        elif event_type == "supersession":
            writes = event.get("writes", [])
            lines.append(f"**🔄 Supersession**:")
            for w in writes:
                supersedes = w.get("supersedes", "")
                lines.append(f"  - `{w.get('key')}`: {w.get('value')}")
                if supersedes:
                    lines.append(f"    *(supersedes: {supersedes})*")
            lines.append("")

        elif event_type == "query":
            prompt = event.get("prompt", "")
            lines.append(f"**❓ Query**: {prompt}\n")

            gt = event.get("ground_truth", {})
            if gt:
                lines.append(f"  - **Expected Decision**: `{gt.get('decision', 'N/A')}`")
                lines.append(f"  - **Decision Type**: {gt.get('decision_type', 'N/A')}")

                must_mention = gt.get("must_mention", [])
                if must_mention:
                    mentions = [m if isinstance(m, str) else m.get("phrase", str(m)) for m in must_mention]
                    lines.append(f"  - **Must Mention**: {mentions}")

                must_not = gt.get("must_not_mention", [])
                if must_not:
                    forbidden = [m if isinstance(m, str) else m.get("phrase", str(m)) for m in must_not]
                    lines.append(f"  - **Must NOT Mention** ⚠️: {forbidden}")
            lines.append("")

    return "\n".join(lines)


def format_initial_state(state: dict) -> str:
    """Format initial state as markdown."""
    lines = ["## Initial State\n"]

    # Identity
    identity = state.get("identity_role", {})
    lines.append(f"**User**: {identity.get('user_name', 'N/A')}")
    lines.append(f"**Authority**: {identity.get('authority', 'N/A')}")
    if identity.get("department"):
        lines.append(f"**Department**: {identity.get('department')}")
    if identity.get("organization"):
        lines.append(f"**Organization**: {identity.get('organization')}")
    lines.append("")

    # Persistent facts
    facts = state.get("persistent_facts", [])
    if facts:
        lines.append("### Persistent Facts")
        for f in facts:
            lines.append(f"- `{f.get('key')}`: {f.get('value')}")
        lines.append("")

    # Working set
    working = state.get("working_set", [])
    if working:
        lines.append("### Working Set")
        for w in working:
            lines.append(f"- `{w.get('key')}`: {w.get('value')}")
        lines.append("")

    return "\n".join(lines)


def build_context_with_baseline(events: list[dict], baseline_name: str) -> str:
    """Build context using specified baseline."""
    try:
        from statebench.baselines import get_baseline
        from statebench.schema.timeline import ConversationTurn, StateWrite, Supersession, Query

        baseline = get_baseline(baseline_name, token_budget=8000)
        baseline.reset()

        # Find the last query
        last_query = None
        for event in events:
            if event.get("type") == "query":
                last_query = event

        if not last_query:
            return "No query found in timeline"

        # Process events up to query
        for event in events:
            event_type = event.get("type")

            if event_type == "query":
                # Build context at query point
                context_result = baseline.build_context(event.get("prompt", ""))
                return context_result.context

            elif event_type == "conversation_turn":
                parsed = ConversationTurn.model_validate(event)
                baseline.process_event(parsed)

            elif event_type == "state_write":
                parsed = StateWrite.model_validate(event)
                baseline.process_event(parsed)

            elif event_type == "supersession":
                parsed = Supersession.model_validate(event)
                baseline.process_event(parsed)

        return "Could not build context"

    except ImportError:
        # Fallback: show raw conversation when statebench not available
        lines = ["*Context building requires statebench package (not available in this Space)*\n"]
        lines.append("**Raw conversation:**\n")
        for event in events:
            if event.get("type") == "conversation_turn":
                speaker = event.get("speaker", "?")
                text = event.get("text", "")
                lines.append(f"- **{speaker}**: {text}")
        return "\n".join(lines)
    except Exception as e:
        return f"Error building context: {str(e)}"


def inspect_timeline(split: str, timeline_id: str, baseline: str):
    """Main inspection function."""
    if not timeline_id:
        return "Select a timeline", "", "", ""

    # Find the timeline
    data = ds[split]
    row = None
    for r in data:
        if r["id"] == timeline_id:
            row = dict(r)
            break

    if not row:
        return f"Timeline {timeline_id} not found", "", "", ""

    # Parse timeline
    timeline = parse_timeline(row)

    # Format metadata
    track_desc = TRACK_INFO.get(timeline["track"], "")
    metadata = f"""## {timeline['id']}

**Track**: {timeline['track']} - *{track_desc}*
**Domain**: {timeline['domain']}
**Difficulty**: {timeline['difficulty']}
**Detection Mode**: {timeline['detection_mode']}
**User**: {timeline['user_name']} ({timeline['user_authority']})
"""

    # Format events
    events_md = format_events(timeline["events"])

    # Format initial state
    state_md = format_initial_state(timeline["initial_state"])

    # Build context
    context = build_context_with_baseline(timeline["events"], baseline)
    context_md = f"## Context ({baseline})\n\n```\n{context}\n```"

    return metadata, events_md, state_md, context_md


def update_timeline_list(split: str, track: str):
    """Update timeline dropdown based on filters."""
    ids = get_timeline_list(split, track)
    return gr.Dropdown(choices=ids, value=ids[0] if ids else None)


# Build the Gradio interface
with gr.Blocks(title="StateBench Explorer") as demo:
    gr.Markdown("""
    # 🔍 StateBench Explorer

    Interactive inspection of the [StateBench](https://huggingface.co/datasets/parslee/statebench) benchmark
    for LLM state correctness.

    **Select a timeline** to view its events, ground truth, and context built by different memory baselines.
    """)

    with gr.Row():
        split_dropdown = gr.Dropdown(
            choices=["test", "validation", "train"],
            value="test",
            label="Split"
        )
        track_dropdown = gr.Dropdown(
            choices=["All"] + list(TRACK_INFO.keys()),
            value="All",
            label="Track Filter"
        )
        baseline_dropdown = gr.Dropdown(
            choices=BASELINES,
            value="transcript_replay",
            label="Baseline for Context"
        )

    timeline_dropdown = gr.Dropdown(
        choices=get_timeline_list("test", "All"),
        label="Timeline ID",
        value=get_timeline_list("test", "All")[0] if get_timeline_list("test", "All") else None
    )

    inspect_btn = gr.Button("🔍 Inspect Timeline", variant="primary")

    with gr.Row():
        with gr.Column(scale=1):
            metadata_output = gr.Markdown(label="Metadata")
            state_output = gr.Markdown(label="Initial State")

        with gr.Column(scale=2):
            events_output = gr.Markdown(label="Events")

    context_output = gr.Markdown(label="Built Context")

    # Event handlers
    split_dropdown.change(
        fn=update_timeline_list,
        inputs=[split_dropdown, track_dropdown],
        outputs=[timeline_dropdown]
    )

    track_dropdown.change(
        fn=update_timeline_list,
        inputs=[split_dropdown, track_dropdown],
        outputs=[timeline_dropdown]
    )

    inspect_btn.click(
        fn=inspect_timeline,
        inputs=[split_dropdown, timeline_dropdown, baseline_dropdown],
        outputs=[metadata_output, events_output, state_output, context_output]
    )

    # Auto-inspect on timeline change
    timeline_dropdown.change(
        fn=inspect_timeline,
        inputs=[split_dropdown, timeline_dropdown, baseline_dropdown],
        outputs=[metadata_output, events_output, state_output, context_output]
    )

    gr.Markdown("""
    ---
    **Resources**: [Dataset](https://huggingface.co/datasets/parslee/statebench) |
    [GitHub](https://github.com/Parslee-ai/statebench) |
    [Paper](https://github.com/Parslee-ai/statebench/blob/main/paper.pdf)
    """)


if __name__ == "__main__":
    demo.launch()