#!/usr/bin/env python3
"""
SYNTHFELD web server — serves the UI + streaming APIs for the Seinfeld persona
model (Llama-3.1-8B SFT), running 4-bit GGUF on CPU via llama.cpp.

Two modes share the same model and prompt path:
  - /chat     : talk to one chosen character (assistant replies as [CHARACTER]).
  - /generate : from a neutral premise, write a 15-20 line scene; each line's
                speaker is a "dice roll" — uniform over the cast, no immediate repeat.

The model was trained to BE a chosen character: the assistant turn is prefixed
with a persona token (e.g. "[JERRY] ..."). At inference we force that prefix
server-side — we render the Llama-3.1 chat template, append the generation prompt,
then append "[CHARACTER]" (no trailing space, matching training) and let the model
continue. Chat lets the user pick who to talk to per message; Generate rolls the
speaker per line and feeds the scene-so-far back as the user turn.

Parity notes (must match build_sft.py / chat_test.py):
  - prompt = chat_template(messages, add_generation_prompt=True) + "[CHAR]"  (NO space)
  - prior assistant turns are fed back WITH their "[CHAR] " prefix
  - generation stops on BOTH <|eot_id|> (128009) and <|end_of_text|> (128001)
"""

import os
import json
import random
import asyncio
import logging
import threading
from pathlib import Path
from contextlib import asynccontextmanager
from typing import List, Optional

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, HTMLResponse
from pydantic import BaseModel
from jinja2 import Environment, BaseLoader
from jinja2.exceptions import TemplateError
from llama_cpp import Llama

# ----------------------------------------------------------------------------- config
APP_DIR     = Path(__file__).resolve().parent
MODEL_PATH  = os.environ.get("MODEL_PATH", str(APP_DIR / "model" / "puffyshirt-Q4_K_M.gguf"))
N_CTX       = int(os.environ.get("N_CTX", "4096"))
N_THREADS   = int(os.environ.get("N_THREADS", str(os.cpu_count() or 4)))

CHARACTERS  = ["JERRY", "GEORGE", "ELAINE", "KRAMER", "NEWMAN"]

EOT_ID      = 128009   # <|eot_id|>  — Llama turn end (the real stop during training)
EOS_ID      = 128001   # <|end_of_text|>
TERMINATORS = {EOT_ID, EOS_ID}

# generation defaults (match chat_test.py / chat_mlx.py)
DEFAULT_TEMPERATURE    = 0.8
DEFAULT_TOP_P          = 0.9
DEFAULT_TOP_K          = 40
DEFAULT_REPEAT_PENALTY = 1.2
DEFAULT_MAX_TOKENS     = 200
MAX_MAX_TOKENS         = 512

# Generate (scene) mode defaults
SCENE_MIN_LINES         = 15
SCENE_MAX_LINES         = 20
DEFAULT_LINE_MAX_TOKENS = 60     # per-line budget (one utterance, not a whole turn)
MAX_LINE_MAX_TOKENS     = 120
SCENE_LINE_RETRIES      = 2      # re-roll a few times if a line comes back empty

# abuse limits
MAX_MESSAGES   = 100
MAX_MESSAGE_LEN = 4000
MAX_PREMISE_LEN = MAX_MESSAGE_LEN
MAX_TOTAL_LEN  = 16000
MIN_TEMPERATURE = 0.0
MAX_TEMPERATURE = 2.0

# script-drift guards: the base model sometimes slides back into raw-script
# narration; cut the turn if it appears (mirrors chat_mlx.trim_to_turn).
STOP_MARKERS = ["New scene", "<|", "\nassistant", "\nJERRY:", "\nGEORGE:",
                "\nELAINE:", "\nKRAMER:", "\nNEWMAN:"]

logging.basicConfig(level=logging.INFO,
                    format="%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S")
logger = logging.getLogger("synthfeld")

# ----------------------------------------------------------------------------- chat template
def _raise(msg, *a, **k):
    raise TemplateError(msg)

_env = Environment(loader=BaseLoader(), trim_blocks=True, lstrip_blocks=True)
_env.globals["raise_exception"] = _raise
_CHAT_TEMPLATE = _env.from_string((APP_DIR / "chat_template.jinja").read_text())
_BOS = "<|begin_of_text|>"


def build_prompt(messages: List[dict], character: str) -> str:
    """Render the Llama-3.1 chat template, then force the persona prefix.

    `messages` is the template-ready list (assistant turns already carry their
    "[CHAR] " prefix). We append "[CHARACTER]" with NO trailing space — the space
    is the model's first generated token (training parity)."""
    rendered = _CHAT_TEMPLATE.render(
        messages=messages, add_generation_prompt=True, bos_token=_BOS)
    return rendered + f"[{character}]"

