#!/usr/bin/env python3
"""Standalone inference for KerzgrAI/PromptTercet-0.5B.

PromptTercet writes the user's side of a chat: give it a conversation and it
asks the next question, give it nothing and it opens one. You, or any
OpenAI-compatible model, write the assistant's side.

Examples:
  python inference.py                                   # it asks, you answer
  python inference.py --system "You are a cooking assistant."
  python inference.py --conversation chat.json          # print the next question
  python inference.py --count 500 --output prompts.jsonl
  python inference.py --count 100 --turns 3 --output chats.jsonl \\
      --answer-url https://openrouter.ai/api/v1 --answer-model openai/gpt-4o-mini
"""

from __future__ import annotations

import argparse
import importlib
import json
import os
import platform
import secrets
import shutil
import subprocess
import sys
import threading
import time
import unicodedata
import urllib.error
import urllib.request
import warnings
from collections.abc import Callable
from dataclasses import dataclass, field
from pathlib import Path
from types import SimpleNamespace
from typing import Any, TextIO

import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from tokenizers import Tokenizer


def _silence_runtime_warnings() -> None:
    patterns = (
        r"tl\.make_block_ptr is deprecated",
        r"Memory efficient kernel not used because",
        r"Memory Efficient attention has been runtime disabled",
        r"Flash attention kernel not used because",
        r"Torch was not compiled with flash attention",
        r"cuDNN attention kernel not used because",
        r"cuDNN attention has been runtime disabled",
    )
    for pattern in patterns:
        warnings.filterwarnings("ignore", message=pattern)


REPO_ID = "KerzgrAI/PromptTercet-0.5B"
MODEL_NAME = "PromptTercet-0.5B"
FLA_COMMIT = "cbb0a72efb55c18ca0ef4f298298317573ad2cb3"
FLA_REPO = "https://github.com/fla-org/flash-linear-attention.git"
FLA_RUNTIME_REQUIREMENT = "transformers>=4.45.0"
# Parallax attention is not in a flash-linear-attention release yet, so fla
# installs from git at the commit the model was trained with.
FLA_MODULES = (
    "fla.layers.gdn2",
    "fla.models.utils",
    "fla.modules.rotary",
    "fla.ops.parallax.parallel",
    "fla.ops.parallax.decode",
)
PATCH_FILES = (
    "fla/__init__.py",
    "fla/ops/__init__.py",
    "fla/layers/__init__.py",
    "fla/ops/simple_gla/__init__.py",
)
TINY_GDN_FILES = (
    "tiny_gdn/__init__.py",
    "tiny_gdn/config.py",
    "tiny_gdn/model.py",
    "tiny_gdn/haiku_layers.py",
    "tiny_gdn/nn_common.py",
    "tiny_gdn/fused_ops.py",
    "tiny_gdn/senryu2.py",
    "tiny_gdn/parallax_ops.py",
    "tiny_gdn/gated_residual_ops.py",
    "tiny_gdn/ngram.py",
    "tiny_gdn/tools.py",
    "tiny_gdn/smoltalk_chat.py",
    "tiny_gdn/chatml.py",
    "tiny_gdn/detokenize.py",
    "tiny_gdn/repetition.py",
)
NGRAM_TABLES_FILE = "ngram_tables.safetensors"
ANSWER_REQUEST_RETRIES = 6
# A conversation's own turns, as PromptTercet's chat template expects them:
# its questions are the template's assistant turns and the replies its user turns.
PROMPTER_ROLES = {"system": "system", "user": "assistant", "assistant": "user"}


@dataclass(frozen=True)
class Source:
    repo_id: str
    local_dir: Path | None
    # A full copy of this repo next to inference.py, used instead of the Hub.
    local_copy: Path | None


def _local_copy() -> Path | None:
    here = Path(__file__).resolve().parent
    config = here / "config.json"
    if not all((here / name).is_file() for name in ("config.json", "model.safetensors", NGRAM_TABLES_FILE)):
        return None
    try:
        name = json.loads(config.read_text(encoding="utf-8")).get("model_name")
    except (OSError, ValueError):
        return None
    return here if name == MODEL_NAME else None


def _download(source: Source, filename: str) -> Path:
    if source.local_copy is not None and (source.local_copy / filename).is_file():
        return source.local_copy / filename
    return Path(
        hf_hub_download(
            repo_id=source.repo_id,
            filename=filename,
            local_dir=str(source.local_dir) if source.local_dir else None,
        )
    )


def _run(cmd: list[str], *, cwd: Path | None = None) -> None:
    print("+", " ".join(cmd), flush=True)
    merged = os.environ.copy()
    merged.setdefault("PYTHONUTF8", "1")
    merged.setdefault("PYTHONIOENCODING", "utf-8")
    subprocess.check_call(cmd, cwd=str(cwd) if cwd else None, env=merged)


