app.py

import json
import tempfile
import shutil
from pathlib import Path

from gradio import Server
from fastapi import File, UploadFile
from fastapi.responses import HTMLResponse, StreamingResponse

from personas import PERSONAS, JUDGE_SYSTEM
from extractor import run_vision_pass, build_context_block
from config import settings
from ui_assets import CSS, APP_HEAD
from llm import stream_text_generation
PERSONA_KEYS = ["senku", "naruto", "light", "saitama"]

app = Server()


# ── Serve the UI ──────────────────────────────────────────────────────────────
@app.get("/", response_class=HTMLResponse)
async def homepage():
    return build_page()


# ── SSE debate endpoint ───────────────────────────────────────────────────────
@app.post("/debate-stream")
async def debate_stream(file: UploadFile = File(...)):
    suffix = Path(file.filename).suffix or ".jpg"
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
    shutil.copyfileobj(file.file, tmp)
    tmp.close()
    image_path = tmp.name

    def event(data: dict) -> str:
        return f"data: {json.dumps(data)}\n\n"

    def generate():
        try:
            yield event({"type": "status", "text": "🔍 Analyzing image..."})
            try:
                scene = run_vision_pass(image_path)
            except Exception as e:
                yield event({"type": "error", "text": str(e)})
                return

            scene_context = build_context_block(scene)
            yield event({"type": "vision_done", "scene": scene})
            yield event({"type": "status", "text": "✓ Scene locked. Starting debate..."})

            history = []
            for round_num in range(1, settings.debate_rounds + 1):
                yield event({"type": "round_start", "round": round_num})
                for persona_key in PERSONA_KEYS:
                    persona = PERSONAS[persona_key]
                    yield event({"type": "turn_start", "persona": persona_key,
                                 "name": persona["name"], "emoji": persona["emoji"],
                                 "round": round_num})

                    history_block = ""
                    if history:
                        debate_so_far = "\n\n".join(
                            f"{PERSONAS[h['persona']]['emoji']} {PERSONAS[h['persona']]['name']}: {h['content']}"
                            for h in history
                        )
                        history_block = f"\n\nDEBATE SO FAR:\n{debate_so_far}"

                    instruction = (
                        "This is Round 1. Give your opening argument about the meaning of this image. "
                        "Be opinionated. Be brief (2-4 sentences). "
                        "Do NOT prefix your response with your name or emoji."
                        if round_num == 1
                        else f"This is Round {round_num}. Respond directly to what others said above. "
                        "Disagree, attack, defend. Be brief (2-4 sentences). "
                        "Do NOT prefix your response with your name or emoji."
                    )

                    messages = [
                        {"role": "system", "content": persona["system"]},
                        {"role": "user", "content": f"{scene_context}{history_block}\n\n{instruction}"},
                    ]
                    stream = stream_text_generation(messages)

                    full_response = ""
                    for token in stream:
                        if token:
                            full_response += token
                            yield event({"type": "token", "persona": persona_key, "text": token})

                    history.append({"persona": persona_key, "round": round_num,
                                    "content": full_response.strip()})
                    yield event({"type": "turn_done", "persona": persona_key})

            # Judge
            yield event({"type": "judge_start"})
            transcript = "\n\n".join(
                f"{PERSONAS[h['persona']]['emoji']} {PERSONAS[h['persona']]['name']}: {h['content']}"
                for h in history
            )
            messages = [
                {"role": "system", "content": JUDGE_SYSTEM},
                {"role": "user", "content": f"{scene_context}\n\nFULL DEBATE TRANSCRIPT:\n{transcript}\n\nDeliver your verdict now."},
            ]
            stream = stream_text_generation(messages)
            
            judge_text = ""
            for token in stream:
                if token:
                    judge_text += token
                    yield event({"type": "judge_token", "text": token})

            yield event({"type": "done", "verdict": judge_text.strip()})

        finally:
            Path(image_path).unlink(missing_ok=True)

    return StreamingResponse(
        generate(),
        media_type="text/event-stream",
        headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
    )


# ── HTML page builder ─────────────────────────────────────────────────────────
def build_page() -> str:
    return f"""<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="UTF-8">
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  <title>🎭 Tiny Debate Club</title>
  <link rel="preconnect" href="https://fonts.googleapis.com">
  <link href="https://fonts.googleapis.com/css2?family=Playfair+Display:ital,wght@0,700;0,900;1,700&family=DM+Mono:wght@400;500&display=swap" rel="stylesheet">
  <style>{CSS}</style>
</head>
<body>
  <div class="tdc-app">
    <div class="tdc-header">
      <h1>🎭 Tiny Debate Club</h1>
      <p>Upload an image · Watch AI personas argue about its meaning</p>
    </div>
    <div class="tdc-body">
      <aside class="tdc-sidebar">
        <label class="upload-area" id="tdc-upload-area" for="tdc-file-input">
          <span class="upload-icon">📷</span>
          <span class="upload-hint">Click or drag &amp; drop<br>to upload an image</span>
          <div class="upload-overlay">
            <span style="font-size:1.4rem">🔄</span>
            <span style="font-size:.72rem;color:#fff;margin-top:.3rem">Change image</span>
          </div>
        </label>
        <input type="file" id="tdc-file-input" accept="image/*" style="display:none">
        <button class="btn" id="tdc-start-btn" disabled>⚡ START DEBATE</button>
        <div class="info-card" id="tdc-scene-card">
          <div class="info-label">Scene</div>
          <div class="info-text" id="tdc-scene-text"></div>
          <div class="info-mood" id="tdc-scene-mood"></div>
        </div>
        <div class="verdict-card" id="tdc-verdict-card">
          <div class="verdict-label">⚖ Verdict</div>
          <div class="verdict-text" id="tdc-verdict-text"></div>
        </div>
        <div class="status-line" id="tdc-status">Ready.</div>
      </aside>
      <div class="tdc-chat-wrap">
        <div class="chat-log" id="tdc-chat-log">
          <div class="empty" id="tdc-empty-state">
            <span class="empty-icon">🎭</span>
            Awaiting image upload
          </div>
        </div>
      </div>
    </div>
  </div>
  {APP_HEAD}
</body>
</html>"""


# ── Launch ────────────────────────────────────────────────────────────────────
demo = app   # HF Spaces looks for `demo`

app.launch(show_error=True)