app.py

from __future__ import annotations

import os
import sys
from pathlib import Path
from typing import Any

import gradio as gr

ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT / "src"))

from id.config import Config, ProviderConfig, load_config  # noqa: E402
from id.engine.loop import Session  # noqa: E402
from id.generator.archive import Archive  # noqa: E402
from id.llm import local as local_llm  # noqa: E402  (registers @spaces.GPU)
from id.llm.client import LLMClient, LLMError  # noqa: E402
from id.llm.prompts import PromptRegistry  # noqa: E402

# In local (no-API) mode, fetch the GGUF on CPU at startup so the first turn
# does not spend its GPU budget downloading ~2.5 GB.
if os.getenv("F_ID_PROVIDER", "local") == "local":
    try:
        local_llm.prefetch()
    except Exception:  # never block app startup on a prefetch hiccup
        pass


def build_config() -> Config:
    cfg = load_config(ROOT / "config.toml")
    # Default: fully self-contained, runs a small model in-process (no API).
    provider = os.getenv("F_ID_PROVIDER", "local")
    model = os.getenv("F_ID_MODEL", "openbmb/MiniCPM4.1-8B")
    base_url = os.getenv("F_ID_BASE_URL")
    api_key_env = os.getenv("F_ID_API_KEY_ENV")

    if provider == "openrouter":
        base_url = base_url or "https://openrouter.ai/api/v1"
        api_key_env = api_key_env or "OPENROUTER_API_KEY"
    elif provider == "local":
        base_url = base_url or cfg.providers["local"].base_url
        api_key_env = api_key_env or "LOCAL_API_KEY"
    else:
        base_url = base_url or "https://api.openai.com/v1"
        api_key_env = api_key_env or "OPENAI_API_KEY"

    cfg.providers[provider] = ProviderConfig(
        base_url=base_url,
        api_key_env=api_key_env,
        default_headers=cfg.providers.get(
            provider,
            ProviderConfig(
                base_url=base_url,
                api_key_env=api_key_env,
            ),
        ).default_headers,
    )
    for tier in cfg.tiers.values():
        tier.provider = provider
        tier.model = model
    cfg.engine.request_timeout = float(os.getenv("F_ID_REQUEST_TIMEOUT", "180"))
    cfg.root = ROOT
    cfg.runtime_dir.mkdir(exist_ok=True)
    return cfg


def make_context() -> tuple[Config, PromptRegistry, LLMClient]:
    cfg = build_config()
    return cfg, PromptRegistry(cfg.prompts_dir), LLMClient(cfg)


def world_entries() -> list[dict[str, Any]]:
    cfg, _, _ = make_context()
    return Archive(cfg.worlds_dir).entries()


def world_choice(entry: dict[str, Any]) -> str:
    title = entry.get("title") or entry.get("world_id", "")
    one_line = entry.get("one_line", "")
    return f"{title} | {one_line}" if one_line else title


WORLD_LABELS = {world_choice(entry): entry["world_id"] for entry in world_entries()}
DEFAULT_WORLD = next(iter(WORLD_LABELS), "")


def session_from_id(session_id: str) -> Session:
    if not session_id:
        raise gr.Error("Start a case first.")
    cfg, prompts, client = make_context()
    return Session.resume(cfg, session_id, prompts, client)


def people(session: Session) -> list[str]:
    return [c.name for c in session.world.characters.values() if c.role != "victim"]


def locations(session: Session) -> list[str]:
    return sorted({s.location for s in session.world.timeline.slices})


def intro_markdown(session: Session) -> str:
    world = session.world
    names = ", ".join(f"{c.name} ({c.role})" for c in world.characters.values())
    body = world.world_md.strip() or world.meta.one_line
    return (
        f"### {world.meta.title or world.meta.world_id}\n\n"
        f"{body[:1400]}\n\n"
        f"**People:** {names}\n\n"
        f"**Session:** `{session.state.session_id}`"
    )