def _pip_install(*args: str) -> None:
    _run([sys.executable, "-m", "pip", "install", *args])


def _fla_importable() -> tuple[bool, str]:
    try:
        for name in FLA_MODULES:
            importlib.import_module(name)
    except Exception as error:  # noqa: BLE001
        return False, str(error)
    return True, ""


def _cache_root() -> Path:
    override = os.environ.get("PROMPTTERCET_CACHE")
    path = Path(override).expanduser().resolve() if override else Path.home() / ".cache" / "prompttercet"
    path.mkdir(parents=True, exist_ok=True)
    return path


def _install_fla(source: Source) -> None:
    print("flash-linear-attention missing/broken — installing automatically …", flush=True)
    _pip_install("einops", "numpy", FLA_RUNTIME_REQUIREMENT)
    if platform.system() != "Windows":
        _pip_install("--no-deps", "--force-reinstall", f"git+{FLA_REPO}@{FLA_COMMIT}")
        return
    if shutil.which("git") is None:
        raise RuntimeError("git is required to auto-install flash-linear-attention. Install Git and put it on PATH.")
    fla_root = _cache_root() / "flash-linear-attention"
    if not (fla_root / ".git").is_dir():
        if fla_root.exists():
            shutil.rmtree(fla_root)
        _run(["git", "clone", "--filter=blob:none", FLA_REPO, str(fla_root)])
    _run(["git", "fetch", "--depth", "1", "origin", FLA_COMMIT], cwd=fla_root)
    _run(["git", "checkout", "--force", FLA_COMMIT], cwd=fla_root)
    print("Applying Windows FLA import patches …", flush=True)
    for relative in PATCH_FILES:
        target = fla_root / relative
        target.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(_download(source, f"windows_fla_patches/{relative}"), target)
    _pip_install("--no-build-isolation", "--no-deps", "-e", str(fla_root))


def _ensure_fla(source: Source) -> None:
    ok, error = _fla_importable()
    if ok:
        return
    print(f"FLA not ready ({error})", flush=True)
    try:
        _install_fla(source)
    except Exception as install_error:
        raise RuntimeError(
            "Automatic flash-linear-attention install failed.\n"
            f"Original import error: {error}\n"
            f"Install error: {install_error}"
        ) from install_error
    for name in list(sys.modules):
        if name == "fla" or name.startswith("fla."):
            del sys.modules[name]
    importlib.invalidate_caches()
    ok, error = _fla_importable()
    if not ok:
        raise RuntimeError(f"flash-linear-attention installed but still failed to import: {error}")
    print("flash-linear-attention ready.", flush=True)


def _ensure_tiny_gdn(source: Source) -> Path:
    for root in (Path(__file__).resolve().parent, source.local_copy, source.local_dir):
        if root is not None and all((root / name).is_file() for name in TINY_GDN_FILES):
            return root
    for name in TINY_GDN_FILES:
        _download(source, name)
    return _download(source, "tiny_gdn/__init__.py").parent.parent


def _import_stack() -> SimpleNamespace:
    # tiny_gdn is fetched from the Hub and FLA may be installed at runtime, so
    # these modules only become importable after _ensure_tiny_gdn/_ensure_fla.
    config = importlib.import_module("tiny_gdn.config")
    senryu2 = importlib.import_module("tiny_gdn.senryu2")
    chatml = importlib.import_module("tiny_gdn.chatml")
    detokenize = importlib.import_module("tiny_gdn.detokenize")
    repetition = importlib.import_module("tiny_gdn.repetition")
    return SimpleNamespace(
        TinyGDNConfig=config.TinyGDNConfig,
        Senryu2ForCausalLM=senryu2.Senryu2ForCausalLM,
        encode_prompter_messages=chatml.encode_prompter_messages,
        IncrementalUtf8Decoder=detokenize.IncrementalUtf8Decoder,
        apply_repetition_penalty=repetition.apply_repetition_penalty,
        structural_token_flags=repetition.structural_token_flags,
    )


@dataclass(frozen=True)
class Sampling:
    temperature: float
    top_p: float
    top_k: int
    repetition_penalty: float
    repetition_window: int
    max_new_tokens: int
    context_length: int


@dataclass(frozen=True)
class Question:
    text: str
    tokens: int
    # "stop" ended the turn itself; "length" hit --max-new-tokens and is cut off.
    stop_reason: str

    @property
    def usable(self) -> bool:
        return self.stop_reason == "stop" and bool(self.text)


