from html import escape
import os
from time import sleep

import gradio as gr

from puppet_theater import (
    DEFAULT_OPENBMB_MODEL_ID,
    TheaterSession,
    create_show_from_premise,
    get_backend_status,
    request_finale,
    run_one_beat,
    summon_actor,
    throw_prop,
    warm_up_openbmb,
)


EMPTY_STAGE = """
<div class="puppet-stage stage-empty">
  <div class="stage-valance"></div>
  <div class="stage-backdrop">
    <div class="stage-marquee">AI Puppet Theater</div>
    <div class="empty-stage-copy">Enter a premise and raise the curtain.</div>
  </div>
  <div class="stage-floorboards"></div>
</div>
"""

EMPTY_TRANSCRIPT = "No show yet. The transcript will appear here."
EMPTY_DIRECTOR_LOG = "No director notes yet."
EMPTY_TRACE = "No trace events yet."
EMPTY_BACKEND = (
    "Active backend: deterministic\n"
    "OpenBMB model id: openbmb/MiniCPM5-1B\n"
    "Model status: unloaded\n"
    "Fallback: deterministic safety path enabled"
)
BACKEND_CHOICES = ["deterministic", "openbmb"]
OPENBMB_MODEL_ID = os.getenv("OPENBMB_MODEL_ID", DEFAULT_OPENBMB_MODEL_ID)
DEFAULT_MAX_NEW_TOKENS = 80
DEFAULT_TEMPERATURE = 0.8
PLAYBACK_DELAY_SECONDS = 0.75
PROP_EMOJI = {
    "rubber duck": "🐤",
    "duck": "🐤",
    "egg": "🥚",
    "flowers": "💐",
    "flower": "💐",
    "tomato": "🍅",
    "crown": "👑",
    "tiny crown": "👑",
    "scroll": "📜",
    "banana": "🍌",
    "mirror": "🪞",
}