def notes_markdown(session: Session) -> str:
    notes = session.notes()
    lines: list[str] = []
    if notes.discovered_clues:
        lines.append("### Clues")
        lines.extend(f"- `{c['id']}`: {c['reveals']}" for c in notes.discovered_clues)
    else:
        lines.append("_No clues discovered yet._")
    if notes.cracked:
        lines.append("\n### Cracked")
        lines.append(", ".join(notes.cracked))
    if notes.ledgers:
        lines.append("\n### Statements")
        for name, claims in notes.ledgers.items():
            lines.append(f"\n**{name}**")
            lines.extend(f"- `{c['claim_id']}`: {c['proposition']}" for c in claims)
    else:
        lines.append("\n_No statements on record yet._")
    return "\n".join(lines)


def status_markdown(session: Session) -> str:
    return (
        f"**Turn:** {session.state.turn}  \n"
        f"**Status:** {session.state.status.value}  \n"
        f"**Clues:** {len(session.state.discovered_clues)}"
    )


def start_case(label: str) -> tuple[str, list[dict[str, str]], str, str, Any, Any, Any, Any]:
    if not label:
        raise gr.Error("No archived worlds are available.")
    world_id = WORLD_LABELS[label]
    cfg, prompts, client = make_context()
    session = Session.start(cfg, world_id, prompts, client)
    chat = [{"role": "assistant", "content": intro_markdown(session)}]
    chars = people(session)
    locs = locations(session)
    return (
        session.state.session_id,
        chat,
        notes_markdown(session),
        status_markdown(session),
        gr.update(choices=chars, value=chars[0] if chars else None),
        gr.update(choices=chars, value=chars[0] if chars else None),
        gr.update(choices=locs, value=locs[0] if locs else None),
        gr.update(choices=chars, value=chars[0] if chars else None),
    )


def talk(session_id: str, character: str, message: str, chat: list[dict[str, str]]):
    if not message.strip():
        raise gr.Error("Enter a question.")
    try:
        session = session_from_id(session_id)
        outcome = session.talk(character, message.strip())
    except (KeyError, LLMError) as exc:
        raise gr.Error(str(exc)) from exc
    suffix = ""
    if outcome.discovered_clues:
        suffix = "\n\n" + "\n".join(
            f"**Clue discovered:** `{cid}`" for cid in outcome.discovered_clues
        )
    label = f"Talk to {character}"
    chat = chat + [
        {"role": "user", "content": f"**{label}:** {message.strip()}"},
        {"role": "assistant", "content": outcome.text + suffix},
    ]
    return chat, "", notes_markdown(session), status_markdown(session)


def look(session_id: str, location: str | None, query: str, chat: list[dict[str, str]]):
    if not query.strip():
        raise gr.Error("Enter something to inspect.")
    try:
        session = session_from_id(session_id)
        answer = session.look(query.strip(), location or None)
    except LLMError as exc:
        raise gr.Error(str(exc)) from exc
    suffix = (
        f"\n\n**Clue discovered:** `{answer.discovered_clue}`" if answer.discovered_clue else ""
    )
    where = f"@{location}" if location else "the scene"
    chat = chat + [
        {"role": "user", "content": f"**Look {where}:** {query.strip()}"},
        {"role": "assistant", "content": answer.text + suffix},
    ]
    return chat, "", notes_markdown(session), status_markdown(session)


def confront(
    session_id: str,
    character: str,
    claim_a: str,
    claim_b: str,
    chat: list[dict[str, str]],
):
    if not claim_a.strip() or not claim_b.strip():
        raise gr.Error("Enter two statement ids from Notes.")
    try:
        session = session_from_id(session_id)
        result = session.confront(character, claim_a.strip(), claim_b.strip())
    except KeyError as exc:
        raise gr.Error(str(exc)) from exc
    label = "Verified contradiction" if result.verified else "No contradiction"
    chat = chat + [
        {"role": "user", "content": f"**Confront {character}:** `{claim_a}` vs `{claim_b}`"},
        {"role": "assistant", "content": f"**{label}.**\n\n{result.reason}"},
    ]
    return chat, notes_markdown(session), status_markdown(session)


