src/midnight_static/writer.py

"""Writer contract for producing schema-valid radio scripts."""

from __future__ import annotations

import json
import os
import re
import sys
from dataclasses import asdict
from functools import cached_property
from typing import Protocol

from .fixtures import fixture_script
from .schema import Genre, Script, script_from_json_dict


class ScriptWriter(Protocol):
    """Interface shared by fixture, Nemotron, and fine-tuned writers."""

    name: str

    def write(
        self, premise: str, genre: Genre = Genre.WEIRD, feedback: str | None = None
    ) -> Script:
        """Return a validated script for the given premise and genre.

        `feedback` carries the error from a previous failed attempt so model
        writers can self-correct on a retry; deterministic writers ignore it.
        """


class FixtureWriter:
    """Deterministic schema-valid writer used until model integration lands."""

    name = "fixture"

    def write(
        self, premise: str, genre: Genre = Genre.WEIRD, feedback: str | None = None
    ) -> Script:
        return fixture_script(premise, genre=genre)


class NemotronWriter:
    """Nemotron-backed writer using Transformers generation.

    The model is lazy-loaded so importing this class never downloads weights.
    Runtime configuration:
    - `MIDNIGHT_NEMOTRON_MODEL`: HF model id.
    - `MIDNIGHT_NEMOTRON_DEVICE_MAP`: Transformers device map, default `auto`.
    - `MIDNIGHT_NEMOTRON_MAX_NEW_TOKENS`: generation cap, default `4000`.

    Note: Nemotron 3 Nano is a reasoning model — it emits thinking tokens before
    the JSON answer, so the cap must be generous (a GPU smoke showed ~8.6k chars
    of output; 1800 tokens truncated the JSON). See modal/smoke_nemotron.py.
    """

    name = "nemotron"

    def __init__(
        self,
        model_id: str | None = None,
        device_map: str | None = None,
        max_new_tokens: int | None = None,
    ) -> None:
        self.model_id = model_id or os.getenv(
            "MIDNIGHT_NEMOTRON_MODEL",
            "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16",
        )
        self.device_map = device_map or os.getenv("MIDNIGHT_NEMOTRON_DEVICE_MAP", "auto")
        self.max_new_tokens = max_new_tokens or int(
            os.getenv("MIDNIGHT_NEMOTRON_MAX_NEW_TOKENS", "4000")
        )

    def write(
        self, premise: str, genre: Genre = Genre.WEIRD, feedback: str | None = None
    ) -> Script:
        prompt = build_writer_prompt(premise, genre, feedback=feedback)
        raw = self._generate_text(prompt)
        return parse_writer_json(raw)

    @cached_property
    def _model_bundle(self):
        try:
            import torch
            from transformers import AutoModelForCausalLM, AutoTokenizer
        except ImportError as exc:  # pragma: no cover - environment guard
            raise RuntimeError(
                "NemotronWriter requires the writer/model dependencies. "
                "Run `uv sync --extra dev --extra tts` for the current scaffold."
            ) from exc

        tokenizer = AutoTokenizer.from_pretrained(self.model_id, trust_remote_code=True)
        dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
        model = AutoModelForCausalLM.from_pretrained(
            self.model_id,
            torch_dtype=dtype,
            device_map=self.device_map,
            trust_remote_code=True,
        )
        return tokenizer, model

    def _generate_text(self, prompt: str) -> str:
        tokenizer, model = self._model_bundle
        messages = [
            {
                "role": "system",
                "content": "You are Midnight Static's radio-drama scriptwriter.",
            },
            {"role": "user", "content": prompt},
        ]
        if hasattr(tokenizer, "apply_chat_template"):
            input_ids = tokenizer.apply_chat_template(
                messages,
                add_generation_prompt=True,
                return_tensors="pt",
            )
        else:
            input_ids = tokenizer(prompt, return_tensors="pt").input_ids
        input_ids = input_ids.to(model.device)
        outputs = model.generate(
            input_ids,
            max_new_tokens=self.max_new_tokens,
            do_sample=True,
            temperature=0.75,
            top_p=0.9,
            pad_token_id=tokenizer.eos_token_id,
        )
        generated = outputs[0][input_ids.shape[-1] :]
        return tokenizer.decode(generated, skip_special_tokens=True)


