app.py

#!/usr/bin/env python3
"""Gradio web UI for Turnabout — Ace Attorney-style agent evaluation environment."""

from __future__ import annotations

import json
import glob
import os
import re as _re
from pathlib import Path

import gradio as gr

from turnabout.envs.text_env import TextCourtEnv
from turnabout.i18n import t

CASES_DIR = Path(__file__).parent / "turnabout" / "cases"


def list_cases() -> list[str]:
    return sorted(glob.glob(str(CASES_DIR / "*.json")))


def load_case_titles() -> dict[str, str]:
    titles = {}
    for path in list_cases():
        try:
            with open(path) as f:
                data = json.load(f)
            titles[data.get("title", Path(path).stem)] = path
        except Exception:
            titles[Path(path).stem] = path
    return titles


class GameSession:
    def __init__(self):
        self.env: TextCourtEnv | None = None
        self.done = False
        self.lang = "en"

    def start(self, case_path: str, difficulty: str, lang: str) -> str:
        self.lang = lang
        self.env = TextCourtEnv(case_path=case_path, difficulty=difficulty, max_steps=200, lang=lang)
        obs = self.env.reset()
        self.done = False
        return obs

    def step(self, action: str) -> tuple[str, float, bool]:
        if self.env is None or self.done:
            return t("ui_no_game", self.lang), 0.0, True
        obs, reward, done, info = self.env.step(action)
        self.done = done
        return obs, reward, done

    def get_valid_action_choices(self) -> list[tuple[str, str]]:
        if self.env is None or self.done:
            return []
        valid = self.env.engine.get_valid_actions()
        return [(a.display(self.lang, self.env.case), str(a)) for a in valid]

    def get_evidence_list(self) -> str:
        if self.env is None:
            return t("ui_no_game", self.lang)
        state = self.env.engine.state
        if not state or not state.inventory:
            return t("ui_no_evidence", self.lang)
        lines = []
        for i, eid in enumerate(sorted(state.inventory), 1):
            ev = self.env.case.get_evidence(eid)
            if ev:
                desc = ev.detail if ev.detail else ev.description
                lines.append(f"**{i}. {ev.name}**\n   {desc}")
        return "\n\n".join(lines)

    def get_status(self) -> str:
        if self.env is None:
            return t("ui_no_game", self.lang)
        state = self.env.engine.state
        if not state:
            return t("ui_no_game", self.lang)
        phase_key = f"phase_{state.phase.value}"
        phase = t(phase_key, self.lang)
        penalties = f"{state.penalties}/{self.env.case.court.penalty_limit}"
        found = len(state.contradictions_found)
        required = len(self.env.case.court.win_condition.required_contradiction_ids)
        return (
            f"**{t('ui_phase', self.lang)}:** {phase}\n\n"
            f"**{t('ui_penalties', self.lang)}:** {penalties}\n\n"
            f"**{t('ui_contradictions', self.lang)}:** {found}/{required}\n\n"
            f"**{t('ui_steps', self.lang)}:** {state.step_count}"
        )

    def get_metrics_text(self) -> str:
        if self.env is None:
            return ""
        metrics = self.env.get_metrics()
        return (
            f"**{t('ui_result', self.lang)}:** {t('ui_won', self.lang) if metrics.won else t('ui_lost', self.lang)}\n\n"
            f"**{t('ui_contradiction_accuracy', self.lang)}:** {metrics.contradiction_accuracy:.1%}\n\n"
            f"**{t('ui_evidence_coverage', self.lang)}:** {metrics.evidence_coverage:.1%}\n\n"
            f"**{t('ui_steps', self.lang)}:** {metrics.total_steps}\n\n"
            f"**{t('ui_composite_score', self.lang)}:** {metrics.composite_score:.3f}"
        )


session = GameSession()


def _msg(role: str, content: str) -> dict:
    return {"role": role, "content": content}


def _lang_code(lang_label: str) -> str:
    return "zh" if lang_label == "中文" else "en"


def on_new_game(case_title: str, difficulty: str, lang_label: str) -> tuple:
    lang = _lang_code(lang_label)
    titles = load_case_titles()
    case_path = titles.get(case_title)
    if not case_path:
        cases = list_cases()
        if not cases:
            return [_msg("assistant", t("ui_no_cases", lang))], "", "", gr.update(choices=[], value=None), gr.update(interactive=True)
        case_path = cases[0]

    obs = session.start(case_path, difficulty, lang)
    history = [_msg("assistant", obs)]
    choices = session.get_valid_action_choices()
    return (
        history,
        session.get_evidence_list(),
        session.get_status(),
        gr.update(choices=choices, value=choices[0][1] if choices else None),
        gr.update(interactive=True),
    )