def _sample(logits: torch.Tensor, sampling: Sampling, generator: torch.Generator) -> int:
    logits = logits.float()
    if sampling.temperature <= 1e-5:
        return int(torch.argmax(logits).item())
    logits = logits / sampling.temperature
    if 0 < sampling.top_k < logits.shape[-1]:
        threshold = torch.topk(logits, sampling.top_k).values[-1]
        logits = logits.masked_fill(logits < threshold, -torch.inf)
    if sampling.top_p < 1.0:
        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
        remove = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1) > sampling.top_p
        remove[1:] = remove[:-1].clone()
        remove[0] = False
        sorted_logits = sorted_logits.masked_fill(remove, -torch.inf)
        logits = torch.full_like(logits, -torch.inf)
        logits.scatter_(0, sorted_indices, sorted_logits)
    probs = torch.softmax(logits, dim=-1).cpu()
    return int(torch.multinomial(probs, 1, generator=generator).item())


def prompter_messages(conversation: list[dict[str, str]]) -> list[dict[str, str]]:
    return [{"role": PROMPTER_ROLES[item["role"]], "content": item["content"]} for item in conversation]


class Prompter:
    """PromptTercet on one GPU. Generations run one at a time, so threads can share it."""

    def __init__(self, model: Any, tokenizer: Tokenizer, stack: SimpleNamespace, device: torch.device, sampling: Sampling, seed: int | None) -> None:
        self.model = model
        self.tokenizer = tokenizer
        self.stack = stack
        self.device = device
        self.sampling = sampling
        self.stop_ids = {int(model.config.eos_token_id)}
        im_end = tokenizer.token_to_id("<|im_end|>")
        if im_end is not None:
            self.stop_ids.add(im_end)
        # The penalty skips whitespace, punctuation and chat tokens, so turns can still end.
        self.structural = stack.structural_token_flags(tokenizer) if sampling.repetition_penalty != 1.0 else ()
        self._seed = seed
        self._calls = 0
        self._lock = threading.Lock()

    def prompt_length(self, conversation: list[dict[str, str]]) -> int:
        return len(self.stack.encode_prompter_messages(self.tokenizer, prompter_messages(conversation)))

    def ask(
        self,
        conversation: list[dict[str, str]],
        *,
        on_text: Callable[[str], None] | None = None,
        stop: threading.Event | None = None,
    ) -> Question:
        prompt_ids = self.stack.encode_prompter_messages(self.tokenizer, prompter_messages(conversation))
        with self._lock:
            if self._seed is None:
                seed = secrets.randbits(63)
            else:
                seed = self._seed + self._calls
                self._calls += 1
            return self._generate(prompt_ids, seed, on_text, stop)

    @torch.inference_mode()
    def _generate(
        self,
        prompt_ids: list[int],
        seed: int,
        on_text: Callable[[str], None] | None,
        stop: threading.Event | None,
    ) -> Question:
        sampling = self.sampling
        token_ids = list(prompt_ids[-sampling.context_length :])
        generated = 0
        decoder = self.stack.IncrementalUtf8Decoder(self.tokenizer)
        generator = torch.Generator(device="cpu")
        generator.manual_seed(seed)
        stop_reason = "length"
        try:
            output = self.model(
                torch.tensor([token_ids], dtype=torch.long, device=self.device),
                return_logits=True,
                logits_to_keep=1,
                use_cache=True,
            )
            for step in range(sampling.max_new_tokens):
                if stop is not None and stop.is_set():
                    stop_reason = "interrupted"
                    break
                if output.logits is None:
                    raise RuntimeError("Model returned no logits")
                logits = output.logits[0, -1].float()
                if self.structural:
                    self.stack.apply_repetition_penalty(
                        logits,
                        token_ids,
                        penalty=sampling.repetition_penalty,
                        window=sampling.repetition_window,
                        structural=self.structural,
                    )
                next_id = _sample(logits, sampling, generator)
                if next_id in self.stop_ids:
                    stop_reason = "stop"
                    break
                token_ids.append(next_id)
                generated += 1
                delta = decoder.push(next_id)
                if delta and on_text is not None:
                    on_text(delta)
                if step + 1 == sampling.max_new_tokens:
                    break
                if output.past_key_values is None:
                    raise RuntimeError("Model returned no decode cache")
                if output.past_key_values.length + 1 > self.model.config.max_position_embeddings:
                    stop_reason = "context_full"
                    break
                output = self.model(
                    torch.tensor([[next_id]], dtype=torch.long, device=self.device),
                    return_logits=True,
                    logits_to_keep=1,
                    past_key_values=output.past_key_values,
                    use_cache=True,
                )
        except KeyboardInterrupt:
            stop_reason = "interrupted"
        text = decoder.finalize().strip()
        return Question(text=text, tokens=generated, stop_reason=stop_reason)