CUSTOM_CSS = """
body,
.gradio-container {
    background:
        radial-gradient(circle at 50% 0%, rgba(127, 29, 29, 0.18), transparent 28rem),
        linear-gradient(180deg, #0b1020 0%, #070914 100%) !important;
    color: #f8efe4 !important;
}
.gradio-container {
    max-width: 1180px !important;
    padding-top: 1rem !important;
    font-family: Inter, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
}
.gradio-container .prose,
.gradio-container label,
.gradio-container span,
.gradio-container p {
    color: #f8efe4;
}
.gradio-container textarea,
.gradio-container input {
    background: rgba(10, 12, 23, 0.82) !important;
    border-color: rgba(246, 196, 83, 0.24) !important;
    color: #f8efe4 !important;
}
.gradio-container textarea::placeholder,
.gradio-container input::placeholder {
    color: #9f8c7a !important;
}
.gradio-container footer {
    color: rgba(203, 183, 161, 0.62) !important;
}
.gradio-container .block,
.gradio-container .form,
.gradio-container .panel,
.gradio-container .tabs,
.gradio-container .tabitem {
    background: rgba(34, 17, 31, 0.56) !important;
    border-color: rgba(246, 196, 83, 0.18) !important;
}
.gradio-container label,
.gradio-container .block-title,
.gradio-container .label-wrap {
    color: #f8efe4 !important;
}
.gradio-container .block-info,
.gradio-container .label-wrap span,
.gradio-container label > span {
    background: rgba(34, 17, 31, 0.88) !important;
    border: 1px solid rgba(246, 196, 83, 0.28) !important;
    border-radius: 6px !important;
    color: #ffd166 !important;
    font-weight: 700 !important;
}
.gradio-container .wrap,
.gradio-container .styler,
.gradio-container .form,
.gradio-container .form > *,
.gradio-container .block > div {
    background-color: transparent !important;
}
.gradio-container select,
.gradio-container [role="listbox"],
.gradio-container [role="combobox"] {
    background: rgba(10, 12, 23, 0.82) !important;
    border-color: rgba(246, 196, 83, 0.24) !important;
    color: #f8efe4 !important;
}
.app-title h1 {
    color: #f8efe4;
    font-family: Georgia, "Times New Roman", serif;
    font-size: 2.15rem;
    letter-spacing: 0;
    margin-bottom: 0;
    text-align: center;
}
.app-title p {
    color: #cbb7a1;
    font-size: 0.95rem;
    margin: 0.15rem 0 0.8rem;
    text-align: center;
}
.gradio-container h3,
.gradio-container h3 span,
.gradio-container .prose h3,
.gradio-container .prose h3 span {
    color: #f8efe4 !important;
}
.premise-panel {
    background: rgba(42, 20, 38, 0.72);
    border-color: rgba(246, 196, 83, 0.3);
    box-shadow: 0 16px 32px rgba(0, 0, 0, 0.2);
    padding: 0.55rem 0.65rem 0.65rem;
}
.premise-panel .block,
.premise-panel .wrap,
.premise-panel .styler,
.premise-panel .form,
.premise-panel .block > div {
    background: rgba(42, 20, 38, 0.78) !important;
}
.control-panel {
    background: rgba(34, 17, 31, 0.76);
    border: 1px solid rgba(246, 196, 83, 0.22);
    border-radius: 8px;
    box-shadow: 0 14px 34px rgba(0, 0, 0, 0.22);
    padding: 0.55rem;
}
.control-panel .block,
.control-panel .wrap,
.control-panel .styler,
.control-panel .form,
.control-panel .block > div {
    background: rgba(34, 17, 31, 0.78) !important;
}
.control-panel .row,
.premise-panel .row {
    background: transparent !important;
}
.control-panel h3 {
    color: #f8efe4;
    margin: 0 0 0.35rem;
    font-size: 1rem;
}
.control-panel .prose,
.control-panel .prose h3,
.control-panel h3 * {
    color: #f8efe4 !important;
}
.puppet-stage {
    min-height: 430px;
    border: 5px solid #3b0a16;
    border-radius: 14px;
    background:
        linear-gradient(90deg, rgba(59, 10, 22, 0.98) 0 10%, transparent 10% 90%, rgba(59, 10, 22, 0.98) 90% 100%),
        linear-gradient(180deg, rgba(42, 20, 38, 0.96), rgba(13, 6, 14, 0.98));
    color: #f8efe4;
    display: flex;
    flex-direction: column;
    align-items: stretch;
    justify-content: stretch;
    position: relative;
    overflow: hidden;
    box-shadow:
        0 24px 48px rgba(0, 0, 0, 0.38),
        inset 0 0 42px rgba(0, 0, 0, 0.58);
}
.puppet-stage::before,
.puppet-stage::after {
    content: "";
    position: absolute;
    top: 0;
    bottom: 0;
    width: 13%;
    background:
        repeating-linear-gradient(90deg, rgba(255, 255, 255, 0.04) 0 14px, transparent 14px 28px),
        linear-gradient(180deg, #8b1e3f 0%, #7f1d1d 54%, #3b0a16 100%);
    box-shadow: inset -16px 0 28px rgba(0, 0, 0, 0.22);
    z-index: 2;
}
.puppet-stage::before {
    left: 0;
}
.puppet-stage::after {
    right: 0;
    transform: scaleX(-1);
}
.stage-valance {
    height: 48px;
    background:
        repeating-linear-gradient(90deg, rgba(255, 255, 255, 0.06) 0 22px, transparent 22px 44px),
        linear-gradient(180deg, #8b1e3f 0%, #7f1d1d 100%);
    border-bottom: 4px solid #f6c453;
    box-shadow: 0 10px 20px rgba(0, 0, 0, 0.34);
    position: relative;
    z-index: 3;
}
.stage-backdrop {
    background:
        radial-gradient(circle at 50% 8%, rgba(255, 224, 150, 0.28), transparent 19rem),
        radial-gradient(circle at 24% 58%, rgba(255, 224, 150, 0.12), transparent 14rem),
        linear-gradient(180deg, #2a1426 0%, #22111f 62%, #130911 100%);
    flex: 1;
    padding: 0.72rem 7.2rem 0.8rem;
    position: relative;
    z-index: 1;
}
.stage-backdrop::after {
    background: linear-gradient(180deg, transparent 0%, rgba(124, 63, 23, 0.46) 100%);
    bottom: 0;
    content: "";
    height: 32%;
    left: 0;
    position: absolute;
    right: 0;
}
.stage-marquee {
    color: #fff7ed;
    font-family: Georgia, "Times New Roman", serif;
    font-size: 1.6rem;
    font-weight: 700;
    letter-spacing: 0;
    text-align: center;
    text-shadow: 0 4px 18px rgba(0, 0, 0, 0.72);
    position: relative;
    z-index: 2;
    overflow-wrap: anywhere;
}
.stage-copy {
    max-width: 54rem;
    color: #cbb7a1;
    font-size: 0.84rem;
    line-height: 1.35;
    margin: 0.25rem auto 0;
    text-align: center;
    position: relative;
    z-index: 2;
}
.stage-copy strong {
    color: #f8efe4;
}
.empty-stage-copy {
    color: #cbb7a1;
    font-size: 1rem;
    margin-top: 5.8rem;
    text-align: center;
    position: relative;
    z-index: 2;
}
.stage-floorboards {
    height: 58px;
    background:
        repeating-linear-gradient(90deg, rgba(255, 255, 255, 0.08) 0 2px, transparent 2px 72px),
        linear-gradient(180deg, #8a4b22 0%, #7c3f17 100%);
    border-top: 2px solid rgba(246, 196, 83, 0.28);
    position: relative;
    z-index: 3;
}
.speech-bubble {
    animation: bubble-in 0.24s ease-out;
    background: rgba(18, 10, 18, 0.82);
    border: 1px solid rgba(246, 196, 83, 0.5);
    border-radius: 16px;
    box-shadow: 0 18px 30px rgba(0, 0, 0, 0.34);
    color: #f8efe4;
    margin: 0.55rem auto 0;
    max-width: 46rem;
    padding: 0.72rem 0.95rem;
    position: relative;
    text-align: center;
    z-index: 4;
}
.speech-bubble::after {
    border-left: 10px solid transparent;
    border-right: 10px solid transparent;
    border-top: 12px solid rgba(246, 196, 83, 0.5);
    bottom: -12px;
    content: "";
    left: 50%;
    position: absolute;
    transform: translateX(-50%);
}
.speech-speaker {
    color: #ffd166;
    font-size: 0.78rem;
    font-weight: 800;
    letter-spacing: 0.08em;
    margin-bottom: 0.18rem;
    text-transform: uppercase;
}
.speech-line {
    color: #f8efe4;
    font-size: 0.96rem;
    line-height: 1.35;
}
.actor-row {
    display: grid;
    grid-template-columns: repeat(auto-fit, minmax(132px, 1fr));
    gap: 0.55rem;
    margin-top: 0.72rem;
    position: relative;
    z-index: 3;
}
.actor-card {
    background: rgba(70, 38, 36, 0.72);
    border: 1px solid rgba(246, 196, 83, 0.45);
    border-radius: 16px 16px 10px 10px;
    box-shadow: 0 14px 28px rgba(0, 0, 0, 0.28);
    min-height: 132px;
    padding: 0.58rem 0.62rem 0.72rem;
    position: relative;
    transform-origin: bottom center;
    text-align: center;
}
.actor-card::after {
    background: #7c3f17;
    border-radius: 0 0 8px 8px;
    bottom: -22px;
    box-shadow: inset 0 -5px 8px rgba(0, 0, 0, 0.2);
    content: "";
    height: 22px;
    left: calc(50% - 8px);
    position: absolute;
    width: 16px;
}
.actor-card.active {
    animation: puppet-bounce 0.78s ease-in-out infinite alternate;
    border-color: #ffd166;
    box-shadow:
        0 0 0 2px rgba(255, 209, 102, 0.22),
        0 0 34px rgba(255, 209, 102, 0.46),
        0 16px 34px rgba(0, 0, 0, 0.34);
}
.actor-avatar {
    background: radial-gradient(circle, rgba(255, 209, 102, 0.2), rgba(59, 10, 22, 0.3));
    border: 1px solid rgba(246, 196, 83, 0.34);
    border-radius: 999px;
    display: inline-grid;
    font-size: 1.7rem;
    height: 3rem;
    place-items: center;
    text-align: center;
    width: 3rem;
}
.actor-name {
    color: #f8efe4;
    font-weight: 700;
    line-height: 1.15;
    margin-top: 0.35rem;
    text-align: center;
}
.speaking-pill {
    background: #ffd166;
    border-radius: 999px;
    color: #3b0a16;
    display: inline-block;
    font-size: 0.64rem;
    font-weight: 800;
    margin-top: 0.26rem;
    padding: 0.12rem 0.44rem;
    text-transform: uppercase;
}
.actor-detail {
    color: #cbb7a1;
    font-size: 0.72rem;
    line-height: 1.28;
    margin-top: 0.35rem;
}
.actor-detail strong {
    color: #f8efe4;
}
.held-prop {
    margin-top: 0.42rem;
}
.held-prop span {
    background: rgba(246, 196, 83, 0.14);
    border: 1px solid rgba(246, 196, 83, 0.32);
    border-radius: 999px;
    color: #ffd166;
    display: inline-block;
    font-size: 0.68rem;
    font-weight: 700;
    padding: 0.12rem 0.42rem;
}
.beat-counter {
    color: #ffd166;
    font-weight: 800;
    margin-top: 0.55rem;
    position: relative;
    text-align: center;
    z-index: 3;
}
.stage-events {
    display: grid;
    gap: 0.4rem;
    margin-top: 0.55rem;
    position: relative;
    z-index: 3;
}
.audience-action,
.prop-pile {
    background: rgba(42, 20, 38, 0.7);
    border: 1px solid rgba(246, 196, 83, 0.25);
    border-radius: 999px;
    color: #f8efe4;
    margin: 0 auto;
    max-width: 48rem;
    padding: 0.38rem;
    text-align: center;
    width: 100%;
}
.audience-action strong,
.prop-pile strong {
    color: #ffd166;
}
.prop-token {
    animation: prop-pop 0.22s ease-out;
    background: rgba(246, 196, 83, 0.17);
    border: 1px solid rgba(246, 196, 83, 0.5);
    border-radius: 999px;
    color: #fff7ed;
    display: inline-block;
    margin: 0.2rem;
    padding: 0.22rem 0.55rem;
}
.gradio-container button.primary,
.gradio-container button.primary-action,
.gradio-container button.run-one-action {
    background: #f97316 !important;
    border-color: #f97316 !important;
    box-shadow: 0 10px 24px rgba(249, 115, 22, 0.25) !important;
    color: #fff7ed !important;
}
.gradio-container button.secondary,
.gradio-container button.secondary-action,
.gradio-container button.audience-action-button {
    background: #3f3148 !important;
    border-color: rgba(246, 196, 83, 0.22) !important;
    color: #f8efe4 !important;
}
.gradio-container button.reset-action {
    background: #3b0a16 !important;
    border-color: rgba(246, 196, 83, 0.24) !important;
    color: #f8efe4 !important;
}
.transcript-box,
.gradio-container .accordion {
    background: rgba(13, 6, 14, 0.58) !important;
    border-color: rgba(246, 196, 83, 0.18) !important;
    color: #f8efe4 !important;
}
@keyframes puppet-bounce {
    from { transform: translateY(0) rotate(-0.4deg); }
    to { transform: translateY(-7px) rotate(0.7deg); }
}
@keyframes bubble-in {
    from { opacity: 0; transform: translateY(8px); }
    to { opacity: 1; transform: translateY(0); }
}
@keyframes prop-pop {
    from { opacity: 0; transform: scale(0.86); }
    to { opacity: 1; transform: scale(1); }
}
@media (max-width: 760px) {
    .puppet-stage {
        min-height: 560px;
    }
    .puppet-stage::before,
    .puppet-stage::after {
        width: 7%;
    }
    .stage-backdrop {
        padding: 0.8rem 1.4rem;
    }
    .stage-marquee {
        font-size: 1.2rem;
    }
    .actor-row {
        grid-template-columns: repeat(2, minmax(0, 1fr));
    }
    .actor-card {
        min-height: 126px;
    }
}

/* Final Gradio chrome overrides: keep the whole app in the theater palette. */
.gradio-container {
    width: min(1200px, calc(100vw - 2rem)) !important;
}
.gradio-container .gr-group {
    background: rgba(34, 17, 31, 0.84) !important;
    border: 1px solid rgba(246, 196, 83, 0.2) !important;
    border-radius: 8px !important;
    color: #f8efe4 !important;
}
.gradio-container .gr-group .form,
.gradio-container .gr-group .block,
.gradio-container .gr-group .wrap,
.gradio-container .gr-group .wrap-inner,
.gradio-container .gr-group .secondary-wrap,
.gradio-container .gr-group .input-container,
.gradio-container .gr-group label {
    background: transparent !important;
    color: #f8efe4 !important;
}
.gradio-container input,
.gradio-container textarea,
.gradio-container select,
.gradio-container .dropdown-container,
.gradio-container .wrap-inner {
    background: rgba(10, 12, 23, 0.9) !important;
    color: #f8efe4 !important;
}
.gradio-container .control-panel input,
.gradio-container .control-panel textarea,
.gradio-container .control-panel .wrap-inner,
.gradio-container .premise-panel textarea {
    border: 1px solid rgba(246, 196, 83, 0.24) !important;
}
.gradio-container button {
    background: #3f3148 !important;
    border: 1px solid rgba(246, 196, 83, 0.24) !important;
    color: #f8efe4 !important;
}
.gradio-container button.primary,
.gradio-container button.primary-action,
.gradio-container button.run-one-action {
    background: #f97316 !important;
    border-color: #f97316 !important;
    color: #fff7ed !important;
}
.gradio-container button.reset-action {
    background: #3b0a16 !important;
    border-color: rgba(246, 196, 83, 0.32) !important;
}
.gradio-container .html-container,
.gradio-container .gradio-style {
    width: 100% !important;
}
.puppet-stage {
    min-height: 500px;
    width: 100%;
}
.puppet-stage::before,
.puppet-stage::after {
    width: clamp(56px, 9%, 110px);
}
.stage-backdrop {
    padding: 0.78rem clamp(4.1rem, 11vw, 8.8rem) 0.72rem;
}
.stage-marquee {
    font-size: clamp(1.25rem, 2.1vw, 1.72rem);
    white-space: normal;
}
.speech-bubble {
    margin-top: 0.48rem;
    max-width: 44rem;
    padding: 0.58rem 0.82rem;
}
.actor-row {
    align-items: end;
    grid-template-columns: repeat(auto-fit, minmax(116px, 1fr));
    gap: 0.62rem;
    margin-top: 0.82rem;
}
.actor-card {
    align-content: start;
    background: radial-gradient(circle at 50% 18%, rgba(246, 196, 83, 0.13), rgba(70, 38, 36, 0.72) 58%);
    border-radius: 18px;
    display: grid;
    justify-items: center;
    min-height: 108px;
    padding: 0.5rem 0.45rem 0.56rem;
}
.actor-card::after {
    bottom: -20px;
    height: 20px;
    width: 14px;
}
.actor-avatar {
    font-size: 2rem;
    height: 3.3rem;
    width: 3.3rem;
}
.actor-name {
    font-size: 0.82rem;
    margin-top: 0.28rem;
}
.actor-detail {
    display: -webkit-box;
    font-size: 0.66rem;
    line-height: 1.18;
    margin-top: 0.2rem;
    max-width: 11rem;
    min-height: 1.55rem;
    overflow: hidden;
    -webkit-box-orient: vertical;
    -webkit-line-clamp: 2;
}
.held-prop {
    margin-top: 0.26rem;
}
.held-prop span {
    font-size: 0.62rem;
    padding: 0.08rem 0.34rem;
}
.speaking-pill {
    font-size: 0.58rem;
    margin-top: 0.18rem;
    padding: 0.08rem 0.36rem;
}
.stage-events {
    gap: 0.32rem;
    margin-top: 0.64rem;
}
.audience-action,
.prop-pile {
    max-width: 45rem;
    padding: 0.3rem 0.55rem;
}
@media (max-width: 760px) {
    .gradio-container {
        width: min(100vw, calc(100vw - 0.75rem)) !important;
    }
    .puppet-stage::before,
    .puppet-stage::after {
        width: 30px;
    }
    .stage-backdrop {
        padding: 0.75rem 2.45rem;
    }
    .actor-row {
        grid-template-columns: repeat(2, minmax(0, 1fr));
        gap: 0.45rem;
    }
    .actor-card {
        min-height: 102px;
        padding-left: 0.28rem;
        padding-right: 0.28rem;
    }
}

/* Compact stage pass: keep the theater look, reduce scrolling, and keep controls close. */
.gradio-container {
    padding-top: 0.65rem !important;
}
.app-title h1 {
    font-size: 1.95rem;
}
.app-title p {
    margin-bottom: 0.55rem;
}
.premise-panel {
    padding: 0.42rem 0.55rem 0.52rem;
}
.stage-output,
.stage-output .html-container,
.stage-output .gradio-style {
    margin-bottom: 0 !important;
}
.puppet-stage {
    min-height: 390px;
}
.stage-valance {
    height: 34px;
    border-bottom-width: 3px;
}
.stage-backdrop {
    padding: 0.48rem clamp(3.9rem, 9vw, 7.3rem) 0.46rem;
}
.stage-marquee {
    font-size: clamp(1.15rem, 1.9vw, 1.52rem);
}
.stage-copy {
    font-size: 0.76rem;
    line-height: 1.25;
    margin-top: 0.14rem;
}
.speech-bubble {
    border-radius: 12px;
    margin-top: 0.34rem;
    max-width: 40rem;
    padding: 0.42rem 0.7rem;
}
.speech-speaker {
    font-size: 0.68rem;
}
.speech-line {
    font-size: 0.86rem;
}
.actor-row {
    grid-template-columns: repeat(auto-fit, minmax(104px, 1fr));
    gap: 0.5rem;
    margin-top: 0.55rem;
}
.actor-card {
    border-radius: 14px;
    min-height: 88px;
    padding: 0.38rem 0.36rem 0.44rem;
}
.actor-card::after {
    bottom: -16px;
    height: 16px;
}
.actor-avatar {
    font-size: 1.65rem;
    height: 2.55rem;
    width: 2.55rem;
}
.actor-name {
    font-size: 0.74rem;
    margin-top: 0.2rem;
}
.actor-detail {
    font-size: 0.6rem;
    line-height: 1.12;
    margin-top: 0.14rem;
    min-height: 1.35rem;
}
.speaking-pill {
    font-size: 0.52rem;
    margin-top: 0.14rem;
}
.held-prop {
    margin-top: 0.18rem;
}
.held-prop span {
    font-size: 0.55rem;
}
.stage-events {
    gap: 0.24rem;
    margin-top: 0.46rem;
}
.audience-action,
.prop-pile {
    font-size: 0.78rem;
    max-width: 39rem;
    padding: 0.22rem 0.5rem;
}
.prop-token {
    margin: 0.08rem;
    padding: 0.12rem 0.4rem;
}
.beat-counter {
    font-size: 0.84rem;
    margin-top: 0.34rem;
}
.stage-floorboards {
    height: 40px;
}
.control-panel {
    margin-top: 0 !important;
    padding: 0.42rem;
}
.control-panel h3 {
    margin-bottom: 0.2rem;
}
.gradio-container .row {
    gap: 0.55rem !important;
}
.stage-output + .row,
.stage-output + div,
.control-panel + .control-panel {
    margin-top: 0.45rem !important;
}
.transcript-section,
.gradio-container .accordion {
    margin-top: 0.55rem !important;
}
@media (max-width: 760px) {
    .puppet-stage {
        min-height: 430px;
    }
    .stage-backdrop {
        padding: 0.52rem 2.15rem;
    }
    .actor-row {
        grid-template-columns: repeat(2, minmax(0, 1fr));
    }
    .speech-line {
        font-size: 0.8rem;
    }
}
"""