def accuse(
    session_id: str,
    culprit: str,
    means: str,
    motive: str,
    opportunity: str,
    chat: list[dict[str, str]],
):
    if not culprit:
        raise gr.Error("Choose a culprit.")
    try:
        session = session_from_id(session_id)
        result = session.accuse(culprit, means.strip(), motive.strip(), opportunity.strip())
    except KeyError as exc:
        raise gr.Error(str(exc)) from exc
    chat = chat + [
        {"role": "user", "content": f"**Accuse:** {culprit}"},
        {"role": "assistant", "content": f"**{result.grade}**\n\n{result.debrief}"},
    ]
    return chat, notes_markdown(session), status_markdown(session)


# Paper palette. We override Gradio's design tokens for BOTH the default and
# `.dark` scopes so the light aesthetic renders identically no matter what the
# visitor's browser/system theme is (the `?__theme=light` redirect is not
# reliable inside the HF Spaces iframe).
_TOKENS = """
  --body-background-fill: #f4f2ee;
  --background-fill-primary: #fdfcfb;
  --background-fill-secondary: #efece7;
  --block-background-fill: #efece7;
  --block-label-background-fill: #efece7;
  --block-title-background-fill: transparent;
  --panel-background-fill: #fdfcfb;
  --stat-background-fill: #efece7;
  --input-background-fill: #ffffff;
  --input-background-fill-focus: #ffffff;
  --input-background-fill-hover: #ffffff;
  --code-background-fill: #efece7;
  --table-even-background-fill: #fdfcfb;
  --table-odd-background-fill: #f4f2ee;

  --body-text-color: #1a1714;
  --body-text-color-subdued: #5f5a53;
  --block-label-text-color: #5f5a53;
  --block-title-text-color: #1a1714;
  --block-info-text-color: #5f5a53;
  --input-placeholder-color: #9a958d;
  --link-text-color: #2f2b27;
  --link-text-color-hover: #000000;

  --border-color-primary: #d8d5cf;
  --border-color-accent: #d8d5cf;
  --border-color-accent-subdued: #e6e3de;
  --block-border-color: #d8d5cf;
  --input-border-color: #d8d5cf;
  --input-border-color-focus: #8c887f;
  --panel-border-color: #d8d5cf;
  --table-border-color: #d8d5cf;

  --color-accent: #2f2b27;
  --color-accent-soft: #efece7;

  --button-primary-background-fill: #2f2b27;
  --button-primary-background-fill-hover: #000000;
  --button-primary-text-color: #ffffff;
  --button-primary-border-color: #2f2b27;
  --button-secondary-background-fill: #efece7;
  --button-secondary-background-fill-hover: #e2dfd9;
  --button-secondary-text-color: #1a1714;
  --button-secondary-border-color: #d8d5cf;
"""

CSS = (
    ":root, html, body, body.dark, .dark, gradio-app, "
    ".gradio-container, .gradio-container.dark {"
    + _TOKENS
    + "}\n"
    + """
:root {
  --font-heading: "EB Garamond", Georgia, serif;
}

/* Force the outermost surfaces light too (vars cascade down, not up). */
html, body, body.dark, gradio-app, .gradio-container, .gradio-container.dark {
  background: var(--body-background-fill) !important;
  color: var(--body-text-color) !important;
}

.gradio-container { max-width: 1120px !important; margin: 0 auto !important; }

h1, h2, h3, h4 { font-family: var(--font-heading) !important; font-weight: 600 !important; }

/* Header */
.app-title { padding: 20px 4px 12px; border-bottom: 1px solid var(--border-color-primary); margin-bottom: 4px; }
.app-title h1 { font-size: 34px; line-height: 1.1; margin: 0; letter-spacing: 0.5px; color: var(--body-text-color); }
.app-title p { color: var(--body-text-color-subdued); margin-top: 6px; font-size: 15px; }

/* Side panels (status + notes) */
.side-panel { background: var(--block-background-fill); border: 1px solid var(--border-color-primary); border-radius: 10px; padding: 14px 16px; }
.side-panel code { background: var(--background-fill-secondary); color: var(--color-accent); padding: 1px 5px; border-radius: 4px; }
"""
)