class AnswerError(RuntimeError):
    pass


@dataclass(frozen=True)
class AnswerModel:
    """An OpenAI-compatible chat completions endpoint that writes the assistant turns."""

    base_url: str
    model: str
    api_key: str
    max_tokens: int
    temperature: float
    extra: dict[str, Any]
    timeout: float

    def reply(self, messages: list[dict[str, str]], stop: threading.Event) -> tuple[str, str]:
        """``(content, finish_reason)`` of one completion, retrying transient failures."""

        body = {
            **self.extra,
            "model": self.model,
            "messages": messages,
            "max_tokens": self.max_tokens,
            "temperature": self.temperature,
        }
        headers = {"Content-Type": "application/json", "User-Agent": "PromptTercet-0.5B/inference.py"}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"
        url = f"{self.base_url.rstrip('/')}/chat/completions"
        last = "no attempt"
        for attempt in range(ANSWER_REQUEST_RETRIES):
            if stop.is_set():
                raise AnswerError("stopped")
            delay = min(2.0**attempt, 30.0)
            try:
                request = urllib.request.Request(url, data=json.dumps(body).encode("utf-8"), headers=headers, method="POST")
                with urllib.request.urlopen(request, timeout=self.timeout) as response:
                    payload = json.loads(response.read().decode("utf-8"))
            except urllib.error.HTTPError as error:
                detail = error.read().decode("utf-8", errors="replace")[:300]
                last = f"HTTP {error.code}: {detail}"
                if error.code not in {408, 409, 425, 429} and error.code < 500:
                    raise AnswerError(last) from error
                retry_after = error.headers.get("Retry-After") if error.headers else None
                if retry_after and retry_after.isdigit():
                    delay = min(float(retry_after), 60.0)
            except (urllib.error.URLError, TimeoutError, ConnectionError, OSError, ValueError) as error:
                last = f"{type(error).__name__}: {error}"
            else:
                if isinstance(payload, dict) and payload.get("error"):
                    last = f"API error: {json.dumps(payload['error'])[:300]}"
                else:
                    try:
                        choice = payload["choices"][0]
                        content = choice["message"].get("content") or ""
                    except (KeyError, IndexError, TypeError, AttributeError) as error:
                        raise AnswerError(f"Unexpected response: {json.dumps(payload)[:300]}") from error
                    return str(content).strip(), str(choice.get("finish_reason") or "")
            if stop.wait(delay):
                raise AnswerError("stopped")
        raise AnswerError(f"{ANSWER_REQUEST_RETRIES} requests failed; last: {last}")


def opening_key(text: str) -> str:
    return unicodedata.normalize("NFC", text).strip()


def load_conversation(path: str) -> list[dict[str, str]]:
    raw = sys.stdin.read() if path == "-" else Path(path).read_text(encoding="utf-8")
    payload = json.loads(raw)
    messages = payload.get("messages") if isinstance(payload, dict) else payload
    if not isinstance(messages, list):
        raise ValueError("A conversation is a JSON list of messages or an object with a 'messages' list")
    conversation: list[dict[str, str]] = []
    for index, item in enumerate(messages):
        if not isinstance(item, dict) or item.get("role") not in PROMPTER_ROLES:
            raise ValueError(f"Message {index} needs a role of system, user or assistant")
        content = item.get("content")
        if not isinstance(content, str) or not content.strip():
            raise ValueError(f"Message {index} needs non-empty text content")
        if item["role"] == "system" and index != 0:
            raise ValueError("Only the first message can be a system prompt")
        conversation.append({"role": item["role"], "content": content})
    if conversation and conversation[-1]["role"] == "user":
        raise ValueError("The conversation ends with a user turn. PromptTercet writes the user turn after an assistant reply.")
    return conversation


class Style:
    def __init__(self, enabled: bool) -> None:
        self.enabled = enabled

    def wrap(self, code: str, text: str) -> str:
        return f"\033[{code}m{text}\033[0m" if self.enabled else text

    def user(self, text: str) -> str:
        return self.wrap("1;36", text)

    def assistant(self, text: str) -> str:
        return self.wrap("1;32", text)

    def note(self, text: str) -> str:
        return self.wrap("2", text)


INTERACTIVE_HELP = (
    "Commands: /retry asks again, /system TEXT sets the system prompt and starts over, "
    "/reset starts over, /save FILE appends the chat as a JSON line, /exit quits."
)