def on_execute(action_str: str, chat_history: list, lang_label: str) -> tuple:
    lang = _lang_code(lang_label)
    if not action_str:
        choices = session.get_valid_action_choices()
        return chat_history, session.get_evidence_list(), session.get_status(), gr.update(choices=choices, value=choices[0][1] if choices else None), gr.update(visible=False), ""

    display_name = action_str
    choices = session.get_valid_action_choices()
    for label, val in choices:
        if val == action_str:
            display_name = label
            break

    obs, reward, done = session.step(action_str)

    reward_text = f"\n\n`[{t('ui_reward', lang)}: {reward:+.2f}]`" if reward != 0 else ""
    response = obs + reward_text

    chat_history = chat_history + [
        _msg("user", display_name),
        _msg("assistant", response),
    ]

    if done:
        metrics_text = session.get_metrics_text()
        chat_history = chat_history + [_msg("assistant", f"---\n**{t('ui_game_over', lang)}**\n\n{metrics_text}")]
        return chat_history, session.get_evidence_list(), session.get_status(), gr.update(choices=[], value=None), gr.update(visible=True), metrics_text

    new_choices = session.get_valid_action_choices()
    return (
        chat_history,
        session.get_evidence_list(),
        session.get_status(),
        gr.update(choices=new_choices, value=new_choices[0][1] if new_choices else None),
        gr.update(visible=False),
        "",
    )


def on_generate_case(theme: str, gen_difficulty: str, progress=gr.Progress()):
    yield "Initializing case generator..."

    try:
        from dotenv import load_dotenv
        load_dotenv()
    except ImportError:
        pass

    if not (os.environ.get("GATEWAY_API_KEY") or os.environ.get("OPENAI_API_KEY")):
        yield (
            "Generation is not configured.\n\n"
            "Set `GATEWAY_API_KEY` for a gateway, or `OPENAI_API_KEY` for "
            "OpenAI. If using a custom gateway, also set `GATEWAY_URL` and "
            "`GATEWAY_MODELS`."
        )
        return

    try:
        from turnabout.generation.generator import CaseGenerator
    except ImportError as e:
        yield f"Import error: {e}"
        return

    progress(0.1, desc="Initializing generator...")

    try:
        generator = CaseGenerator(backend="openai")
    except Exception as e:
        yield f"Failed to create generator: {e}"
        return

    try:
        case = None
        for value, desc, maybe_case in generator.generate_events(
            theme=theme or None,
            difficulty=gen_difficulty,
        ):
            progress(value, desc=desc)
            if maybe_case is None:
                yield (
                    "**Generating case...**\n\n"
                    f"{desc}\n\n"
                    "LLM case generation can take a few minutes."
                )
            else:
                case = maybe_case
    except Exception as e:
        yield f"Generation failed: {e}"
        return

    if case is None:
        yield "Generation failed: no case was produced."
        return

    progress(0.9, desc="Saving case...")

    slug = _re.sub(r"[^a-z0-9]+", "_", case.title.lower()).strip("_")
    out_path = CASES_DIR / f"{slug}.json"
    with open(out_path, "w") as f:
        json.dump(case.model_dump(), f, indent=2, ensure_ascii=False)

    yield f"Case generated successfully!\n\n**{case.title}**\n{case.description}\n\nSaved to: `{out_path.name}`\n\nRefresh the page to see the new case in the Play section."


COMMANDS_HELP_EN = """
### Investigation Commands
- `move <location>` — Travel to a location
- `examine <object>` — Examine an object
- `talk <person>` — Talk to a character
- `present <evidence> to <person>` — Show evidence to someone
- `go to court` — Proceed to court

### Court Commands
- `press` — Press the witness on current statement
- `present <evidence>` — Present evidence against current statement
- `next` / `prev` — Navigate between statements
"""

COMMANDS_HELP_ZH = """
### 调查阶段命令
- `move <地点>` — 前往某地点
- `examine <物品>` — 调查物品
- `talk <人物>` — 与人物对话
- `present <证据> to <人物>` — 向人物出示证据
- `go to court` — 前往法庭

### 法庭命令
- `press` — 追问当前证言
- `present <证据>` — 出示证据反驳当前证言
- `next` / `prev` — 切换证言
"""


