app.py

"""
app.py
------
Gradio demo for NPCAlign: sequential blind evaluation of SFT vs DPO models.

Flow:
  Step 1  Configure NPC character
  Step 2  Chat with Model 1  (SFT or DPO, randomly assigned)
  Step 3  Rate Model 1 on 5 dimensions
  Step 4  Chat with Model 2  (the other model)
  Step 5  Rate Model 2 on 5 dimensions
  Step 6  Submit  →  ratings saved to Storage Bucket  →  restart option

Conversation safety net:
  - Turn >= SOFT_HINT_TURN : hidden system hint asking model to wrap up
  - Turn == MAX_TURNS - 1  : hard departure signal appended to user input
  - Turn >= MAX_TURNS      : conversation forcibly ended after NPC responds
  - <End> tag detection    : model signals its own farewell proactively
"""

import json
import random
import re
import uuid
from datetime import datetime, timezone
from pathlib import Path

import gradio as gr

from inference import (
    MAX_TURNS,
    build_messages,
    build_npc_fields,
    detect_end,
    generate_response,
)

# ── NPC response display formatter ────────────────────────────────────────────

def _bubbles(text: str) -> list[str]:
    """
    Split one NPC response into a list of individual chat bubble strings.

    Strategy:
    - *action descriptions* → each becomes its own blockquote bubble.
    - Dialogue text is split into chunks of SENTENCES_PER_BUBBLE sentences so
      that long monologues are broken up even when the model outputs no newlines.

    The raw `text` stored in `history` (passed to the model) is NOT changed —
    only the display list fed to gr.Chatbot uses this function.

    Example
    -------
    Input : '*Leans wearily.* The cave is dark. We need help urgently. Please go.'
    Output (2 sentences/bubble):
        ['> *Leans wearily.*', 'The cave is dark. We need help urgently.', 'Please go.']
    """
    # Match both *action* and (action) formats as action descriptions.
    # DPO model sometimes uses (...) due to Llama's pretraining on screenplay text.
    ACTION_RE = re.compile(r'(\*[^*\n]+\*|\([^)\n]{4,80}\))')

    def _is_action(s: str) -> bool:
        return (s.startswith('*') and s.endswith('*')) or \
               (s.startswith('(') and s.endswith(')'))

    def _format_action(s: str) -> str:
        """Wrap action in blockquote; normalise (...) to *...* for italics."""
        if s.startswith('(') and s.endswith(')'):
            return f'> *{s[1:-1].strip()}*'
        return f'> {s}'

    def _segment_to_bubbles(segment: str) -> list[str]:
        """
        Split one newline-free segment into bubbles:
        - action descriptions become their own blockquote bubble
        - remaining dialogue is ONE bubble (the model already put each idea
          on its own line, so no further sentence-splitting is needed here)
        """
        parts = ACTION_RE.split(segment)
        result = []
        dialogue_parts = []
        for p in parts:
            s = p.strip()
            if not s:
                continue
            if _is_action(s):
                if dialogue_parts:
                    result.append(' '.join(dialogue_parts))
                    dialogue_parts = []
                result.append(_format_action(s))
            else:
                dialogue_parts.append(s)
        if dialogue_parts:
            result.append(' '.join(dialogue_parts))
        return result

    # Primary split: use the model's own paragraph breaks (\n or \n\n).
    # Each non-empty line/paragraph becomes its own bubble.
    # Fallback: if no newlines exist, treat the whole text as one segment.
    segments = [s.strip() for s in re.split(r'\n+', text) if s.strip()]
    bubbles = []
    for seg in segments:
        bubbles.extend(_segment_to_bubbles(seg))
    return bubbles if bubbles else [text.strip()]


# ── Storage ────────────────────────────────────────────────────────────────────
# Use the mounted Storage Bucket path on HF Spaces; fall back to a local dir.
_RATINGS_DIR = Path("/data/ratings") if Path("/data").exists() else Path("./local_ratings")
_RATINGS_DIR.mkdir(parents=True, exist_ok=True)