def render_stage(session: TheaterSession | None) -> str:
    if session is None:
        return EMPTY_STAGE

    actor_cards = []
    latest_beat = session.transcript[-1] if session.transcript else None
    latest_speaker = latest_beat.speaker if latest_beat else None
    for actor in session.actors:
        active_class = " active" if actor.name == latest_speaker else ""
        active_label = '<div class="speaking-pill">Now speaking</div>' if actor.name == latest_speaker else ""
        role_line = actor.goal.split(".", maxsplit=1)[0]
        held_prop = actor.held_prop or "nothing"
        held_emoji = PROP_EMOJI.get(held_prop.lower(), "🎁") if actor.held_prop else ""
        actor_cards.append(
            f"""
            <div class="actor-card{active_class}">
              <div class="actor-avatar">{escape(actor.avatar)}</div>
              <div class="actor-name">{escape(actor.name)}</div>
              {active_label}
              <div class="actor-detail">{escape(role_line)}</div>
              <div class="held-prop"><span>Holding: {escape((held_emoji + " ") if held_emoji else "")}{escape(held_prop)}</span></div>
            </div>
            """
        )
    latest_line = ""
    if latest_beat is not None:
        latest_line = f"""
        <div class="speech-bubble">
          <div class="speech-speaker">{escape(latest_beat.speaker)}</div>
          <div class="speech-line">{escape(latest_beat.line)}</div>
        </div>
        """
    audience_action = ""
    if session.latest_audience_action is not None:
        audience_action = f"""
        <div class="audience-action">
          <strong>Audience:</strong> {escape(session.latest_audience_action)}
        </div>
        """
    prop_pile = ""
    if session.props:
        prop_tokens = "".join(
            f'<span class="prop-token">{escape(PROP_EMOJI.get(prop.lower(), "🎁"))} {escape(prop)}</span>'
            for prop in session.props
        )
        prop_pile = f"""
        <div class="prop-pile">
          <strong>Props on stage:</strong> {prop_tokens}
        </div>
        """

    return f"""
    <div class="puppet-stage stage-live">
      <div class="stage-valance"></div>
      <div class="stage-backdrop">
        <div class="stage-marquee">{escape(session.show_title)}</div>
        <div class="stage-copy">
          <strong>Setting:</strong> {escape(session.setting)}<br />
          <strong>Premise:</strong> {escape(session.premise)}
        </div>
        {latest_line}
        <div class="actor-row">
          {''.join(actor_cards)}
        </div>
        <div class="stage-events">
          {audience_action}
          {prop_pile}
        </div>
        <div class="beat-counter">Beat {session.beat_index} of {session.max_beats}</div>
      </div>
      <div class="stage-floorboards"></div>
    </div>
    """


