"""
model/backend.py — language-model abstraction for the affect judgments.

Backends (STORY_SHAPES_BACKEND env var, default "llamacpp"):
  - "llamacpp"  : recommended; runs locally and on a GPU Space. A quantized GGUF
                  run in-process via llama-cpp-python (the llama.cpp runtime) —
                  no `transformers`. Structured output is enforced by a GBNF
                  grammar built from the JSON schema, so JSON is guaranteed
                  valid. ~5 GB at Q4_K_M, so it fits beside the FLUX painter on
                  one 24 GB GPU.
  - "modal_llm" : HTTP POST to a deployed Modal GPU endpoint (modal_llm.py),
                  where `transformers` is pinned to 4.x (the version MiniCPM4.1
                  needs). Use when the Space itself has no GPU.

(The painter is independent — see model/painter.py / STORY_SHAPES_PAINT_BACKEND.)

MiniCPM4.1-8B uses the <think> tag convention; thinking is suppressed by the
grammar (JSON paths) or a /no_think system line (prose paths), and stripped
defensively either way.
"""
import os, json, re, urllib.request

import spaces

import logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s  %(levelname)s  %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger("story_shapes")

BACKEND = os.environ.get("STORY_SHAPES_BACKEND", "llamacpp")
LLM_URL = os.environ.get("STORY_SHAPES_LLM_MODAL_URL", "")   # Modal LLM endpoint URL

# llama.cpp backend: a quantized GGUF run in-process via llama-cpp-python.
LLAMACPP_REPO = os.environ.get("STORY_SHAPES_LLAMACPP_REPO", "openbmb/MiniCPM4.1-8B-GGUF")
LLAMACPP_FILE = os.environ.get("STORY_SHAPES_LLAMACPP_FILE", "*Q4_K_M.gguf")
LLAMACPP_CTX  = int(os.environ.get("STORY_SHAPES_LLAMACPP_CTX", "4096"))
# -1 offloads all layers to GPU; set 0 for CPU-only.
LLAMACPP_GPU_LAYERS = int(os.environ.get("STORY_SHAPES_LLAMACPP_GPU_LAYERS", "-1"))

# Active model identifier, for /health and the eval harness (informational —
# the Modal endpoint pins its own model id internally).
MODEL = os.environ.get("STORY_SHAPES_MODEL", LLAMACPP_REPO)

PROMPT_MARKDOWN_DIVIDER = "================================================================================"

# ---- JSON schemas (source for the llama.cpp GBNF / schema-constrained output) ----
UNIT = {"type": "number", "minimum": 0, "maximum": 1}
MATERIAL_ENUM = {"type": "string", "enum": ["paper", "ink", "glass", "enamel", "chalk", "metal"]}
CORE_SCHEMA = {
    "type": "object",
    "properties": {
        "valence": UNIT, "arousal": UNIT, "dominance": UNIT, "coherence": UNIT,
        "deserves_shape": {"type": "boolean"},
        "label": {"type": "string"}, "comment": {"type": "string"},
        "material": MATERIAL_ENUM,
    },
    "required": ["valence", "arousal", "dominance", "coherence",
                 "deserves_shape", "label", "comment", "material"],
}
SEGMENT_SCHEMA = {
    "type": "object",
    "properties": {
        "segments": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "phrase":    {"type": "string"},
                    "valence":   UNIT,
                    "arousal":   UNIT,
                    "dominance": UNIT,
                    "words":     {"type": "integer"},
                },
                "required": ["phrase", "valence", "arousal", "dominance", "words"]
            }
        },
        "coherence":      UNIT,
        "deserves_shape": {"type": "boolean"},
        "label":          {"type": "string"},
        "comment":        {"type": "string"},
        "material":       MATERIAL_ENUM,
    },
    "required": ["segments", "coherence", "deserves_shape", "label", "comment", "material"],
}
TARGET_SCHEMA = {
    "type": "object",
    "properties": {"valence": UNIT, "arousal": UNIT, "dominance": UNIT,
                   "label": {"type": "string"}, "comment": {"type": "string"}},
    "required": ["valence", "arousal", "dominance", "label", "comment"],
}
ATTEMPT_SCHEMA = {
    "type": "object",
    "properties": {"valence": UNIT, "arousal": UNIT, "dominance": UNIT,
                   "comment": {"type": "string"}},
    "required": ["valence", "arousal", "dominance", "comment"],
}
REVEAL_SCHEMA = {
    "type": "object",
    "properties": {
        "coherence": UNIT,
        "companions": {"type": "array", "items": {
            "type": "object",
            "properties": {"valence": UNIT, "arousal": UNIT, "dominance": UNIT,
                           "label": {"type": "string"}},
            "required": ["valence", "arousal", "dominance", "label"]}},
        "layout": {"type": "array", "items": {
            "type": "object",
            "properties": {"ref": {"type": "string"},
                           "region": {"type": "string"},
                           "scale": UNIT,
                           "rotation": {"type": "integer"},
                           "depth": {"type": "integer"}},
            "required": ["ref", "region", "scale", "rotation", "depth"]}},
    },
    "required": ["coherence", "companions", "layout"],
}
# Tiny wrappers so the prose calls (continuation, title) also go through the
# structured-output path — this is what keeps model "thinking" from leaking:
# llama.cpp's grammar forbids a <think> block outright, and for Modal the JSON
# object is extracted even if the model emits one first.
SENTENCE_SCHEMA = {
    "type": "object",
    "properties": {"sentence": {"type": "string"}},
    "required": ["sentence"],
}
TITLE_SCHEMA = {
    "type": "object",
    "properties": {"title": {"type": "string"}},
    "required": ["title"],
}