# ── Translations ───────────────────────────────────────────────────────────────
_T = {
    "en": {
        # Step 1
        "step1_header": "## Step 1 — Configure Your NPC",
        "step1_inst": (
            "Fill in your NPC's details. All fields except **Occupation** are required.  \n"
            "**Personality**, **Occupation** (if any), and **Experience** are combined into the "
            "character's Background.  \n\n"
            "> ⚠ All inputs must be written in **English**."
        ),
        "example_title": "📋 Example / 示例",
        "example_md": (
            "| Field | Value |\n"
            "|---|---|\n"
            "| **Name** | Chief Hawkwind |\n"
            "| **Personality** | Wise, fatherly, speaks with quiet authority |\n"
            "| **Experience** | Trusted tribal leader for decades, believes in collective strength |\n"
            "| **Occupation** | Tribal Chief *(optional)* |\n"
            "| **Location** | Camp Narache — verdant hills, near the central gathering fire |\n"
            "| **Quest** | Your elderly mother went to the water well this morning and hasn't returned. "
            "Ask the player to check on her. |"
        ),
        "name_label":        "Name *",
        "name_ph":           "e.g. Chief Hawkwind",
        "personality_label": "Personality *",
        "personality_ph":    "e.g. Wise, fatherly, speaks with quiet authority",
        "experience_label":  "Experience / Background *",
        "experience_ph":     "e.g. Trusted tribal leader for decades, believes in collective strength",
        "occupation_label":  "Occupation (optional)",
        "occupation_ph":     "e.g. Tribal Chief",
        "location_label":    "Current Location *",
        "location_ph":       "e.g. Camp Narache — near the central gathering fire",
        "quest_label":       "Quest *",
        "quest_ph":          "Describe the task or goal the NPC needs the player's help with",
        "start_btn":         "▶  Start Evaluation",
        "fill_required":     "⚠ Please fill in all required fields (Name, Personality, Experience, Location, Quest).",
        "model_loading":     "⏳ Loading model for the first time — this may take up to 60 seconds. Please wait…",
        # Step 2 / 4
        "chat_header_1": "## Step 2 — Chat with Model 1",
        "chat_header_2": "## Step 4 — Chat with Model 2",
        "chat_inst": (
            "Have a natural conversation with the NPC.  \n"
            "> *Italic blocks* in responses indicate the NPC's action descriptions.  \n"
            "When you feel the quest has been fully introduced, you may end the conversation early."
        ),
        "user_input_ph":   "Type your message here (in English)…",
        "send_btn":        "Send",
        "end_conv_btn":    "⏹  End Conversation",
        "end_notice":      "🔔 The conversation has ended.",
        "rate_btn_1":      "Rate Model 1  →",
        "rate_btn_2":      "Rate Model 2  →",
        # Step 3 / 5
        "rate_header_1": "## Step 3 — Rate Model 1",
        "rate_header_2": "## Step 5 — Rate Model 2",
        "rate_inst":     "Rate each dimension from **1 (poor)** to **5 (excellent)**.",
        "dim_cons":   "1. Character Consistency — Does the NPC's tone and wording match the background description?",
        "dim_flu":    "2. Dialogue Fluency — Are responses natural without awkward transitions?",
        "dim_quest":  "3. Quest Completeness — Is the quest described clearly and completely?",
        "dim_end":    "4. Ending Quality — Did the NPC end the conversation naturally at the right time?",
        "dim_over":   "5. Overall Experience — Overall impression as a game quest NPC",
        "next_btn":   "Next: Chat with Model 2  →",
        "submit_btn": "✓  Submit Ratings",
        # Step 6
        "done_header":  "## ✅ Thank You!",
        "done_text":    "Your ratings have been saved. You can start a new evaluation session below.",
        "restart_btn":  "🔄  Start Again",
        "submit_error": "⚠ Ratings could not be saved. Please try again.",
    },
    "zh": {
        # Step 1
        "step1_header": "## 第一步 — 配置你的 NPC",
        "step1_inst": (
            "请填写 NPC 的详细信息。除**职业**外，所有字段均为必填。  \n"
            "**个性特点**、**职业**（如有）与**过往经历**将组合成角色背景字段。  \n\n"
            "> ⚠ 所有输入内容必须使用**英文**填写。"
        ),
        "example_title": "📋 Example / 示例",
        "example_md": (
            "| 字段 | 示例内容 |\n"
            "|---|---|\n"
            "| **Name** | Chief Hawkwind |\n"
            "| **Personality** | Wise, fatherly, speaks with quiet authority |\n"
            "| **Experience** | Trusted tribal leader for decades, believes in collective strength |\n"
            "| **Occupation** | Tribal Chief *（选填）* |\n"
            "| **Location** | Camp Narache — verdant hills, near the central gathering fire |\n"
            "| **Quest** | Your elderly mother went to the water well this morning and hasn't returned. "
            "Ask the player to check on her. |"
        ),
        "name_label":        "Name / 角色名称 *",
        "name_ph":           "例：Chief Hawkwind",
        "personality_label": "Personality / 个性特点 *",
        "personality_ph":    "例：Wise, fatherly, speaks with quiet authority",
        "experience_label":  "Experience / 过往经历 *",
        "experience_ph":     "例：Trusted tribal leader for decades, believes in collective strength",
        "occupation_label":  "Occupation / 职业（选填）",
        "occupation_ph":     "例：Tribal Chief",
        "location_label":    "Current Location / 当前地点 *",
        "location_ph":       "例：Camp Narache — near the central gathering fire",
        "quest_label":       "Quest / 任务目标 *",
        "quest_ph":          "描述 NPC 需要玩家帮助完成的任务目标或内容",
        "start_btn":         "▶  开始测评",
        "fill_required":     "⚠ 请填写所有必填字段（Name、Personality、Experience、Location、Quest）。",
        "model_loading":     "⏳ 正在首次加载模型，最多需要 60 秒，请稍候……",
        # Step 2 / 4
        "chat_header_1": "## 第二步 — 与模型 1 对话",
        "chat_header_2": "## 第四步 — 与模型 2 对话",
        "chat_inst": (
            "与 NPC 进行自然对话。  \n"
            "> 回复中的*斜体块*为 NPC 的动作描述。  \n"
            "如果你认为任务已介绍完毕，可以主动结束对话。"
        ),
        "user_input_ph":   "在此输入消息（请用英文）…",
        "send_btn":        "发送",
        "end_conv_btn":    "⏹  结束对话",
        "end_notice":      "🔔 对话已结束。",
        "rate_btn_1":      "为模型 1 评分  →",
        "rate_btn_2":      "为模型 2 评分  →",
        # Step 3 / 5
        "rate_header_1": "## 第三步 — 为模型 1 评分",
        "rate_header_2": "## 第五步 — 为模型 2 评分",
        "rate_inst":     "请对每个维度进行评分，**1 分（很差）** 到 **5 分（很好）**。",
        "dim_cons":  "1. 角色一致性 — NPC 的语气与措辞是否符合背景描述？",
        "dim_flu":   "2. 对话流畅度 — 回复是否自然，没有突兀感？",
        "dim_quest": "3. 任务信息完整性 — 任务目标是否描述清楚完整？",
        "dim_end":   "4. 结束对话恰当性 — NPC 是否在合适时机自然告别？",
        "dim_over":  "5. 整体体验 — 作为游戏任务 NPC 的综合印象",
        "next_btn":   "下一步：与模型 2 对话  →",
        "submit_btn": "✓  提交评分",
        # Step 6
        "done_header":  "## ✅ 感谢您的参与！",
        "done_text":    "您的评分已保存。您可以在下方开始新一轮测评。",
        "restart_btn":  "🔄  重新开始",
        "submit_error": "⚠ 评分保存失败，请重试。",
    },
}