def render_transcript(session: TheaterSession | None) -> str:
    if session is None:
        return EMPTY_TRANSCRIPT

    transcript_lines = [
        "Transcript:",
        "No puppet lines yet. The first beat will be added in the next milestone.",
    ]
    if session.transcript:
        transcript_lines = ["Transcript:"]
        for index, beat in enumerate(session.transcript, start=1):
            transcript_lines.append(f"{index}. {beat.speaker}: {beat.line}")

    return "\n".join(transcript_lines)


def render_director_log(session: TheaterSession | None) -> str:
    if session is None:
        return EMPTY_DIRECTOR_LOG
    return "\n".join(f"- {entry}" for entry in session.director_log)


def render_trace(session: TheaterSession | None) -> str:
    if session is None:
        return EMPTY_TRACE
    return "\n".join(f"- {entry}" for entry in session.trace_events)


def normalize_backend_name(backend_name: str | None) -> str:
    return backend_name if backend_name in BACKEND_CHOICES else "deterministic"


def normalize_max_new_tokens(max_new_tokens: int | float | None) -> int:
    if max_new_tokens is None:
        return DEFAULT_MAX_NEW_TOKENS
    return max(16, min(160, int(max_new_tokens)))


def normalize_temperature(temperature: int | float | None) -> float:
    if temperature is None:
        return DEFAULT_TEMPERATURE
    return max(0.0, min(1.5, float(temperature)))