class LlamaCppNemotronWriter:
    """The real Nemotron 3 Nano 4B running **locally on CPU** via llama.cpp.

    Uses the GGUF build (`nemotron_h` architecture) so it needs neither a GPU nor
    the `mamba-ssm` CUDA kernels — which is what lets the live cpu-basic Space run
    the genuine model and keep the off-the-grid claim. Reasoning is disabled
    (the chat template's empty ``<think></think>``) for speed, and a wall-clock
    deadline + token cap bound the latency; on overrun the streamed text is
    returned as-is, so an incomplete script fails validation and
    ``generate_script`` falls back to the fixture.

    Runtime configuration:
    - `MIDNIGHT_GGUF_REPO` / `MIDNIGHT_GGUF_FILE`: override the GGUF source.
    - `MIDNIGHT_GGUF_MAX_TOKENS`: generation cap, default `800`.
    - `MIDNIGHT_GGUF_DEADLINE`: wall-clock budget in seconds, default `150`.
    - `MIDNIGHT_GGUF_THREADS`: llama.cpp threads, default `0` (llama.cpp auto).
    """

    name = "nemotron"

    def __init__(
        self,
        repo_id: str | None = None,
        filename: str | None = None,
        max_tokens: int | None = None,
        deadline_seconds: float | None = None,
    ) -> None:
        self.repo_id = repo_id or os.getenv(
            "MIDNIGHT_GGUF_REPO", "lmstudio-community/NVIDIA-Nemotron-3-Nano-4B-GGUF"
        )
        self.filename = filename or os.getenv(
            "MIDNIGHT_GGUF_FILE", "NVIDIA-Nemotron-3-Nano-4B-Q4_K_M.gguf"
        )
        self.max_tokens = max_tokens or int(os.getenv("MIDNIGHT_GGUF_MAX_TOKENS", "800"))
        self.deadline_seconds = deadline_seconds or float(
            os.getenv("MIDNIGHT_GGUF_DEADLINE", "150")
        )

    @cached_property
    def _llm(self):
        from huggingface_hub import hf_hub_download
        from llama_cpp import Llama

        threads = int(os.getenv("MIDNIGHT_GGUF_THREADS", "0")) or None
        model_path = hf_hub_download(self.repo_id, self.filename)
        return Llama(
            model_path=model_path,
            n_ctx=4096,
            n_threads=threads,
            n_gpu_layers=0,
            verbose=False,
        )

    def write(
        self, premise: str, genre: Genre = Genre.WEIRD, feedback: str | None = None
    ) -> Script:
        import time

        prompt = _nemotron_chatml(
            system="You are Midnight Static's radio-drama scriptwriter. Reply with ONLY the JSON script, no prose.",
            user=build_writer_prompt(premise, genre, feedback=feedback),
        )
        deadline = time.monotonic() + self.deadline_seconds
        chunks: list[str] = []
        for part in self._llm.create_completion(
            prompt,
            max_tokens=self.max_tokens,
            temperature=0.7,
            top_p=0.9,
            stop=["<|im_end|>"],
            stream=True,
        ):
            chunks.append(part["choices"][0]["text"])
            if time.monotonic() > deadline:
                break
        return parse_writer_json("".join(chunks))


def _nemotron_chatml(system: str, user: str) -> str:
    """Reasoning-off ChatML prompt for Nemotron 3 Nano (empty <think></think>)."""
    return (
        f"<|im_start|>system\n{system}<|im_end|>\n"
        f"<|im_start|>user\n{user}<|im_end|>\n"
        f"<|im_start|>assistant\n<think></think>"
    )


def get_writer(name: str = "fixture") -> ScriptWriter:
    if name == "fixture":
        return FixtureWriter()
    if name == "nemotron":
        # Live default: the real model on CPU via llama.cpp (GGUF).
        return LlamaCppNemotronWriter()
    if name == "nemotron-hf":
        # Transformers/GPU path (needs CUDA + mamba-ssm); used off the Space.
        return NemotronWriter()
    raise ValueError(f"unknown writer: {name}")


# Failures we treat as "the model produced an unusable script": malformed JSON,
# missing/extra keys, or cross-field validation errors from the schema.
WRITER_FAILURES = (ValueError, KeyError, TypeError, json.JSONDecodeError)

# Failures that mean the selected writer cannot run in this environment at all —
# e.g. Nemotron on the cpu-basic Space (no GPU, no mamba-ssm CUDA kernels) or a
# missing torch/transformers. Retrying won't help; we drop straight to the
# fixture writer so the station never shows an error page.
WRITER_UNAVAILABLE = (ImportError, RuntimeError, OSError)