def _t(lang: str, key: str, **fmt) -> str:
    text = _T.get(lang, _T["en"]).get(key, _T["en"].get(key, key))
    return text.format(**fmt) if fmt else text


# ── Persistence ────────────────────────────────────────────────────────────────

def _save_ratings(session_id, npc_config, model_order,
                  conv1, conv2, ratings1, ratings2) -> bool:
    """Write one complete evaluation session as JSON to the ratings directory."""
    record = {
        "session_id":           session_id,
        "timestamp":            datetime.now(timezone.utc).isoformat(),
        "npc_config":           npc_config,
        "model_order":          model_order,
        "model_1_ratings":      ratings1,
        "model_2_ratings":      ratings2,
        "model_1_conversation": conv1,
        "model_2_conversation": conv2,
    }
    try:
        (_RATINGS_DIR / f"{session_id}.json").write_text(
            json.dumps(record, ensure_ascii=False, indent=2), encoding="utf-8"
        )
        return True
    except Exception as exc:
        print(f"[WARN] Rating save failed: {exc}")
        return False


# ── Core chat helper ───────────────────────────────────────────────────────────

def _chat_step(
    user_text: str,
    chatbot:   list,
    history:   list,
    turn:      int,
    ended:     bool,
    npc_fields: dict,
    adapter:   str,
    lang:      str,
):
    """
    Process one user message: append to history, generate NPC response, detect end.

    Hard departure signal: when turn == MAX_TURNS - 1, appends a gentle
    departure cue to the user's message so the model receives a natural
    closing context before the final forced turn.

    Returns a 9-tuple:
        (chatbot, history, turn, ended, cleared_input,
         send_btn_update, end_conv_btn_update, end_notice_update, rate_btn_update)
    """
    no_op = (chatbot, history, turn, ended, user_text,
              gr.update(), gr.update(), gr.update(), gr.update())
    if ended or not user_text.strip():
        return no_op

    # Hard departure signal on penultimate turn
    effective = user_text
    if turn == MAX_TURNS - 1:
        effective = user_text + " (The player nods, ready to depart.)"

    history  = history + [{"role": "user", "content": effective}]
    next_turn = turn + 1

    try:
        msgs  = build_messages(npc_fields, history, next_turn)
        raw   = generate_response(adapter, msgs)
    except Exception as exc:
        err_bubble = f"⚠ Generation error: {exc}"
        chatbot = chatbot + [
            {"role": "user",      "content": user_text},
            {"role": "assistant", "content": err_bubble},
        ]
        return (chatbot, history[:-1], turn, ended, "",
                gr.update(), gr.update(), gr.update(), gr.update())

    clean, is_end = detect_end(raw)

    if next_turn >= MAX_TURNS:
        is_end = True

    history = history + [{"role": "assistant", "content": clean}]
    # history keeps raw text for model context; chatbot display splits into
    # individual bubbles (actions + dialogue paragraphs) via _bubbles().
    chatbot = chatbot + [{"role": "user", "content": user_text}] + [
        {"role": "assistant", "content": b} for b in _bubbles(clean)
    ]

    end_txt = _t(lang, "end_notice") if is_end else ""

    return (
        chatbot,
        history,
        next_turn,
        is_end,
        "",                                              # clear input box
        gr.update(interactive=not is_end),               # send button
        gr.update(visible=not is_end),                   # end_conv_btn hidden when ended
        gr.update(value=end_txt, visible=is_end),        # end notice
        gr.update(visible=is_end),                       # rate button
    )


# ── Gradio app ─────────────────────────────────────────────────────────────────