def apply_backend_selection(
    session: TheaterSession | None,
    backend_name: str | None,
    max_new_tokens: int | float | None = None,
    temperature: int | float | None = None,
) -> TheaterSession | None:
    if session is None:
        return None
    session.backend_name = normalize_backend_name(backend_name)
    session.backend_model_id = OPENBMB_MODEL_ID if session.backend_name == "openbmb" else None
    session.backend_max_new_tokens = normalize_max_new_tokens(max_new_tokens)
    session.backend_temperature = normalize_temperature(temperature)
    return session


def render_backend_settings(
    session: TheaterSession | None,
    backend_name: str | None = None,
    max_new_tokens: int | float | None = None,
    temperature: int | float | None = None,
) -> str:
    selected_backend = normalize_backend_name(backend_name)
    active_backend = session.backend_name if session is not None else selected_backend
    model_id = session.backend_model_id if session is not None else None
    if active_backend == "openbmb":
        model_id = model_id or OPENBMB_MODEL_ID
    status = get_backend_status(active_backend)
    openbmb_status = get_backend_status("openbmb")
    configured_max_new_tokens = (
        session.backend_max_new_tokens if session is not None else normalize_max_new_tokens(max_new_tokens)
    )
    configured_temperature = (
        session.backend_temperature if session is not None else normalize_temperature(temperature)
    )
    latency = f"{status.latest_latency_ms}ms" if status.latest_latency_ms is not None else "none yet"
    fallback_reason = status.latest_fallback_reason or "none"
    return (
        f"Active backend: {active_backend}\n"
        "Available backends: deterministic, openbmb\n"
        f"OpenBMB model id: {model_id or 'not selected'}\n"
        f"Model status: {status.load_status}\n"
        f"OpenBMB status: {openbmb_status.load_status}\n"
        f"Latest latency: {latency}\n"
        f"Latest fallback reason: {fallback_reason}\n"
        f"Generation: max_new_tokens={configured_max_new_tokens}, temperature={configured_temperature:.2f}\n"
        "Fallback behavior: invalid model output falls back to deterministic actor lines"
    )