# ----------------------------------------------------------------------------- model
llm: Optional[Llama] = None
gen_lock = asyncio.Lock()   # llama.cpp is single-context; serialize generations


@asynccontextmanager
async def lifespan(app: FastAPI):
    global llm
    logger.info(f"Loading GGUF: {MODEL_PATH}  (n_ctx={N_CTX}, n_threads={N_THREADS})")
    llm = Llama(model_path=MODEL_PATH, n_ctx=N_CTX, n_threads=N_THREADS,
                verbose=False, logits_all=False)
    # sanity: each persona token must encode to exactly one id (atomic special token)
    for c in CHARACTERS:
        ids = llm.tokenize(f"[{c}]".encode("utf-8"), add_bos=False, special=True)
        flag = "ok" if len(ids) == 1 else "!! SPLIT — check tokenizer/GGUF !!"
        logger.info(f"persona [{c}] -> {ids}  {flag}")
    logger.info("Model ready.")
    yield


app = FastAPI(lifespan=lifespan)
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True,
                   allow_methods=["*"], allow_headers=["*"])

# ----------------------------------------------------------------------------- schema
class Msg(BaseModel):
    role: str
    content: str
    character: Optional[str] = None   # for assistant turns: who said it


class ChatRequest(BaseModel):
    messages: List[Msg]
    character: str
    temperature: Optional[float] = None
    max_tokens: Optional[int] = None


class GenerateRequest(BaseModel):
    premise: str                                  # neutral scene setup (not a character line)
    num_lines: Optional[int] = None               # default: random in [15, 20]
    temperature: Optional[float] = None
    max_tokens_per_line: Optional[int] = None


def validate(req: ChatRequest):
    if req.character not in CHARACTERS:
        raise HTTPException(400, f"Unknown character. Choose one of {CHARACTERS}")
    if not req.messages:
        raise HTTPException(400, "At least one message is required")
    if len(req.messages) > MAX_MESSAGES:
        raise HTTPException(400, f"Too many messages (max {MAX_MESSAGES})")
    if req.messages[-1].role != "user":
        raise HTTPException(400, "Last message must be a user turn")
    total = 0
    for i, m in enumerate(req.messages):
        if m.role not in ("user", "assistant"):
            raise HTTPException(400, f"Message {i}: role must be user|assistant")
        if not m.content:
            raise HTTPException(400, f"Message {i}: empty content")
        if len(m.content) > MAX_MESSAGE_LEN:
            raise HTTPException(400, f"Message {i}: too long (max {MAX_MESSAGE_LEN} chars)")
        if m.role == "assistant" and m.character not in CHARACTERS:
            raise HTTPException(400, f"Message {i}: assistant turn needs a valid character")
        total += len(m.content)
    if total > MAX_TOTAL_LEN:
        raise HTTPException(400, f"Conversation too long (max {MAX_TOTAL_LEN} chars)")
    if req.temperature is not None and not (MIN_TEMPERATURE <= req.temperature <= MAX_TEMPERATURE):
        raise HTTPException(400, f"temperature must be in [{MIN_TEMPERATURE}, {MAX_TEMPERATURE}]")


def validate_generate(req: GenerateRequest):
    premise = (req.premise or "").strip()
    if not premise:
        raise HTTPException(400, "A scene premise is required")
    if len(premise) > MAX_PREMISE_LEN:
        raise HTTPException(400, f"Premise too long (max {MAX_PREMISE_LEN} chars)")
    if req.num_lines is not None and not (SCENE_MIN_LINES <= req.num_lines <= SCENE_MAX_LINES):
        raise HTTPException(400, f"num_lines must be in [{SCENE_MIN_LINES}, {SCENE_MAX_LINES}]")
    if req.temperature is not None and not (MIN_TEMPERATURE <= req.temperature <= MAX_TEMPERATURE):
        raise HTTPException(400, f"temperature must be in [{MIN_TEMPERATURE}, {MAX_TEMPERATURE}]")


def to_template_messages(messages: List[Msg]) -> List[dict]:
    """Re-attach the persona prefix to prior assistant turns so the context the
    model sees matches its training distribution."""
    out = []
    for m in messages:
        if m.role == "assistant":
            out.append({"role": "assistant", "content": f"[{m.character}] {m.content}"})
        else:
            out.append({"role": "user", "content": m.content})
    return out