THEME = gr.themes.Base(
    primary_hue="neutral",
    secondary_hue="stone",
    neutral_hue="stone",
    font=[gr.themes.GoogleFont("IBM Plex Sans"), "system-ui", "sans-serif"],
    font_mono=[gr.themes.GoogleFont("IBM Plex Mono"), "monospace"],
).set(
    body_background_fill="#f4f2ee",
    body_text_color="#1a1714",
    block_background_fill="#fdfcfb",
    block_label_text_color="#5f5a53",
    border_color_primary="#d8d5cf",
    input_background_fill="#ffffff",
    button_primary_background_fill="#2f2b27",
    button_primary_background_fill_hover="#000000",
    button_primary_text_color="#ffffff",
)

with gr.Blocks(title="ID") as demo:
    gr.HTML(
        """
        <div class="app-title">
          <h1>ID</h1>
          <p>An LLM-driven investigation game built as a small Gradio Space.</p>
        </div>
        """
    )
    session_state = gr.State("")

    with gr.Row():
        with gr.Column(scale=2):
            world = gr.Dropdown(
                choices=list(WORLD_LABELS),
                value=DEFAULT_WORLD,
                label="Case",
                interactive=True,
            )
        with gr.Column(scale=0, min_width=140):
            start = gr.Button("Start", variant="primary")

    with gr.Row(equal_height=False):
        with gr.Column(scale=3):
            chat = gr.Chatbot(label="Transcript", height=560, resizable=True)
            with gr.Tab("Talk"):
                character = gr.Dropdown(label="Character", choices=[])
                message = gr.Textbox(label="Question", lines=3)
                talk_btn = gr.Button("Ask", variant="primary")
            with gr.Tab("Look"):
                location = gr.Dropdown(label="Location", choices=[])
                query = gr.Textbox(label="Inspect", lines=2)
                look_btn = gr.Button("Look", variant="primary")
            with gr.Tab("Confront"):
                confront_character = gr.Dropdown(label="Character", choices=[])
                with gr.Row():
                    claim_a = gr.Textbox(label="Statement A")
                    claim_b = gr.Textbox(label="Statement B")
                confront_btn = gr.Button("Confront", variant="primary")
            with gr.Tab("Accuse"):
                accuse_character = gr.Dropdown(label="Culprit", choices=[])
                means = gr.Textbox(label="Means")
                motive = gr.Textbox(label="Motive")
                opportunity = gr.Textbox(label="Opportunity")
                accuse_btn = gr.Button("Accuse", variant="primary")
        with gr.Column(scale=2):
            status = gr.Markdown(label="Status", elem_classes="side-panel")
            notes = gr.Markdown(label="Notes", elem_classes="side-panel")

    start.click(
        start_case,
        inputs=[world],
        outputs=[
            session_state,
            chat,
            notes,
            status,
            character,
            confront_character,
            location,
            accuse_character,
        ],
    )
    talk_btn.click(
        talk,
        inputs=[session_state, character, message, chat],
        outputs=[chat, message, notes, status],
    )
    look_btn.click(
        look,
        inputs=[session_state, location, query, chat],
        outputs=[chat, query, notes, status],
    )
    confront_btn.click(
        confront,
        inputs=[session_state, confront_character, claim_a, claim_b, chat],
        outputs=[chat, notes, status],
    )
    accuse_btn.click(
        accuse,
        inputs=[session_state, accuse_character, means, motive, opportunity, chat],
        outputs=[chat, notes, status],
    )
    demo.load(
        start_case,
        inputs=[world],
        outputs=[
            session_state,
            chat,
            notes,
            status,
            character,
            confront_character,
            location,
            accuse_character,
        ],
    )

if __name__ == "__main__":
    demo.launch(theme=THEME, css=CSS)