def render_outputs(session: TheaterSession | None):
    return (
        render_stage(session),
        render_transcript(session),
        render_director_log(session),
        render_trace(session),
        render_backend_settings(session),
    )


def create_show(
    premise: str,
    session: TheaterSession | None,
    backend_name: str,
    max_new_tokens: int | float,
    temperature: int | float,
):
    premise = premise.strip()
    selected_backend = normalize_backend_name(backend_name)
    selected_max_new_tokens = normalize_max_new_tokens(max_new_tokens)
    selected_temperature = normalize_temperature(temperature)
    if not premise:
        return (
            None,
            EMPTY_STAGE,
            "No premise yet. Add a premise to raise the curtain.",
            EMPTY_DIRECTOR_LOG,
            EMPTY_TRACE,
            render_backend_settings(None, selected_backend, selected_max_new_tokens, selected_temperature),
        )

    session = create_show_from_premise(
        premise,
        backend_name=selected_backend,
        backend_model_id=OPENBMB_MODEL_ID if selected_backend == "openbmb" else None,
        backend_max_new_tokens=selected_max_new_tokens,
        backend_temperature=selected_temperature,
    )
    return session, *render_outputs(session)


def reset_show():
    return (
        None,
        "",
        "rubber duck",
        "",
        EMPTY_STAGE,
        EMPTY_TRANSCRIPT,
        EMPTY_DIRECTOR_LOG,
        EMPTY_TRACE,
        "deterministic",
        DEFAULT_MAX_NEW_TOKENS,
        DEFAULT_TEMPERATURE,
        True,
        EMPTY_BACKEND,
    )