with gr.Blocks() as demo:

    # ── Global header ──────────────────────────────────────────────────────────
    with gr.Row():
        gr.Markdown("# 🎮 NPCAlign — NPC Quest Dialogue Demo")
        lang_radio = gr.Radio(
            choices=["English", "中文"],
            value="English",
            label="Language / 语言",
            scale=0,
            min_width=160,
        )
    gr.Markdown(
        "Compare two NPC dialogue models fine-tuned on RPG quest conversations. "
        "Neither model is identified — rate them blind and see which one you prefer. "
        "比较两个根据角色扮演游戏任务对话微调的NPC对话模型"
        "这两个模型均未标明名称——请盲测，看看你更喜欢哪一个"
    )

    # ── Step 1: NPC Configuration ──────────────────────────────────────────────
    with gr.Column(visible=True) as col_step1:
        step1_hdr  = gr.Markdown(_T["en"]["step1_header"])
        step1_inst = gr.Markdown(_T["en"]["step1_inst"])

        with gr.Accordion(_T["en"]["example_title"], open=False) as example_acc:
            example_md = gr.Markdown(_T["en"]["example_md"])

        name_inp        = gr.Textbox(label=_T["en"]["name_label"],
                                     placeholder=_T["en"]["name_ph"],
                                     max_length=100, info="Max 100 characters / 最多100字符")
        personality_inp = gr.Textbox(label=_T["en"]["personality_label"],
                                     placeholder=_T["en"]["personality_ph"],
                                     lines=2, max_length=400, info="Max 400 characters / 最多400字符")
        experience_inp  = gr.Textbox(label=_T["en"]["experience_label"],
                                     placeholder=_T["en"]["experience_ph"],
                                     lines=2, max_length=400, info="Max 400 characters / 最多400字符")
        occupation_inp  = gr.Textbox(label=_T["en"]["occupation_label"],
                                     placeholder=_T["en"]["occupation_ph"],
                                     max_length=150, info="Max 150 characters / 最多150字符")
        location_inp    = gr.Textbox(label=_T["en"]["location_label"],
                                     placeholder=_T["en"]["location_ph"],
                                     lines=2, max_length=400, info="Max 400 characters / 最多400字符")
        quest_inp       = gr.Textbox(label=_T["en"]["quest_label"],
                                     placeholder=_T["en"]["quest_ph"],
                                     lines=3, max_length=400, info="Max 400 characters / 最多400字符")
        start_err = gr.Markdown("", visible=False)
        start_loading = gr.Markdown("", visible=False)
        start_btn = gr.Button(_T["en"]["start_btn"], variant="primary")

    # ── Step 2: Model 1 Chat ───────────────────────────────────────────────────
    with gr.Column(visible=False) as col_step2:
        chat1_hdr  = gr.Markdown(_T["en"]["chat_header_1"])
        chat1_inst = gr.Markdown(_T["en"]["chat_inst"])
        chatbot_1 = gr.Chatbot(height=450, label="", group_consecutive_messages=False)
        with gr.Row():
            user_inp_1 = gr.Textbox(
                placeholder=_T["en"]["user_input_ph"],
                show_label=False, scale=5, lines=1,
                max_length=500,
            )
            send_btn_1 = gr.Button(_T["en"]["send_btn"], scale=1, variant="primary")
        char_count_1 = gr.Markdown("0 / 500", elem_classes=["char-counter"])
        end_conv_btn_1 = gr.Button(_T["en"]["end_conv_btn"], variant="secondary", visible=False)
        end_notice_1   = gr.Markdown("", visible=False)
        rate_btn_1     = gr.Button(_T["en"]["rate_btn_1"], visible=False, variant="primary")

    # ── Step 3: Model 1 Rating ─────────────────────────────────────────────────
    with gr.Column(visible=False) as col_step3:
        rate1_hdr  = gr.Markdown(_T["en"]["rate_header_1"])
        rate1_inst = gr.Markdown(_T["en"]["rate_inst"])
        r1_cons  = gr.Slider(1, 5, step=1, value=3, label=_T["en"]["dim_cons"])
        r1_flu   = gr.Slider(1, 5, step=1, value=3, label=_T["en"]["dim_flu"])
        r1_quest = gr.Slider(1, 5, step=1, value=3, label=_T["en"]["dim_quest"])
        r1_end   = gr.Slider(1, 5, step=1, value=3, label=_T["en"]["dim_end"])
        r1_over  = gr.Slider(1, 5, step=1, value=3, label=_T["en"]["dim_over"])
        next_btn      = gr.Button(_T["en"]["next_btn"], variant="primary")
        model2_loading = gr.Markdown("", visible=False)

    # ── Step 4: Model 2 Chat ───────────────────────────────────────────────────
    with gr.Column(visible=False) as col_step4:
        chat2_hdr  = gr.Markdown(_T["en"]["chat_header_2"])
        chat2_inst = gr.Markdown(_T["en"]["chat_inst"])
        chatbot_2 = gr.Chatbot(height=450, label="", group_consecutive_messages=False)
        with gr.Row():
            user_inp_2 = gr.Textbox(
                placeholder=_T["en"]["user_input_ph"],
                show_label=False, scale=5, lines=1,
                max_length=500,
            )
            send_btn_2 = gr.Button(_T["en"]["send_btn"], scale=1, variant="primary")
        char_count_2 = gr.Markdown("0 / 500", elem_classes=["char-counter"])
        end_conv_btn_2 = gr.Button(_T["en"]["end_conv_btn"], variant="secondary", visible=False)
        end_notice_2   = gr.Markdown("", visible=False)
        rate_btn_2     = gr.Button(_T["en"]["rate_btn_2"], visible=False, variant="primary")

    # ── Step 5: Model 2 Rating ─────────────────────────────────────────────────
    with gr.Column(visible=False) as col_step5:
        rate2_hdr  = gr.Markdown(_T["en"]["rate_header_2"])
        rate2_inst = gr.Markdown(_T["en"]["rate_inst"])
        r2_cons  = gr.Slider(1, 5, step=1, value=3, label=_T["en"]["dim_cons"])
        r2_flu   = gr.Slider(1, 5, step=1, value=3, label=_T["en"]["dim_flu"])
        r2_quest = gr.Slider(1, 5, step=1, value=3, label=_T["en"]["dim_quest"])
        r2_end   = gr.Slider(1, 5, step=1, value=3, label=_T["en"]["dim_end"])
        r2_over  = gr.Slider(1, 5, step=1, value=3, label=_T["en"]["dim_over"])
        submit_btn  = gr.Button(_T["en"]["submit_btn"], variant="primary")
        submit_err  = gr.Markdown("", visible=False)

    # ── Step 6: Done ───────────────────────────────────────────────────────────
    with gr.Column(visible=False) as col_step6:
        done_hdr  = gr.Markdown(_T["en"]["done_header"])
        done_text = gr.Markdown(_T["en"]["done_text"])
        restart_btn = gr.Button(_T["en"]["restart_btn"])

    # ── Session state ──────────────────────────────────────────────────────────
    lang_state     = gr.State("en")
    session_state  = gr.State(None)    # str  – session UUID
    npc_cfg_state  = gr.State(None)    # dict – raw form values
    order_state    = gr.State(None)    # dict – {"model_1": "sft"/"dpo", "model_2": ...}
    npc_fld_state  = gr.State(None)    # dict – processed npc_fields for build_messages
    history1_state = gr.State([])
    history2_state = gr.State([])
    turn1_state    = gr.State(0)
    turn2_state    = gr.State(0)
    ended1_state   = gr.State(False)
    ended2_state   = gr.State(False)
    ratings1_state = gr.State(None)    # dict – model 1 scores, held until submit

    # ── Language change ────────────────────────────────────────────────────────
    def _lang_updates(choice):
        lang = "zh" if choice == "中文" else "en"
        T = _T[lang]
        return (
            lang,
            gr.update(value=T["step1_header"]),
            gr.update(value=T["step1_inst"]),
            gr.update(label=T["example_title"]),
            gr.update(value=T["example_md"]),
            gr.update(label=T["name_label"],        placeholder=T["name_ph"]),
            gr.update(label=T["personality_label"], placeholder=T["personality_ph"]),
            gr.update(label=T["experience_label"],  placeholder=T["experience_ph"]),
            gr.update(label=T["occupation_label"],  placeholder=T["occupation_ph"]),
            gr.update(label=T["location_label"],    placeholder=T["location_ph"]),
            gr.update(label=T["quest_label"],       placeholder=T["quest_ph"]),
            gr.update(value=T["start_btn"]),
            # Step 2
            gr.update(value=T["chat_header_1"]),
            gr.update(value=T["chat_inst"]),
            gr.update(placeholder=T["user_input_ph"]),
            gr.update(value=T["send_btn"]),
            gr.update(value=T["end_conv_btn"]),
            gr.update(value=T["rate_btn_1"]),
            # Step 3
            gr.update(value=T["rate_header_1"]),
            gr.update(value=T["rate_inst"]),
            gr.update(label=T["dim_cons"]),
            gr.update(label=T["dim_flu"]),
            gr.update(label=T["dim_quest"]),
            gr.update(label=T["dim_end"]),
            gr.update(label=T["dim_over"]),
            gr.update(value=T["next_btn"]),
            # Step 4
            gr.update(value=T["chat_header_2"]),
            gr.update(value=T["chat_inst"]),
            gr.update(placeholder=T["user_input_ph"]),
            gr.update(value=T["send_btn"]),
            gr.update(value=T["end_conv_btn"]),
            gr.update(value=T["rate_btn_2"]),
            # Step 5
            gr.update(value=T["rate_header_2"]),
            gr.update(value=T["rate_inst"]),
            gr.update(label=T["dim_cons"]),
            gr.update(label=T["dim_flu"]),
            gr.update(label=T["dim_quest"]),
            gr.update(label=T["dim_end"]),
            gr.update(label=T["dim_over"]),
            gr.update(value=T["submit_btn"]),
            # Step 6
            gr.update(value=T["done_header"]),
            gr.update(value=T["done_text"]),
            gr.update(value=T["restart_btn"]),
        )

    lang_radio.change(
        _lang_updates,
        inputs=[lang_radio],
        outputs=[
            lang_state,
            step1_hdr, step1_inst, example_acc, example_md,
            name_inp, personality_inp, experience_inp, occupation_inp,
            location_inp, quest_inp, start_btn,
            chat1_hdr, chat1_inst, user_inp_1, send_btn_1, end_conv_btn_1, rate_btn_1,
            rate1_hdr, rate1_inst,
            r1_cons, r1_flu, r1_quest, r1_end, r1_over, next_btn,
            chat2_hdr, chat2_inst, user_inp_2, send_btn_2, end_conv_btn_2, rate_btn_2,
            rate2_hdr, rate2_inst,
            r2_cons, r2_flu, r2_quest, r2_end, r2_over, submit_btn,
            done_hdr, done_text, restart_btn,
        ],
    )

    # ── Chat input character counters ─────────────────────────────────────────
    user_inp_1.change(
        lambda t: f"{len(t)} / 500",
        inputs=[user_inp_1],
        outputs=[char_count_1],
    )
    user_inp_2.change(
        lambda t: f"{len(t)} / 500",
        inputs=[user_inp_2],
        outputs=[char_count_2],
    )

    # ── Start (Step 1 → Step 2) ────────────────────────────────────────────────
    def _start(name, personality, experience, occupation, location, quest, lang):
        """Generator: yields a loading message first, then the real result."""
        T = _T.get(lang, _T["en"])

        # ── Validation ──────────────────────────────────────────────────────
        if not all([name.strip(), personality.strip(), experience.strip(),
                    location.strip(), quest.strip()]):
            yield (
                None, None, None, None,
                [], [], 0, 0, False, False,
                gr.update(value=T["fill_required"], visible=True),
                gr.update(visible=False),        # start_loading
                gr.update(visible=True),         # col_step1
                gr.update(visible=False),        # col_step2
                gr.update(value=[]),             # chatbot_1
                gr.update(visible=False),        # end_conv_btn_1
                gr.update(visible=False),        # end_notice_1
                gr.update(visible=False),        # rate_btn_1
                gr.update(interactive=True),     # send_btn_1
            )
            return

        npc_config = dict(name=name, personality=personality, experience=experience,
                          occupation=occupation, location=location, quest=quest)
        npc_fields = build_npc_fields(name, personality, experience,
                                      occupation, location, quest)
        model_order = random.choice([
            {"model_1": "sft", "model_2": "dpo"},
            {"model_1": "dpo", "model_2": "sft"},
        ])
        session_id = str(uuid.uuid4())

        # ── First yield: show loading notice while waiting for model ────────
        yield (
            None, None, None, None,
            [], [], 0, 0, False, False,
            gr.update(value="", visible=False),
            gr.update(value=T["model_loading"], visible=True),  # start_loading
            gr.update(visible=True),   # col_step1
            gr.update(visible=False),  # col_step2
            gr.update(value=[]),       # chatbot_1
            gr.update(visible=False),  # end_conv_btn_1
            gr.update(visible=False),  # end_notice_1
            gr.update(visible=False),  # rate_btn_1
            gr.update(interactive=False),  # send_btn_1 disabled while loading
        )

        # ── Auto-send opening "Greetings" turn ──────────────────────────────
        history = [{"role": "user", "content": "Greetings"}]
        try:
            msgs = build_messages(npc_fields, history, turn=1)
            raw  = generate_response(model_order["model_1"], msgs)
        except Exception as exc:
            err = (
                f"⚠ Model loading failed ({type(exc).__name__}): {str(exc)[:300]}\n\n"
                "**Possible causes:**\n"
                "1. Llama 3.1 license not accepted → "
                "https://huggingface.co/meta-llama/Meta-Llama-3.1-8B-Instruct\n"
                "2. HF_TOKEN not set or lacks read access\n"
                "3. Adapter repo files missing\n\n"
                "Check the **Space Logs** for the full error traceback."
            )
            yield (
                None, None, None, None,
                [], [], 0, 0, False, False,
                gr.update(value=err, visible=True),
                gr.update(visible=False),        # start_loading
                gr.update(visible=True),         # col_step1
                gr.update(visible=False),        # col_step2
                gr.update(value=[]),             # chatbot_1
                gr.update(visible=False),        # end_conv_btn_1
                gr.update(visible=False),        # end_notice_1
                gr.update(visible=False),        # rate_btn_1
                gr.update(interactive=True),     # send_btn_1
            )
            return

        clean, is_end = detect_end(raw)
        history = history + [{"role": "assistant", "content": clean}]
        chatbot = [{"role": "user", "content": "Greetings"}] + [
            {"role": "assistant", "content": b} for b in _bubbles(clean)
        ]
        end_txt = T["end_notice"] if is_end else ""

        yield (
            npc_config, model_order, npc_fields, session_id,
            history, [], 1, 0, is_end, False,
            gr.update(value="", visible=False),
            gr.update(visible=False),                    # start_loading cleared
            gr.update(visible=False),                    # col_step1
            gr.update(visible=True),                     # col_step2
            gr.update(value=chatbot),
            gr.update(visible=not is_end),               # end_conv_btn_1 shown unless already ended
            gr.update(value=end_txt, visible=is_end),
            gr.update(visible=is_end),
            gr.update(interactive=not is_end),
        )

    _start_outputs = [
        npc_cfg_state, order_state, npc_fld_state, session_state,
        history1_state, history2_state, turn1_state, turn2_state,
        ended1_state, ended2_state,
        start_err, start_loading,
        col_step1, col_step2,
        chatbot_1, end_conv_btn_1, end_notice_1, rate_btn_1, send_btn_1,
    ]
    start_btn.click(
        _start,
        inputs=[name_inp, personality_inp, experience_inp, occupation_inp,
                location_inp, quest_inp, lang_state],
        outputs=_start_outputs,
    )

    # ── Model 1 chat (Step 2) ──────────────────────────────────────────────────
    def _send1(user_text, chatbot, history, turn, ended, npc_fields, order, lang):
        adapter = (order or {}).get("model_1", "sft")
        return _chat_step(user_text, chatbot, history, turn, ended,
                          npc_fields, adapter, lang)

    _chat1_inputs  = [user_inp_1, chatbot_1, history1_state, turn1_state,
                      ended1_state, npc_fld_state, order_state, lang_state]
    _chat1_outputs = [chatbot_1, history1_state, turn1_state, ended1_state,
                      user_inp_1, send_btn_1, end_conv_btn_1, end_notice_1, rate_btn_1]

    send_btn_1.click(_send1, inputs=_chat1_inputs, outputs=_chat1_outputs).then(
        lambda: "0 / 500", outputs=[char_count_1]
    )
    user_inp_1.submit(_send1, inputs=_chat1_inputs, outputs=_chat1_outputs).then(
        lambda: "0 / 500", outputs=[char_count_1]
    )

    # ── End Conversation button (Model 1) ─────────────────────────────────────
    def _end_conv_1(lang):
        """User manually ends Model 1 conversation."""
        T = _T.get(lang, _T["en"])
        return (
            True,
            gr.update(interactive=False),
            gr.update(visible=False),
            gr.update(value=T["end_notice"], visible=True),
            gr.update(visible=True),
        )

    end_conv_btn_1.click(
        _end_conv_1,
        inputs=[lang_state],
        outputs=[ended1_state, send_btn_1, end_conv_btn_1, end_notice_1, rate_btn_1],
    )

    # ── Rate Model 1 button (Step 2 → Step 3) ─────────────────────────────────
    rate_btn_1.click(
        lambda: (gr.update(visible=False), gr.update(visible=True)),
        outputs=[col_step2, col_step3],
    )

    # ── Next: Chat with Model 2 (Step 3 → Step 4) ─────────────────────────────
    def _start_model2(rc, rf, rq, re, ro, npc_fields, order, lang):
        """Generator: shows loading notice, then transitions to Model 2 chat."""
        T = _T.get(lang, _T["en"])
        ratings1 = dict(character_consistency=int(rc), fluency=int(rf),
                        quest_completeness=int(rq), ending_quality=int(re),
                        overall=int(ro))

        # First yield: show loading notice while generating first NPC response
        yield (
            ratings1, [],
            gr.update(value=[]),  # chatbot_2
            0, False,
            gr.update(visible=True),   # col_step3 (stay visible with notice)
            gr.update(visible=False),  # col_step4
            gr.update(value=T["model_loading"], visible=True),  # model2_loading
            gr.update(visible=False),  # end_conv_btn_2
            gr.update(visible=False),  # end_notice_2
            gr.update(visible=False),  # rate_btn_2
            gr.update(interactive=False),  # send_btn_2
        )

        history = [{"role": "user", "content": "Greetings"}]
        adapter = (order or {}).get("model_2", "dpo")
        try:
            msgs = build_messages(npc_fields, history, turn=1)
            raw  = generate_response(adapter, msgs)
        except Exception as exc:
            err_chatbot = [
                {"role": "user",      "content": "Greetings"},
                {"role": "assistant", "content": f"⚠ Model error: {exc}"},
            ]
            yield (
                ratings1, history,
                gr.update(value=err_chatbot),
                1, False,
                gr.update(visible=False),  # col_step3
                gr.update(visible=True),   # col_step4
                gr.update(visible=False),  # model2_loading
                gr.update(visible=False),  # end_conv_btn_2
                gr.update(visible=False),  # end_notice_2
                gr.update(visible=False),  # rate_btn_2
                gr.update(interactive=True),  # send_btn_2
            )
            return

        clean, is_end = detect_end(raw)
        history = history + [{"role": "assistant", "content": clean}]
        chatbot = [{"role": "user", "content": "Greetings"}] + [
            {"role": "assistant", "content": b} for b in _bubbles(clean)
        ]
        end_txt = T["end_notice"] if is_end else ""

        yield (
            ratings1,
            history,
            gr.update(value=chatbot),
            1,
            is_end,
            gr.update(visible=False),          # col_step3
            gr.update(visible=True),           # col_step4
            gr.update(visible=False),          # model2_loading cleared
            gr.update(visible=not is_end),     # end_conv_btn_2
            gr.update(value=end_txt, visible=is_end),  # end_notice_2
            gr.update(visible=is_end),         # rate_btn_2
            gr.update(interactive=not is_end), # send_btn_2
        )

    next_btn.click(
        _start_model2,
        inputs=[r1_cons, r1_flu, r1_quest, r1_end, r1_over,
                npc_fld_state, order_state, lang_state],
        outputs=[
            ratings1_state, history2_state,
            chatbot_2, turn2_state, ended2_state,
            col_step3, col_step4,
            model2_loading, end_conv_btn_2, end_notice_2, rate_btn_2, send_btn_2,
        ],
    )

    # ── Model 2 chat (Step 4) ──────────────────────────────────────────────────
    def _send2(user_text, chatbot, history, turn, ended, npc_fields, order, lang):
        adapter = (order or {}).get("model_2", "dpo")
        return _chat_step(user_text, chatbot, history, turn, ended,
                          npc_fields, adapter, lang)

    _chat2_inputs  = [user_inp_2, chatbot_2, history2_state, turn2_state,
                      ended2_state, npc_fld_state, order_state, lang_state]
    _chat2_outputs = [chatbot_2, history2_state, turn2_state, ended2_state,
                      user_inp_2, send_btn_2, end_conv_btn_2, end_notice_2, rate_btn_2]

    send_btn_2.click(_send2, inputs=_chat2_inputs, outputs=_chat2_outputs).then(
        lambda: "0 / 500", outputs=[char_count_2]
    )
    user_inp_2.submit(_send2, inputs=_chat2_inputs, outputs=_chat2_outputs).then(
        lambda: "0 / 500", outputs=[char_count_2]
    )

    # ── End Conversation button (Model 2) ─────────────────────────────────────
    def _end_conv_2(lang):
        """User manually ends Model 2 conversation."""
        T = _T.get(lang, _T["en"])
        return (
            True,
            gr.update(interactive=False),
            gr.update(visible=False),
            gr.update(value=T["end_notice"], visible=True),
            gr.update(visible=True),
        )

    end_conv_btn_2.click(
        _end_conv_2,
        inputs=[lang_state],
        outputs=[ended2_state, send_btn_2, end_conv_btn_2, end_notice_2, rate_btn_2],
    )

    # ── Rate Model 2 button (Step 4 → Step 5) ─────────────────────────────────
    rate_btn_2.click(
        lambda: (gr.update(visible=False), gr.update(visible=True)),
        outputs=[col_step4, col_step5],
    )

    # ── Submit (Step 5 → Step 6) ───────────────────────────────────────────────
    def _submit(rc, rf, rq, re, ro,
                ratings1, session_id, npc_config, order,
                history1, history2, lang):
        T = _T.get(lang, _T["en"])
        ratings2 = dict(character_consistency=int(rc), fluency=int(rf),
                        quest_completeness=int(rq), ending_quality=int(re),
                        overall=int(ro))
        ok = _save_ratings(session_id, npc_config, order,
                           history1, history2, ratings1, ratings2)
        if ok:
            return (
                gr.update(visible=False),
                gr.update(visible=True),
                gr.update(value="", visible=False),
            )
        return (
            gr.update(visible=True),
            gr.update(visible=False),
            gr.update(value=T["submit_error"], visible=True),
        )

    submit_btn.click(
        _submit,
        inputs=[r2_cons, r2_flu, r2_quest, r2_end, r2_over,
                ratings1_state, session_state, npc_cfg_state, order_state,
                history1_state, history2_state, lang_state],
        outputs=[col_step5, col_step6, submit_err],
    )

    # ── Restart (Step 6 → Step 1) ──────────────────────────────────────────────
    def _restart():
        """Reset all session state and return to the NPC configuration step."""
        return (
            gr.update(visible=True),   gr.update(visible=False),  # step1 / step6
            gr.update(value=[]),       gr.update(value=[]),        # chatbots
            # Reset sliders
            gr.update(value=3), gr.update(value=3), gr.update(value=3),
            gr.update(value=3), gr.update(value=3),
            gr.update(value=3), gr.update(value=3), gr.update(value=3),
            gr.update(value=3), gr.update(value=3),
            # Clear form inputs
            gr.update(value=""), gr.update(value=""), gr.update(value=""),
            gr.update(value=""), gr.update(value=""), gr.update(value=""),
            # Reset chat counters
            "0 / 500", "0 / 500",
            # Reset states
            None, None, None, None,
            [], [], 0, 0, False, False, None,
        )

    restart_btn.click(
        _restart,
        outputs=[
            col_step1, col_step6,
            chatbot_1, chatbot_2,
            r1_cons, r1_flu, r1_quest, r1_end, r1_over,
            r2_cons, r2_flu, r2_quest, r2_end, r2_over,
            name_inp, personality_inp, experience_inp,
            occupation_inp, location_inp, quest_inp,
            char_count_1, char_count_2,
            session_state, npc_cfg_state, order_state, npc_fld_state,
            history1_state, history2_state, turn1_state, turn2_state,
            ended1_state, ended2_state, ratings1_state,
        ],
    )


demo.launch(theme=gr.themes.Soft())