def build_app():
    case_titles = load_case_titles()
    case_choices = list(case_titles.keys()) if case_titles else ["No cases found"]

    with gr.Blocks(title="Turnabout Bench") as app:
        gr.Markdown(t("ui_title", "en"), elem_id="app_title")

        lang_radio = gr.Radio(
            choices=["English", "中文"],
            value="English",
            label="Language / 语言",
            interactive=True,
        )

        gr.Markdown("## Play")
        with gr.Row():
            with gr.Column(scale=3):
                chatbot = gr.Chatbot(
                    label=t("ui_game_label", "en"),
                    height=480,
                )
                action_radio = gr.Radio(
                    choices=[],
                    label=t("ui_select_action", "en"),
                    interactive=True,
                )
                execute_btn = gr.Button(t("ui_execute", "en"), variant="primary")

            with gr.Column(scale=1):
                with gr.Group():
                    case_dropdown = gr.Dropdown(
                        choices=case_choices,
                        value=case_choices[0] if case_choices else None,
                        label=t("ui_case", "en"),
                    )
                    difficulty_radio = gr.Radio(
                        choices=["easy", "hard"],
                        value="easy",
                        label=t("ui_difficulty", "en"),
                    )
                    new_game_btn = gr.Button(t("ui_new_game", "en"), variant="primary")

                status_md = gr.Markdown(
                    t("ui_click_new_game", "en"),
                    label=t("ui_status", "en"),
                )

                evidence_md = gr.Markdown(
                    "",
                    label=t("ui_evidence", "en"),
                )

                metrics_box = gr.Markdown("", visible=False, label="Final Metrics")

                with gr.Accordion(t("ui_commands_ref", "en"), open=False):
                    gr.Markdown(COMMANDS_HELP_EN)

        new_game_btn.click(
            fn=on_new_game,
            inputs=[case_dropdown, difficulty_radio, lang_radio],
            outputs=[chatbot, evidence_md, status_md, action_radio, execute_btn],
        )

        execute_btn.click(
            fn=on_execute,
            inputs=[action_radio, chatbot, lang_radio],
            outputs=[chatbot, evidence_md, status_md, action_radio, metrics_box, metrics_box],
        )

        gr.Markdown(
            "---\n"
            "## Generate Case\n"
            "Use the configured LLM gateway to generate new cases automatically."
        )
        with gr.Row():
            with gr.Column():
                theme_input = gr.Textbox(
                    label=t("ui_theme", "en"),
                    placeholder="e.g. 'art heist', 'corporate espionage', 'poisoned cake'",
                )
                gen_difficulty = gr.Radio(
                    choices=["easy", "hard"],
                    value="easy",
                    label=t("ui_difficulty", "en"),
                )
                generate_btn = gr.Button(t("ui_generate_btn", "en"), variant="primary")
            with gr.Column():
                gen_output = gr.Markdown(label=t("ui_result", "en"))

        generate_btn.click(
            fn=on_generate_case,
            inputs=[theme_input, gen_difficulty],
            outputs=[gen_output],
        )

        gr.Markdown("""
---
## About

Turnabout Bench is an interactive evaluation environment for testing agent long-horizon reasoning,
inspired by the **Ace Attorney** game series.

**Key Features:**
- **Investigation Phase** (hard mode): Explore locations, collect evidence, talk to witnesses
- **Court Phase**: Cross-examine witnesses, find contradictions in testimony
- **Dual Interface**: Text-based for LLM agents, Gymnasium-compatible for RL agents
- **Evaluation Metrics**: Contradiction accuracy, evidence coverage, step efficiency
- **LLM Case Generation**: Automatically generate new cases via API

**Usage:**
```python
# LLM Agent
from turnabout.envs.text_env import TextCourtEnv
env = TextCourtEnv(case_path="...", difficulty="easy")
obs = env.reset()
obs, reward, done, info = env.step("present evidence_name")

# RL Agent (Gymnasium)
import gymnasium as gym
import turnabout.envs
env = gym.make("Turnabout-Discrete-v0", case_path="...")
```
""")

    return app.queue()


if __name__ == "__main__":
    app = build_app()
    app.launch(share=False)