def _load_prompt(name):
    here = os.path.dirname(os.path.dirname(__file__))
    with open(os.path.join(here, "prompts", name), encoding="utf-8") as f:
        return f.read()

# extract just the system prompt block from a prompt .md (between the SYSTEM markers)
def _system_block(md, marker="SYSTEM PROMPT"):
    if marker in md:
        after = md.split(marker, 1)[1]
        return after.split(PROMPT_MARKDOWN_DIVIDER, 2)[1].strip()
    return md

# ---------------------------------------------------------------------------
# Shared JSON helpers (used by both backends)
# ---------------------------------------------------------------------------
def _try_unescape(candidate: str):
    """Recover an over-escaped JSON object — some models emit the whole object
    with literal \\n and \\" as if it were a quoted string. Returns the parsed
    dict, or None if neither recovery works."""
    for attempt in (
        lambda c: json.loads(json.loads('"' + c + '"')),               # body is an escaped JSON string
        lambda c: json.loads(c.encode("utf-8").decode("unicode_escape")),  # generic unescape
    ):
        try:
            return attempt(candidate)
        except Exception:
            continue
    return None

def _extract_json(text: str) -> dict:
    """Extract the first {...} JSON object, handling markdown fences, trailing
    commas, and over-escaped output (literal \\n / \\" used as formatting) —
    the most common model formatting mistakes."""
    text = re.sub(r"```(?:json)?", "", text).strip()
    m = re.search(r"\{.*\}", text, re.DOTALL)
    if not m:
        raise ValueError(f"No JSON object found in output: {text!r}")
    candidate = m.group(0)
    candidate = re.sub(r",\s*([}\]])", r"\1", candidate)  # trailing commas
    try:
        return json.loads(candidate)
    except json.JSONDecodeError:
        fixed = _try_unescape(candidate)
        if fixed is not None:
            return fixed
        return json.loads(candidate)  # re-raise the original error with context

def _schema_hint(schema: dict) -> str:
    """A plain-language key list — NOT a JSON skeleton. (A skeleton like
    {"keys": {...}} gets copied verbatim by the model, nesting its real answer
    under a stray wrapper.)"""
    parts = []
    for k in schema.get("required", []):
        prop = schema.get("properties", {}).get(k, {})
        t = f"one of {prop['enum']}" if "enum" in prop else prop.get("type", "value")
        parts.append(f"{k} ({t})")
    return ", ".join(parts)

def _unwrap(parsed: dict, schema: dict) -> dict:
    """If the model nested its answer one level deep (e.g. under "keys" or the
    schema's own name) so none of the required keys are at the top level, dig
    into the single child object that does carry them."""
    req = schema.get("required", [])
    if isinstance(parsed, dict) and req and not any(k in parsed for k in req):
        for v in parsed.values():
            if isinstance(v, dict) and any(k in v for k in req):
                return v
    return parsed