def advance_one_beat(
    session: TheaterSession | None,
    backend_name: str,
    max_new_tokens: int | float,
    temperature: int | float,
):
    if session is None:
        return (
            None,
            EMPTY_STAGE,
            "Create a show before running a beat.",
            EMPTY_DIRECTOR_LOG,
            EMPTY_TRACE,
            render_backend_settings(None, backend_name, max_new_tokens, temperature),
        )

    session = apply_backend_selection(session, backend_name, max_new_tokens, temperature)
    session = run_one_beat(session)
    return session, *render_outputs(session)


def advance_full_act(
    session: TheaterSession | None,
    backend_name: str,
    max_new_tokens: int | float,
    temperature: int | float,
    use_deterministic_full_act: bool,
):
    if session is None:
        yield (
            None,
            EMPTY_STAGE,
            "Create a show before running the full act.",
            EMPTY_DIRECTOR_LOG,
            EMPTY_TRACE,
            render_backend_settings(None, backend_name, max_new_tokens, temperature),
        )
        return

    session = apply_backend_selection(session, backend_name, max_new_tokens, temperature)
    selected_backend = session.backend_name
    deterministic_full_act = selected_backend == "openbmb" and use_deterministic_full_act
    if deterministic_full_act:
        session.director_log.append(
            "OpenBMB is selected, so Run Full Act will use deterministic actor lines for this playback."
        )
        session.trace_events.append("full_act_openbmb_deterministic_playback")

    if session.beat_index >= session.max_beats:
        if deterministic_full_act:
            session.backend_name = "deterministic"
        session = run_one_beat(session)
        if deterministic_full_act:
            session.backend_name = selected_backend
        yield session, *render_outputs(session)
        return

    while session.beat_index < session.max_beats:
        if deterministic_full_act:
            session.backend_name = "deterministic"
        session = run_one_beat(session)
        if deterministic_full_act:
            session.backend_name = selected_backend
        yield session, *render_outputs(session)
        if session.beat_index < session.max_beats:
            sleep(PLAYBACK_DELAY_SECONDS)


def throw_audience_prop(
    session: TheaterSession | None,
    prop_name: str,
    backend_name: str,
    max_new_tokens: int | float,
    temperature: int | float,
):
    if session is None:
        return (
            None,
            EMPTY_STAGE,
            "Create a show before throwing a prop.",
            EMPTY_DIRECTOR_LOG,
            EMPTY_TRACE,
            render_backend_settings(None, backend_name, max_new_tokens, temperature),
        )

    session = apply_backend_selection(session, backend_name, max_new_tokens, temperature)
    session = throw_prop(session, prop_name)
    return session, *render_outputs(session)


def summon_audience_actor(
    session: TheaterSession | None,
    actor_name: str,
    backend_name: str,
    max_new_tokens: int | float,
    temperature: int | float,
):
    if session is None:
        return (
            None,
            EMPTY_STAGE,
            "Create a show before summoning an actor.",
            EMPTY_DIRECTOR_LOG,
            EMPTY_TRACE,
            render_backend_settings(None, backend_name, max_new_tokens, temperature),
        )

    session = apply_backend_selection(session, backend_name, max_new_tokens, temperature)
    session = summon_actor(session, actor_name)
    return session, *render_outputs(session)


def request_audience_finale(
    session: TheaterSession | None,
    backend_name: str,
    max_new_tokens: int | float,
    temperature: int | float,
):
    if session is None:
        return (
            None,
            EMPTY_STAGE,
            "Create a show before requesting a finale.",
            EMPTY_DIRECTOR_LOG,
            EMPTY_TRACE,
            render_backend_settings(None, backend_name, max_new_tokens, temperature),
        )

    session = apply_backend_selection(session, backend_name, max_new_tokens, temperature)
    session = request_finale(session)
    return session, *render_outputs(session)


def warm_up_backend(
    session: TheaterSession | None,
    max_new_tokens: int | float,
    temperature: int | float,
):
    selected_max_new_tokens = normalize_max_new_tokens(max_new_tokens)
    selected_temperature = normalize_temperature(temperature)
    status = warm_up_openbmb(
        max_new_tokens=selected_max_new_tokens,
        temperature=selected_temperature,
    )
    if session is not None:
        session.backend_max_new_tokens = selected_max_new_tokens
        session.backend_temperature = selected_temperature
        if status.load_status == "loaded":
            session.director_log.append(f"OpenBMB warm-up loaded {status.model_id}.")
            session.trace_events.append(f"openbmb_warmup_loaded:{status.model_id}:latency_ms={status.latest_latency_ms}")
        else:
            reason = status.latest_fallback_reason or "unknown error"
            session.director_log.append(f"OpenBMB warm-up failed: {reason}.")
            session.trace_events.append(f"openbmb_warmup_failed:{status.model_id}:{reason}")
    return (
        session,
        render_director_log(session),
        render_trace(session),
        render_backend_settings(session, "openbmb", selected_max_new_tokens, selected_temperature),
    )