def run_interactive(prompter: Prompter, system: str, style: Style) -> int:
    print(
        f"{MODEL_NAME} plays the user and you play the assistant: answer each question it asks.\n"
        f"{INTERACTIVE_HELP} Ctrl+C stops a question.",
        flush=True,
    )

    def fresh(prompt: str) -> list[dict[str, str]]:
        return [{"role": "system", "content": prompt}] if prompt.strip() else []

    conversation = fresh(system)
    ask_now = True
    while True:
        if ask_now:
            ask_now = False
            if prompter.prompt_length(conversation) > prompter.sampling.context_length:
                print(style.note(f"(the chat is over {prompter.sampling.context_length} tokens; the oldest are dropped)"))
            print(f"\n{style.user('user>')} ", end="", flush=True)
            try:
                question = prompter.ask(conversation, on_text=lambda delta: print(delta, end="", flush=True))
            except KeyboardInterrupt:
                question = Question(text="", tokens=0, stop_reason="interrupted")
            print(flush=True)
            if question.stop_reason == "interrupted":
                print(style.note("(stopped; /retry asks again)"))
            elif not question.text:
                print(style.note("(empty question; /retry asks again)"))
            else:
                if question.stop_reason == "length":
                    print(style.note(f"(cut off at --max-new-tokens {prompter.sampling.max_new_tokens}; /retry asks again)"))
                conversation.append({"role": "user", "content": question.text})
        try:
            line = input(f"{style.assistant('assistant>')} ")
        except (EOFError, KeyboardInterrupt):
            print()
            return 0
        text = line.strip()
        if not text:
            continue
        if text.startswith("/"):
            command, _, rest = text.partition(" ")
            if command in {"/exit", "/quit"}:
                return 0
            if command == "/help":
                print(INTERACTIVE_HELP)
            elif command == "/reset":
                conversation = fresh(system)
                ask_now = True
                print(style.note("(new chat)"))
            elif command == "/system":
                system = rest.strip()
                conversation = fresh(system)
                ask_now = True
                print(style.note("(system prompt set; new chat)" if system else "(system prompt cleared; new chat)"))
            elif command == "/retry":
                if conversation and conversation[-1]["role"] == "user":
                    conversation.pop()
                ask_now = True
            elif command == "/save":
                path = Path(rest.strip() or "conversations.jsonl")
                if not any(item["role"] == "user" for item in conversation):
                    print(style.note("(nothing to save yet)"))
                    continue
                with path.open("a", encoding="utf-8") as handle:
                    handle.write(json.dumps({"messages": conversation}, ensure_ascii=False) + "\n")
                print(style.note(f"(appended to {path})"))
            else:
                print(style.note(f"(unknown command {command}; /help lists them)"))
            continue
        if not conversation or conversation[-1]["role"] != "user":
            print(style.note("(no question to answer; /retry asks again)"))
            continue
        conversation.append({"role": "assistant", "content": text})
        ask_now = True


def run_next_question(prompter: Prompter, conversation: list[dict[str, str]]) -> int:
    question = prompter.ask(conversation)
    if question.stop_reason == "interrupted":
        return 130
    print(question.text, flush=True)
    if not question.usable:
        print(f"[warn] question ended with {question.stop_reason}", file=sys.stderr, flush=True)
        return 1
    return 0


@dataclass
class BatchState:
    count: int
    written: int = 0
    in_flight: int = 0
    failed: int = 0
    redone_questions: int = 0
    repeated_openings: int = 0
    redone_answers: int = 0
    last_error: str = ""
    seen: set[str] = field(default_factory=set)
    lock: threading.Lock = field(default_factory=threading.Lock)

    def claim(self) -> bool:
        with self.lock:
            if self.written + self.in_flight >= self.count:
                return False
            self.in_flight += 1
            return True

    def bump(self, name: str) -> None:
        with self.lock:
            setattr(self, name, getattr(self, name) + 1)

    def note(self, error: str) -> None:
        with self.lock:
            self.last_error = error

    def finish(self, ok: bool, stopping: bool) -> bool:
        """Settle one claimed chat; True when failures say the batch should stop."""

        with self.lock:
            self.in_flight -= 1
            if ok:
                self.written += 1
            elif not stopping:
                self.failed += 1
            return self.failed >= 20 and self.failed > self.written

    def summary(self) -> str:
        parts = [f"{self.written}/{self.count} written"]
        if self.failed:
            parts.append(f"{self.failed} failed")
        if self.repeated_openings:
            parts.append(f"{self.repeated_openings} repeated openings resampled")
        if self.redone_questions:
            parts.append(f"{self.redone_questions} cut-off questions resampled")
        if self.redone_answers:
            parts.append(f"{self.redone_answers} cut-off answers redone")
        return ", ".join(parts)