# ----------------------------------------------------------------------------- generation
def _generate_blocking(prompt_ids, character, params, loop, queue):
    """Runs in a worker thread; pushes ('token'|'error'|'done', payload) to queue."""
    out_tokens: List[int] = []
    last_clean = ""

    def emit(kind, payload):
        loop.call_soon_threadsafe(queue.put_nowait, (kind, payload))

    try:
        for tid in llm.generate(
            prompt_ids,
            top_k=params["top_k"],
            top_p=params["top_p"],
            temp=params["temperature"],
            repeat_penalty=params["repeat_penalty"],
            reset=True,
        ):
            if tid in TERMINATORS:
                break
            out_tokens.append(tid)
            if len(out_tokens) >= params["max_tokens"]:
                break

            text = llm.detokenize(out_tokens).decode("utf-8", errors="replace")
            if text.endswith("�"):
                continue   # mid multi-byte char; wait for the next token

            # script-drift stop: truncate and finish
            cut = min((text.find(mk) for mk in STOP_MARKERS if mk in text), default=-1)
            if cut != -1:
                new = text[len(last_clean):cut]
                if new:
                    emit("token", new)
                last_clean = text[:cut]
                break

            new = text[len(last_clean):]
            if new:
                emit("token", new)
                last_clean = text
    except Exception as e:                       # noqa: BLE001
        emit("error", str(e))
    finally:
        emit("done", last_clean.strip())


async def stream_chat(req: ChatRequest):
    params = {
        "temperature": req.temperature if req.temperature is not None else DEFAULT_TEMPERATURE,
        "top_p": DEFAULT_TOP_P,
        "top_k": DEFAULT_TOP_K,
        "repeat_penalty": DEFAULT_REPEAT_PENALTY,
        "max_tokens": min(req.max_tokens or DEFAULT_MAX_TOKENS, MAX_MAX_TOKENS),
    }
    prompt = build_prompt(to_template_messages(req.messages), req.character)
    prompt_ids = llm.tokenize(prompt.encode("utf-8"), add_bos=False, special=True)

    queue: asyncio.Queue = asyncio.Queue()
    loop = asyncio.get_running_loop()

    full = []
    async with gen_lock:
        threading.Thread(
            target=_generate_blocking,
            args=(prompt_ids, req.character, params, loop, queue),
            daemon=True,
        ).start()
        while True:
            kind, payload = await queue.get()
            if kind == "token":
                full.append(payload)
                yield f"data: {json.dumps({'token': payload}, ensure_ascii=False)}\n\n"
            elif kind == "error":
                logger.error(f"generation error: {payload}")
                yield f"data: {json.dumps({'error': payload})}\n\n"
                break
            else:  # done
                logger.info(f"[{req.character}] {''.join(full).strip()[:300]}")
                yield f"data: {json.dumps({'done': True})}\n\n"
                break


@app.post("/chat")
async def chat(req: ChatRequest):
    validate(req)
    logger.info("=" * 20)
    for m in req.messages:
        tag = m.role.upper() + (f"[{m.character}]" if m.character else "")
        logger.info(f"{tag}: {m.content[:200]}")
    logger.info(f"-> generating as [{req.character}]")
    return StreamingResponse(stream_chat(req), media_type="text/event-stream")


# ----------------------------------------------------------------------------- scene generation
# A line ends at the first newline, the next persona/scene tag, or a template token.
LINE_STOP_MARKERS = ["\n", "[", "<|"]


def roll_speaker(last: Optional[str]) -> str:
    """Dice roll for the next speaker: uniform over the cast, never twice in a row."""
    pool = [c for c in CHARACTERS if c != last] or CHARACTERS
    return random.choice(pool)


def _generate_line_blocking(prompt_ids, params, loop, queue):
    """Generate ONE line (worker thread). Same streaming contract as
    _generate_blocking, but stops at the first newline / next tag so a single
    utterance is produced (mirrors gen_nextline.trim). 'done' payload is the
    stripped line."""
    out_tokens: List[int] = []
    last_clean = ""

    def emit(kind, payload):
        loop.call_soon_threadsafe(queue.put_nowait, (kind, payload))

    try:
        for tid in llm.generate(
            prompt_ids,
            top_k=params["top_k"],
            top_p=params["top_p"],
            temp=params["temperature"],
            repeat_penalty=params["repeat_penalty"],
            reset=True,
        ):
            if tid in TERMINATORS:
                break
            out_tokens.append(tid)
            if len(out_tokens) >= params["max_tokens"]:
                break

            text = llm.detokenize(out_tokens).decode("utf-8", errors="replace")
            if text.endswith("�"):
                continue   # mid multi-byte char; wait for the next token

            # a leading newline is just the model warming up — ignore until there's content
            cut = min((text.find(mk) for mk in LINE_STOP_MARKERS
                       if mk in text and text.find(mk) > 0), default=-1)
            if cut != -1:
                new = text[len(last_clean):cut]
                if new:
                    emit("token", new)
                last_clean = text[:cut]
                break

            new = text[len(last_clean):]
            if new:
                emit("token", new)
                last_clean = text
    except Exception as e:                       # noqa: BLE001
        emit("error", str(e))
    finally:
        emit("done", last_clean.strip())