with gr.Blocks(title="AI Puppet Theater") as app:
    session_state = gr.State(None)

    gr.Markdown(
        """
        # AI Puppet Theater
        Create a tiny deterministic puppet show, then interrupt it from the audience.
        """,
        elem_classes=["app-title"],
    )

    with gr.Group(elem_classes=["control-panel", "premise-panel"]):
        premise_input = gr.Textbox(
            label="Premise",
            placeholder="A moon detective interrogates a suspicious toaster...",
            lines=1,
        )
        with gr.Row():
            create_button = gr.Button("Create Show", variant="primary", elem_classes=["primary-action"])
            reset_button = gr.Button("Reset", elem_classes=["reset-action"])

    stage_output = gr.HTML(value=EMPTY_STAGE, label="Stage", elem_classes=["stage-output"])

    with gr.Row():
        with gr.Group(elem_classes=["control-panel"]):
            gr.Markdown("### Show Controls")
            with gr.Row():
                run_one_button = gr.Button(
                    "Run One Beat",
                    variant="primary",
                    elem_classes=["run-one-action"],
                )
                run_full_button = gr.Button("Run Full Act", elem_classes=["secondary-action"])

        with gr.Group(elem_classes=["control-panel"]):
            gr.Markdown("### Audience")
            prop_input = gr.Dropdown(
                choices=["rubber duck", "egg", "flowers", "tomato", "tiny crown", "scroll"],
                value="rubber duck",
                allow_custom_value=True,
                label="Prop",
            )
            with gr.Row():
                throw_prop_button = gr.Button(
                    "Throw Prop",
                    elem_classes=["audience-action-button"],
                )
                request_finale_button = gr.Button(
                    "Request Finale",
                    elem_classes=["audience-action-button"],
                )
            actor_input = gr.Textbox(label="Summon Actor", placeholder="Professor Button")
            summon_actor_button = gr.Button(
                "Summon Actor",
                elem_classes=["audience-action-button"],
            )

    with gr.Accordion("Transcript", open=False, elem_classes=["transcript-section"]):
        transcript_output = gr.Textbox(
            value=EMPTY_TRANSCRIPT,
            label="Transcript",
            lines=7,
            interactive=False,
            elem_classes=["transcript-box"],
        )
    with gr.Accordion("Behind the Curtain", open=False):
        director_output = gr.Textbox(
            value=EMPTY_DIRECTOR_LOG,
            label="Director Log",
            lines=6,
            interactive=False,
        )
    with gr.Accordion("Trace / Debug", open=False):
        trace_output = gr.Textbox(
            value=EMPTY_TRACE,
            label="Trace Events",
            lines=6,
            interactive=False,
        )
    with gr.Accordion("Backend", open=False):
        backend_select = gr.Dropdown(
            choices=BACKEND_CHOICES,
            value="deterministic",
            label="Actor Line Backend",
            interactive=True,
        )
        with gr.Row():
            max_new_tokens_input = gr.Slider(
                minimum=16,
                maximum=160,
                value=DEFAULT_MAX_NEW_TOKENS,
                step=8,
                label="Max New Tokens",
                interactive=True,
            )
            temperature_input = gr.Slider(
                minimum=0.0,
                maximum=1.5,
                value=DEFAULT_TEMPERATURE,
                step=0.1,
                label="Temperature",
                interactive=True,
            )
        deterministic_full_act_input = gr.Checkbox(
            value=True,
            label="Use deterministic actor lines for OpenBMB full-act playback",
            interactive=True,
        )
        warm_up_button = gr.Button("Warm up OpenBMB", elem_classes=["secondary-action"])
        backend_output = gr.Textbox(
            value=EMPTY_BACKEND,
            label="Model Settings",
            lines=8,
            interactive=False,
        )

    create_button.click(
        create_show,
        inputs=[premise_input, session_state, backend_select, max_new_tokens_input, temperature_input],
        outputs=[session_state, stage_output, transcript_output, director_output, trace_output, backend_output],
    )
    run_one_button.click(
        advance_one_beat,
        inputs=[session_state, backend_select, max_new_tokens_input, temperature_input],
        outputs=[session_state, stage_output, transcript_output, director_output, trace_output, backend_output],
    )
    run_full_button.click(
        advance_full_act,
        inputs=[
            session_state,
            backend_select,
            max_new_tokens_input,
            temperature_input,
            deterministic_full_act_input,
        ],
        outputs=[session_state, stage_output, transcript_output, director_output, trace_output, backend_output],
    )
    throw_prop_button.click(
        throw_audience_prop,
        inputs=[session_state, prop_input, backend_select, max_new_tokens_input, temperature_input],
        outputs=[session_state, stage_output, transcript_output, director_output, trace_output, backend_output],
    )
    summon_actor_button.click(
        summon_audience_actor,
        inputs=[session_state, actor_input, backend_select, max_new_tokens_input, temperature_input],
        outputs=[session_state, stage_output, transcript_output, director_output, trace_output, backend_output],
    )
    request_finale_button.click(
        request_audience_finale,
        inputs=[session_state, backend_select, max_new_tokens_input, temperature_input],
        outputs=[session_state, stage_output, transcript_output, director_output, trace_output, backend_output],
    )
    warm_up_button.click(
        warm_up_backend,
        inputs=[session_state, max_new_tokens_input, temperature_input],
        outputs=[session_state, director_output, trace_output, backend_output],
    )
    reset_button.click(
        reset_show,
        outputs=[
            session_state,
            premise_input,
            prop_input,
            actor_input,
            stage_output,
            transcript_output,
            director_output,
            trace_output,
            backend_select,
            max_new_tokens_input,
            temperature_input,
            deterministic_full_act_input,
            backend_output,
        ],
    )


if __name__ == "__main__":
    app.launch(css=CUSTOM_CSS, theme=gr.themes.Soft())