def _apply_defaults(parsed: dict, schema: dict) -> dict:
    """Fill missing required fields; clamp unit floats; validate enum strings."""
    unit_keys = {k for k, v in schema.get("properties", {}).items()
                 if isinstance(v, dict) and v.get("minimum") == 0 and v.get("maximum") == 1}
    for k in unit_keys:
        if k in parsed:
            parsed[k] = max(0.0, min(1.0, float(parsed[k])))
    for k in schema.get("required", []):
        prop = schema.get("properties", {}).get(k, {})
        if k not in parsed:
            if prop.get("type") == "boolean": parsed[k] = False
            elif "enum" in prop: parsed[k] = prop["enum"][0]
            elif prop.get("type") == "string": parsed[k] = ""
            else: parsed[k] = 0.5
        elif "enum" in prop and parsed[k] not in prop["enum"]:
            parsed[k] = prop["enum"][0]   # coerce invalid enum value to first valid
    return parsed

# ---------------------------------------------------------------------------
# Modal LLM backend — HTTP POST to the deployed modal_llm.py endpoint
# ---------------------------------------------------------------------------
def _modal_generate(messages, max_new_tokens=1024, temperature=0.6) -> str:
    if not LLM_URL:
        raise RuntimeError(
            "STORY_SHAPES_LLM_MODAL_URL is not set. "
            "Run `modal deploy modal_llm.py` and set the printed URL."
        )
    body = json.dumps({
        "messages": messages,
        "max_new_tokens": max_new_tokens,
        "temperature": temperature,
    }).encode()
    req = urllib.request.Request(
        LLM_URL, data=body,
        headers={"Content-Type": "application/json"}, method="POST",
    )
    log.info("modal_llm →  max_tokens=%d  temp=%.2f", max_new_tokens, temperature)
    with urllib.request.urlopen(req, timeout=180) as r:
        resp = json.loads(r.read())
    return resp["text"]

def _modal_llm_chat(system: str, user: str, schema: dict, temperature: float = 0.6) -> dict:
    """Generate JSON via the Modal LLM endpoint; parse, validate, retry up to 3×."""
    hint = _schema_hint(schema)
    messages = [
        {"role": "system", "content": system},
        {"role": "user",   "content": f"{user}\n\nReturn ONLY a flat JSON object with these keys: {hint}"},
    ]
    for attempt in range(3):
        raw = _modal_generate(messages, temperature=temperature)
        log.info("Modal LLM raw (attempt %d) <- %s", attempt + 1, raw[:400])
        try:
            parsed = _apply_defaults(_unwrap(_extract_json(raw), schema), schema)
            log.info("Modal LLM json -> %s", parsed)
            return parsed
        except (ValueError, json.JSONDecodeError, KeyError) as e:
            log.warning("Modal LLM parse attempt %d failed: %s", attempt + 1, e)
            if attempt == 2:
                raise RuntimeError(
                    f"Modal LLM failed to produce valid JSON after 3 attempts. "
                    f"Last output: {raw!r}") from e

# ---------------------------------------------------------------------------
# llama.cpp backend — quantized GGUF in-process via llama-cpp-python
# ---------------------------------------------------------------------------
_llama = None

def _get_llama():
    global _llama
    if _llama is not None:
        return _llama
    # Import torch first: it loads the CUDA runtime libraries (libcudart, etc.)
    # into the process, which the prebuilt llama-cpp-python CUDA wheel's
    # libllama.so needs to resolve when it loads. Harmless if torch is absent.
    try:
        import torch  # noqa: F401
    except Exception:
        pass
    from llama_cpp import Llama
    log.info("loading llama.cpp model %s / %s …", LLAMACPP_REPO, LLAMACPP_FILE)
    _llama = Llama.from_pretrained(
        repo_id=LLAMACPP_REPO,
        filename=LLAMACPP_FILE,
        n_ctx=LLAMACPP_CTX,
        n_gpu_layers=LLAMACPP_GPU_LAYERS,
        verbose=False,
    )
    log.info("llama.cpp model loaded.")
    return _llama