def seen_openings(path: Path) -> set[str]:
    seen: set[str] = set()
    if not path.is_file():
        return seen
    with path.open(encoding="utf-8") as handle:
        for line in handle:
            try:
                messages = json.loads(line).get("messages", [])
            except (ValueError, AttributeError):
                continue
            first = next((item for item in messages if isinstance(item, dict) and item.get("role") == "user"), None)
            if first and isinstance(first.get("content"), str):
                seen.add(opening_key(first["content"]))
    return seen


def run_batch(prompter: Prompter, args: argparse.Namespace, answer: AnswerModel | None) -> int:
    output_path = Path(args.output) if args.output else None
    state = BatchState(count=args.count)
    if output_path is not None:
        state.seen = seen_openings(output_path)
        if state.seen:
            print(f"{output_path} already has {len(state.seen)} openings; new ones will differ.", file=sys.stderr)
    sink: TextIO = output_path.open("a", encoding="utf-8") if output_path else sys.stdout
    write_lock = threading.Lock()
    stop = threading.Event()
    system = args.system.strip()
    answer_system = system if args.answer_system is None else args.answer_system.strip()

    def ask(conversation: list[dict[str, str]], opening: bool) -> str | None:
        for _ in range(args.attempts):
            question = prompter.ask(conversation, stop=stop)
            if question.stop_reason == "interrupted" or stop.is_set():
                return None
            if not question.usable:
                state.bump("redone_questions")
                continue
            if opening:
                key = opening_key(question.text)
                with state.lock:
                    if key in state.seen:
                        state.repeated_openings += 1
                        continue
                    state.seen.add(key)
            return question.text
        state.note(f"no usable question in {args.attempts} tries")
        return None

    def reply(conversation: list[dict[str, str]]) -> str | None:
        assert answer is not None
        messages = ([{"role": "system", "content": answer_system}] if answer_system else []) + [
            item for item in conversation if item["role"] != "system"
        ]
        for _ in range(args.attempts):
            try:
                content, finish = answer.reply(messages, stop)
            except AnswerError as error:
                if not stop.is_set():
                    state.note(f"answer model: {error}")
                return None
            if content and finish != "length":
                return content
            state.bump("redone_answers")
        state.note(f"{args.attempts} answers in a row hit --answer-max-tokens {args.answer_max_tokens} or came back empty")
        return None

    def conversation_once() -> bool:
        """One chat of exactly --turns questions, written only when complete."""

        conversation = [{"role": "system", "content": system}] if system else []
        for turn in range(args.turns):
            question = ask(conversation, opening=turn == 0)
            if question is None:
                return False
            conversation.append({"role": "user", "content": question})
            if answer is not None:
                text = reply(conversation)
                if text is None:
                    return False
                conversation.append({"role": "assistant", "content": text})
        if stop.is_set():
            return False
        with write_lock:
            sink.write(json.dumps({"messages": conversation}, ensure_ascii=False) + "\n")
            sink.flush()
        return True

    def worker() -> None:
        while not stop.is_set() and state.claim():
            ok = False
            try:
                ok = conversation_once()
            except Exception as error:  # noqa: BLE001
                state.note(f"{type(error).__name__}: {error}")
            if state.finish(ok, stop.is_set()):
                state.note(f"stopped after {state.failed} failed chats; last: {state.last_error}")
                stop.set()

    workers = [threading.Thread(target=worker, daemon=True) for _ in range(max(1, args.workers))]
    for thread in workers:
        thread.start()
    tty = sys.stderr.isatty()
    shown = ""
    shown_at = 0.0
    interrupted = False
    try:
        while any(thread.is_alive() for thread in workers):
            for thread in workers:
                thread.join(timeout=0.5)
            line = state.summary()
            if line != shown and (tty or time.monotonic() - shown_at >= 15):
                shown, shown_at = line, time.monotonic()
                print(f"\r[{MODEL_NAME}] {line}\033[K" if tty else f"[{MODEL_NAME}] {line}", end="" if tty else "\n", file=sys.stderr, flush=True)
        if tty:
            print(file=sys.stderr)
    except KeyboardInterrupt:
        interrupted = True
        stop.set()
        print("\nStopping after the requests in flight (Ctrl+C again quits now) …", file=sys.stderr, flush=True)
        try:
            for thread in workers:
                thread.join()
        except KeyboardInterrupt:
            pass
    finally:
        if output_path is not None:
            sink.close()
    print(f"[{MODEL_NAME}] done: {state.summary()}", file=sys.stderr, flush=True)
    if state.last_error and state.written < state.count:
        print(f"[{MODEL_NAME}] last failure: {state.last_error}", file=sys.stderr, flush=True)
    if interrupted:
        return 130
    return 0 if state.written >= state.count else 1


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=f"{MODEL_NAME}: writes the user's side of a chat",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    mode = parser.add_argument_group("what to do (default: an interactive chat where you answer)")
    mode.add_argument("--system", default="", help="System prompt of the chat; it steers what PromptTercet asks about")
    mode.add_argument("--conversation", default=None, help="JSON chat (list of messages, or {messages: [...]}); print the next question. '-' reads stdin")
    mode.add_argument("--count", type=int, default=0, help="Batch mode: write this many chats as JSON lines")
    mode.add_argument("--turns", type=int, default=1, help="Questions per chat in batch mode; above 1 needs --answer-url")
    mode.add_argument("--output", default=None, help="JSONL file for batch mode (default: stdout). Appends, and new openings differ from those already in it")
    mode.add_argument("--attempts", type=int, default=8, help="Samples per turn before a chat is given up: cut-off or repeated questions, cut-off answers")
    mode.add_argument("--workers", type=int, default=None, help="Chats in flight in batch mode (default: 8 with an answer model, else 1)")

    answers = parser.add_argument_group("answer model for batch mode (any OpenAI-compatible API)")
    answers.add_argument("--answer-url", default=None, help="Base URL, e.g. https://openrouter.ai/api/v1 or http://localhost:8000/v1")
    answers.add_argument("--answer-model", default=None, help="Model name sent to the API")
    answers.add_argument("--answer-key-env", default="OPENAI_API_KEY", help="Environment variable holding the API key (unset sends none)")
    answers.add_argument("--answer-system", default=None, help="System prompt only the answer model sees, not saved (default: --system)")
    answers.add_argument("--answer-max-tokens", type=int, default=4096, help="Answers that hit this are discarded and redone")
    answers.add_argument("--answer-temperature", type=float, default=0.7)
    answers.add_argument("--answer-extra", default="{}", help='Extra request fields as JSON, e.g. \'{"reasoning": {"effort": "low"}}\'')
    answers.add_argument("--answer-timeout", type=float, default=600.0, help="Seconds per request")

    sampling = parser.add_argument_group("PromptTercet sampling")
    sampling.add_argument("--temperature", type=float, default=0.6)
    sampling.add_argument("--top-p", type=float, default=0.95)
    sampling.add_argument("--top-k", type=int, default=0, help="0 turns top-k off")
    sampling.add_argument(
        "--repetition-penalty",
        type=float,
        default=1.36,
        help="Penalty on words already used twice in the chat, which keeps follow-ups from repeating it (1 turns it off)",
    )
    sampling.add_argument("--repetition-window", type=int, default=0, help="Tokens the penalty looks back over (0: the whole chat)")
    sampling.add_argument("--max-new-tokens", type=int, default=1024, help="Questions that hit this are cut off")
    sampling.add_argument("--context-length", type=int, default=8192, help="Chat tokens kept, oldest dropped first (up to 32768)")
    sampling.add_argument("--seed", type=int, default=None, help="Fix the sampling seed (default: a fresh seed per question)")

    runtime = parser.add_argument_group("runtime")
    runtime.add_argument("--device", default="cuda", help=f"CUDA device (cuda, cuda:1, …); {MODEL_NAME} needs a GPU")
    runtime.add_argument("--ngram-device", default=None, help="Where the 2.2 GB n-gram tables live: the GPU (default, --device) or cpu")
    runtime.add_argument("--no-color", action="store_true")
    runtime.add_argument("--local-dir", default=None, help="Download the model files into this folder")
    runtime.add_argument("--repo-id", default=REPO_ID)
    return parser.parse_args()