def generate_script(
    writer: ScriptWriter,
    premise: str,
    genre: Genre = Genre.WEIRD,
    *,
    retries: int = 1,
    fixture_fallback: bool = True,
) -> Script:
    """Produce a script with one retry on invalid output, then a fixture fallback.

    Mirrors the ARCHITECTURE.md policy: a schema/JSON failure triggers one retry
    with the error fed back to the writer; a second failure improvises with the
    genre fixture (woven around the user's premise) so the station never shows an
    error page.
    """
    last_error: Exception | None = None
    feedback: str | None = None
    for _ in range(retries + 1):
        try:
            return writer.write(premise, genre=genre, feedback=feedback)
        except WRITER_FAILURES as exc:
            last_error = exc
            feedback = _format_feedback(exc)
            print(
                f"[writer] {getattr(writer, 'name', '?')} produced unusable output "
                f"({type(exc).__name__}: {exc}); will retry/fallback",
                file=sys.stderr,
                flush=True,
            )
        except WRITER_UNAVAILABLE as exc:
            # The writer can't load here (e.g. Nemotron on CPU). No point
            # retrying — go straight to the fixture fallback below.
            last_error = exc
            print(
                f"[writer] {getattr(writer, 'name', '?')} unavailable "
                f"({type(exc).__name__}: {exc}); falling back to fixture",
                file=sys.stderr,
                flush=True,
            )
            break
    if fixture_fallback:
        return fixture_script(premise, genre=genre)
    raise last_error  # type: ignore[misc]


def _format_feedback(error: Exception) -> str:
    return (
        "Your previous attempt was rejected: "
        f"{type(error).__name__}: {error}. "
        "Return one corrected JSON object that satisfies every schema rule."
    )


def script_to_json_dict(script: Script) -> dict:
    script.validate()
    return asdict(script)


def parse_writer_json(raw: str) -> Script:
    payload = json.loads(extract_json_object(raw))
    return script_from_json_dict(payload)


def extract_json_object(raw: str) -> str:
    """Extract the first top-level JSON object from model output."""
    text = raw.strip()
    if text.startswith("```"):
        text = re.sub(r"^```(?:json)?\s*", "", text)
        text = re.sub(r"\s*```$", "", text)
    start = text.find("{")
    if start < 0:
        raise ValueError("writer output did not contain a JSON object")
    depth = 0
    in_string = False
    escaped = False
    for index, char in enumerate(text[start:], start=start):
        if escaped:
            escaped = False
            continue
        if char == "\\":
            escaped = True
            continue
        if char == '"':
            in_string = not in_string
            continue
        if in_string:
            continue
        if char == "{":
            depth += 1
        elif char == "}":
            depth -= 1
            if depth == 0:
                return text[start : index + 1]
    raise ValueError("writer output contained incomplete JSON")


def build_writer_prompt(premise: str, genre: Genre, feedback: str | None = None) -> str:
    premise = " ".join(premise.strip().split())[:300] or "a signal with no sender"
    correction = f"\nCorrection required: {feedback}\n" if feedback else ""
    voices = [
        "am_michael",
        "am_onyx",
        "af_bella",
        "am_puck",
        "am_fenrir",
        "bm_george",
        "hf_alpha",
        "hm_omega",
    ]
    deliveries = ["neutral", "slow", "urgent", "whisper", "booming", "deadpan", "agitated"]
    return f"""
Write a 60-90 second vintage radio-drama script as strict JSON only.

Premise: {premise}
Genre: {genre.value}
{correction}
Rules:
- Return one JSON object and no markdown.
- Use exactly the schema keys shown below.
- `genre` and `music.genre` must be "{genre.value}".
- Use 2-4 cast members.
- Cast `voice` must be one of: {", ".join(voices)}.
- Every line `cast` must reference a cast member id.
- Every line `delivery` must be one of: {", ".join(deliveries)}.
- Use 2-3 scenes and 8-18 total dialogue lines.
- Keep content PG-13. Avoid real public figures, explicit gore, and slurs.
- For hindi_melodrama, dialogue must be romanized Hinglish only.

Schema example:
{{
  "title": "The Case of the Vanishing Signal",
  "logline": "A one-sentence hook.",
  "genre": "{genre.value}",
  "cast": [
    {{"id": "announcer", "name": "The Announcer", "voice": "am_michael", "description": "Warm station voice."}},
    {{"id": "caller", "name": "The Caller", "voice": "af_bella", "description": "Nervous lead."}}
  ],
  "scenes": [
    {{
      "title": "Scene title",
      "sfx": ["radio static", "distant thunder"],
      "lines": [
        {{"cast": "announcer", "text": "Line text.", "delivery": "booming"}}
      ]
    }}
  ],
  "music": {{
    "genre": "{genre.value}",
    "opening_sting": "short music prompt",
    "bed": "loopable bed prompt",
    "closing_sting": "short closing prompt"
  }},
  "estimated_seconds": 75
}}
""".strip()


SMOKE_PREMISES = [
    "a phone that only receives calls from 1962",
    "two rival lighthouse keepers, one light",
    "static",
    "the last petrol pump before the salt flats",
    "a wedding where every guest recognizes the wrong bride",
]