@spaces.GPU(duration=60)
def _llamacpp_chat(system: str, user: str, schema: dict, temperature: float = 0.6) -> dict:
    """Schema-constrained JSON via llama.cpp's GBNF (built from the JSON schema).
    The grammar guarantees valid JSON and also suppresses <think> blocks (the
    first token must open the object)."""
    llm = _get_llama()
    for attempt in range(3):
        resp = llm.create_chat_completion(
            messages=[
                {"role": "system", "content": system + "\n/no_think"},
                {"role": "user",   "content": user},
            ],
            response_format={"type": "json_object", "schema": schema},
            max_tokens=LLAMACPP_CTX // 2,
            temperature=temperature, top_p=0.95,
        )
        raw = (resp["choices"][0]["message"]["content"] or "").strip()
        log.info("llamacpp raw (attempt %d) <- %s", attempt + 1, raw[:400])
        try:
            parsed = _apply_defaults(_unwrap(_extract_json(raw), schema), schema)
            log.info("llamacpp json -> %s", parsed)
            return parsed
        except (ValueError, json.JSONDecodeError, KeyError) as e:
            log.warning("llamacpp parse attempt %d failed: %s", attempt + 1, e)
            if attempt == 2:
                raise RuntimeError(
                    f"llama.cpp failed to produce valid JSON after 3 attempts. "
                    f"Last output: {raw!r}") from e

def _chat(system, user, schema, temperature: float = 0.6):
    log.info("system prompt: %s", system)
    log.info("user prompt: %s", user)
    log.info("output schema: %s", schema)
    if BACKEND == "modal_llm":
        return _modal_llm_chat(system, user, schema, temperature)
    return _llamacpp_chat(system, user, schema, temperature)

# Co-writing system prompt, shared by the pass-the-pen continuation.
_CONTINUE_SYSTEM = (
    "You are co-writing a story one story beat at a time. "
    "Continue with exactly 1-3 sentences that follow naturally. "
    "If the story is empty, start it however you like but do not return an empty string."
    "Output only the sentence(s), no quotes, no preamble."
)

# ---------------------------------------------------------------------------
# Public interface
# ---------------------------------------------------------------------------
def judge_beat(text, story, mode="exploration", pacing=None):
    log.info("judge_beat [%s]  beat=%r  pacing=%s", mode, text, pacing)
    if mode == "puzzle":
        md = _load_prompt("prompt_puzzle.md")
        system = _system_block(md.split("A) TARGET CALL", 1)[1])
        user = f'Story so far:\n{story or "(none yet)"}\nNew beat:\n"{text}"\nReturn the JSON object.'
        return _chat(system, user, TARGET_SCHEMA)
    md = _load_prompt("prompt_core.md")
    system = _system_block(md)
    pace = ""
    if pacing:
        placed = pacing.get("shapes_placed", 0)
        first_shape_note = " — no shape yet, be generous with the first one" if placed == 0 else " ([has shape] beats are settled — do not flag them again)"
        pace = (f'\nPacing: {placed} shapes already placed{first_shape_note}, '
                f'{pacing.get("shapes_left","?")} shapes left in budget, '
                f'~{pacing.get("beats_left","?")} beats remain. '
                f'Only flag deserves_shape for a genuinely strong beat; '
                f'spend the remaining budget across the strongest beats left.')
    user = (f'Story so far (each beat tagged with its shape status):\n'
            f'{story or "(none yet)"}\n'
            f'New beat:\n"{text}"{pace}\nReturn the JSON object.')
    return _chat(system, user, CORE_SCHEMA)

def judge_beat_segmented(text, story, pacing=None):
    """
    Like judge_beat (exploration mode) but chunks the sentence into 2-4 phrases,
    each with its own affect. Used for per-segment essence (phase 2).
    """
    log.info("judge_beat_segmented  beat=%r  pacing=%s", text, pacing)
    md = _load_prompt("prompt_core.md")
    system = _system_block(md)
    seg_instruction = (
        "\n\nADDITIONAL TASK: Before returning the JSON, chunk the sentence into "
        "2-4 meaningful phrases (not single words). Judge the affect of EACH phrase "
        "independently. Return them in order as the `segments` array. Each segment "
        "needs `phrase` (the text), `valence`, `arousal`, `dominance` (each 0..1), "
        "and `words` (integer word count of that phrase). The top-level valence/"
        "arousal/dominance should reflect the whole sentence as usual."
    )
    pace = ""
    if pacing:
        pace = (f'\nPacing: {pacing.get("shapes_placed",0)} shapes placed '
                f'([has shape] beats are settled), '
                f'{pacing.get("shapes_left","?")} left, '
                f'~{pacing.get("beats_left","?")} beats remain.')
    user = (f'Story so far (each beat tagged with its shape status):\n'
            f'{story or "(none yet)"}\n'
            f'New beat:\n"{text}"{pace}\nReturn the JSON object.')
    return _chat(system + seg_instruction, user, SEGMENT_SCHEMA)