def validate(args: argparse.Namespace) -> str | None:
    if args.count < 0 or args.turns < 1 or args.attempts < 1:
        return "--count must be 0 or more, --turns and --attempts at least 1"
    if args.workers is not None and args.workers < 1:
        return "--workers must be at least 1"
    if args.count and args.conversation:
        return "Use --count or --conversation, not both"
    if bool(args.answer_url) != bool(args.answer_model):
        return "--answer-url and --answer-model go together"
    if args.answer_url and not args.count:
        return "--answer-url is for batch mode; add --count"
    if args.answer_url and not args.answer_url.startswith(("http://", "https://")):
        return "--answer-url must start with http:// or https://"
    if args.answer_max_tokens < 1 or args.answer_timeout <= 0:
        return "--answer-max-tokens and --answer-timeout must be positive"
    if args.count and args.turns > 1 and not args.answer_url:
        return "--turns above 1 needs an answer model: add --answer-url and --answer-model"
    if not 1 <= args.max_new_tokens <= 32768 or not 256 <= args.context_length <= 32768:
        return "--max-new-tokens must be 1-32768 and --context-length 256-32768"
    if args.temperature < 0 or not 0 < args.top_p <= 1 or args.top_k < 0 or args.repetition_penalty <= 0:
        return "Sampling needs --temperature >= 0, 0 < --top-p <= 1, --top-k >= 0 and --repetition-penalty > 0"
    if args.repetition_window < 0:
        return "--repetition-window must be 0 (the whole chat) or more"
    try:
        extra = json.loads(args.answer_extra)
    except json.JSONDecodeError as error:
        return f"--answer-extra is not JSON: {error}"
    if not isinstance(extra, dict):
        return "--answer-extra must be a JSON object"
    return None