async def stream_scene(req: GenerateRequest):
    params = {
        "temperature": req.temperature if req.temperature is not None else DEFAULT_TEMPERATURE,
        "top_p": DEFAULT_TOP_P,
        "top_k": DEFAULT_TOP_K,
        "repeat_penalty": DEFAULT_REPEAT_PENALTY,
        "max_tokens": min(req.max_tokens_per_line or DEFAULT_LINE_MAX_TOKENS, MAX_LINE_MAX_TOKENS),
    }
    num_lines = req.num_lines or random.randint(SCENE_MIN_LINES, SCENE_MAX_LINES)
    premise = req.premise.strip()
    scene_header = f"[Scene: {premise}]"          # in-distribution stage direction

    loop = asyncio.get_running_loop()
    scene_lines: List[str] = []
    last_speaker: Optional[str] = None

    async with gen_lock:
        # announce the premise so the UI can show it as the scene direction
        yield f"data: {json.dumps({'premise': premise}, ensure_ascii=False)}\n\n"

        produced = 0
        while produced < num_lines:
            speaker = roll_speaker(last_speaker)
            user_content = "\n".join([scene_header, *scene_lines])
            prompt = build_prompt([{"role": "user", "content": user_content}], speaker)
            prompt_ids = llm.tokenize(prompt.encode("utf-8"), add_bos=False, special=True)

            # try a line; re-roll the speaker a few times if it comes back empty
            line = ""
            for attempt in range(SCENE_LINE_RETRIES + 1):
                first = attempt == 0
                if first:
                    yield f"data: {json.dumps({'speaker': speaker, 'index': produced})}\n\n"
                queue: asyncio.Queue = asyncio.Queue()
                threading.Thread(target=_generate_line_blocking,
                                 args=(prompt_ids, params, loop, queue), daemon=True).start()
                err = None
                while True:
                    kind, payload = await queue.get()
                    if kind == "token":
                        if first:    # only stream the first attempt's tokens to the UI
                            yield f"data: {json.dumps({'token': payload}, ensure_ascii=False)}\n\n"
                    elif kind == "error":
                        err = payload
                        break
                    else:            # done
                        line = payload
                        break
                if err:
                    logger.error(f"scene generation error: {err}")
                    yield f"data: {json.dumps({'error': err})}\n\n"
                    return
                if line:
                    break
                # empty -> re-roll a different speaker and overwrite the placeholder line
                speaker = roll_speaker(last_speaker)
                prompt = build_prompt([{"role": "user", "content": user_content}], speaker)
                prompt_ids = llm.tokenize(prompt.encode("utf-8"), add_bos=False, special=True)
                yield f"data: {json.dumps({'speaker': speaker, 'index': produced, 'reroll': True})}\n\n"

            if not line:
                # couldn't get a non-empty line; skip this slot without recording a speaker
                yield f"data: {json.dumps({'line_skipped': True, 'index': produced})}\n\n"
                produced += 1
                continue

            scene_lines.append(f"[{speaker}] {line}")
            last_speaker = speaker
            logger.info(f"  [{speaker}] {line[:160]}")
            yield f"data: {json.dumps({'line_done': True, 'speaker': speaker, 'line': line}, ensure_ascii=False)}\n\n"
            produced += 1

        yield f"data: {json.dumps({'done': True})}\n\n"


@app.post("/generate")
async def generate_scene(req: GenerateRequest):
    validate_generate(req)
    logger.info("=" * 20)
    logger.info(f"GENERATE premise: {req.premise[:200]}")
    return StreamingResponse(stream_scene(req), media_type="text/event-stream")


@app.get("/")
async def root():
    return HTMLResponse((APP_DIR / "index.html").read_text(encoding="utf-8"))


@app.get("/health")
async def health():
    return {"status": "ok", "ready": llm is not None,
            "characters": CHARACTERS, "modes": ["chat", "generate"]}