def judge_attempt(target_affect, shape_sentence):
    log.info("judge_attempt  target=%s  attempt=%r", target_affect, shape_sentence)
    md = _load_prompt("prompt_puzzle.md")
    system = _system_block(md.split("B) ATTEMPT CALL", 1)[1])
    user = (f"Target feeling (for your judgment only, do not reveal): "
            f"V={target_affect['valence']} A={target_affect['arousal']} D={target_affect['dominance']}\n"
            f'Player\'s shape sentence:\n"{shape_sentence}"\nReturn the JSON object.')
    return _chat(system, user, ATTEMPT_SCHEMA)

_START_FALLBACK = "The morning the letter arrived, nothing in the house looked the same."

# Rotating "opening lenses" — without one, a low-temperature cold start collapses
# to the model's single favourite opening ("The old clock…"). Each nudges a
# different angle without dictating the actual content.
_OPENING_LENSES = [
    "Open in the middle of an action, already in motion.",
    "Open on a single vivid sensory detail (a sound, a smell, a texture).",
    "Open with a line of dialogue.",
    "Open on an unsettling or out-of-place image.",
    "Open with a character doing something ordinary, moments before it changes.",
    "Open on a specific place at a specific time of day.",
    "Open with something small that's gone wrong.",
]

def continue_story(story):
    log.info("continue_story (story %d chars)", len(story))
    starting = not story.strip()
    if starting:
        # Cold opens (empty story) are the hard case: with no context the model
        # both returns empty and, at low temperature, defaults to one stock
        # opening. A random lens + higher temperature break that mode.
        import random
        lens = random.choice(_OPENING_LENSES)
        user = ("Begin a brand-new short story. Write 1-3 vivid opening sentences "
                "that drop the reader into a character, place, or moment where "
                f"something is about to happen. {lens} "
                "Invent freely; never return empty.")
    else:
        user = f"Story so far:\n{story}\nContinue it with 1-3 sentences that follow naturally."
    # Structured output ({"sentence": ...}) so model reasoning can never leak:
    # llama.cpp's grammar forbids a <think> block; for Modal, _extract_json pulls
    # the object out even if the model thinks first. Higher temperature for this
    # creative call (judge calls stay at 0.6); a rare empty result is retried.
    out = ""
    for attempt in range(3):
        result = _chat(_CONTINUE_SYSTEM, user, SENTENCE_SCHEMA, temperature=0.95)
        out = (result.get("sentence") or "").strip().strip('"')
        if out:
            break
        log.warning("continue_story empty (attempt %d) — retrying", attempt + 1)
    if not out and starting:
        out = _START_FALLBACK
    log.info("continue_story -> %r", out)
    return out

def title_story(story: str) -> str:
    """Name the finished story: one short evocative title (for the share card)."""
    log.info("title_story (story %d chars)", len(story))
    system = (
        "You name a very short story. Read it and return ONE evocative title of "
        "2 to 5 words. Title Case. No quotation marks, no trailing punctuation, no "
        "explanation, no label like 'Title:'."
    )
    user = f"Story:\n{story or '(empty)'}"
    result = _chat(system, user, TITLE_SCHEMA)
    title = (result.get("title") or "").splitlines()[0].strip() if result.get("title") else ""
    title = title.strip('"').strip("'").strip("*").strip().rstrip(".").strip()
    # Strip "Title:" or "Title :" prefix that some models emit despite instructions
    import re as _re
    title = _re.sub(r"(?i)^title\s*:\s*", "", title).strip()
    log.info("title_story -> %r", title)
    return title or "A Story in Shapes"

def reveal(story, shapes_labels):
    log.info("reveal  labels=%s", shapes_labels)
    md = _load_prompt("prompt_reveal.md")
    system = _system_block(md)
    shapes_str = "\n".join(f"{i}: {lbl}" for i, lbl in enumerate(shapes_labels))
    user = f"Full story:\n{story}\nShapes already made:\n{shapes_str}\nReturn the JSON object."
    return _chat(system, user, REVEAL_SCHEMA)