def main() -> int:
    _silence_runtime_warnings()
    args = parse_args()
    problem = validate(args)
    if problem:
        print(problem, file=sys.stderr)
        return 2
    conversation: list[dict[str, str]] | None = None
    if args.conversation is not None:
        try:
            conversation = load_conversation(args.conversation)
        except (OSError, ValueError) as error:
            print(f"Cannot read --conversation: {error}", file=sys.stderr)
            return 2
        if args.system.strip() and not (conversation and conversation[0]["role"] == "system"):
            conversation.insert(0, {"role": "system", "content": args.system.strip()})

    try:
        device = torch.device(args.device)
        ngram_device = torch.device(args.ngram_device) if args.ngram_device else device
    except RuntimeError as error:
        print(f"Invalid device: {error}", file=sys.stderr)
        return 2
    if device.type != "cuda" or not torch.cuda.is_available():
        print(
            f"{MODEL_NAME} needs a CUDA GPU: its Gated DeltaNet-2 and Parallax layers run on "
            f"flash-linear-attention Triton kernels. torch {torch.__version__} reports "
            f"CUDA available={torch.cuda.is_available()}.",
            file=sys.stderr,
        )
        return 1
    if ngram_device.type not in {"cuda", "cpu"}:
        print("--ngram-device must be a CUDA device or cpu", file=sys.stderr)
        return 2

    source = Source(
        repo_id=args.repo_id,
        local_dir=Path(args.local_dir).resolve() if args.local_dir else None,
        local_copy=_local_copy() if args.repo_id == REPO_ID else None,
    )
    where = source.local_copy or f"huggingface.co/{source.repo_id}"
    print(f"Loading {MODEL_NAME} from {where} …", file=sys.stderr, flush=True)
    try:
        package_root = _ensure_tiny_gdn(source)
        if str(package_root) not in sys.path:
            sys.path.insert(0, str(package_root))
        _ensure_fla(source)
        stack = _import_stack()
    except Exception as error:  # noqa: BLE001
        print(f"Setup failed: {error}", file=sys.stderr)
        return 1

    config = stack.TinyGDNConfig.from_json(_download(source, "config.json"))
    with torch.device(device):
        model = stack.Senryu2ForCausalLM(config)
    model = model.to(dtype=torch.bfloat16)
    model.load_state_dict(load_file(str(_download(source, "model.safetensors")), device=str(device)), strict=True)
    if model.ngram is not None:
        model.ngram.load_tables(_download(source, NGRAM_TABLES_FILE), device=ngram_device, dtype=torch.bfloat16)
    model.eval()
    model.requires_grad_(False)
    tokenizer = Tokenizer.from_file(str(_download(source, "tokenizer.json")))
    sampling = Sampling(
        temperature=args.temperature,
        top_p=args.top_p,
        top_k=args.top_k,
        repetition_penalty=args.repetition_penalty,
        repetition_window=args.repetition_window,
        max_new_tokens=args.max_new_tokens,
        context_length=min(args.context_length, config.max_position_embeddings),
    )
    prompter = Prompter(model, tokenizer, stack, device, sampling, args.seed)
    print(
        f"Ready: {config.num_hidden_layers} layers on {device}, n-gram tables on {ngram_device}.",
        file=sys.stderr,
        flush=True,
    )

    if conversation is not None:
        return run_next_question(prompter, conversation)
    if args.count:
        answer = None
        if args.answer_url:
            answer = AnswerModel(
                base_url=args.answer_url,
                model=args.answer_model,
                api_key=os.environ.get(args.answer_key_env, "").strip(),
                max_tokens=args.answer_max_tokens,
                temperature=args.answer_temperature,
                extra=json.loads(args.answer_extra),
                timeout=args.answer_timeout,
            )
        if args.workers is None:
            args.workers = 8 if answer else 1
        return run_batch(prompter, args, answer)
    style = Style(not args.no_color and sys.stdout.isatty())
    return run_interactive(prompter, args.system, style)


if __name__ == "__main__":
    raise SystemExit(main())