#!/usr/bin/env python3
# fanfictioner: turn a story idea into a text-free illustrated book.
# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

import argparse
import hashlib
import http.client
import json
import os
import random
import re
import shlex
import shutil
import subprocess
import sys
import tempfile
import time
import urllib.request
from pathlib import Path

LLM_URL = "http://localhost:8181"
LLM_MODEL = "Gemma4-12B-qat-mtp"
SD_DIR = Path("/data/LLM-models/SD")
CONF = Path.home() / ".config/book-reader.conf"

# Settings verified by hand on 2026-10-04 (Krea2 re-checked against upstream on 2026-10-05).
# Model paths are relative to SD_DIR. All profiles run at cfg 1.0, so no negative prompts.
PROFILES = {
    "zimage": ["--diffusion-model", "z_image_turbo-Q8_0.gguf",
               "--vae", "vae/flux1-ae.safetensors",
               "--llm", "Qwen3-4B-Instruct-2507-Q8_0.gguf",
               "--cfg-scale", "1.0", "--steps", "8", "--diffusion-fa", "--offload-to-cpu"],
    "krea2": ["--diffusion-model", "Krea-2-Turbo-Q6_K.gguf",
              "--vae", "vae/wan_2.1_vae.safetensors",
              "--llm", "Qwen3VL-4B-Instruct-Q8_0.gguf",
              "--cfg-scale", "1.0", "--steps", "8", "--diffusion-fa", "--offload-to-cpu"],
    # te=cpu alone puts diffusion on the iGPU and crashes: always name all three backends.
    # --offload-to-cpu here OOM-killed the desktop: never add it.
    "qwen_edit": ["--diffusion-model", "Qwen-Image-2.1-viggle-turbo-v0.3-6step-Q8_0.gguf",
                  "--vae", "vae/qwen_image_2.1_vae_bf16.safetensors",
                  "--llm", "Qwen3VL-8B-Instruct-Q8_0.gguf",
                  "--llm_vision", "mmproj-Qwen3VL-8B-Instruct-F16.gguf",
                  "--cfg-scale", "1.0", "--steps", "6",
                  "--sigmas", "1.0,0.9375,0.875,0.75,0.5,0.25,0.0", "--sampling-method", "euler",
                  "--fa", "--backend", "te=cpu,diffusion=vulkan0,vae=vulkan0", "--mmap", "--vae-tiling"],
}
MODEL_FLAGS = {"--diffusion-model", "--vae", "--llm", "--llm_vision"}
SIZES = {"portrait": (832, 1216), "landscape": (1216, 832), "turnaround": (1216, 832)}
# Edit inputs are shrunk first: full-size refs blow past 12 GB VRAM.
SHRINK = {"portrait": "576x832", "landscape": "832x576", "ref": "560x384"}
MIN_RAM_KB = 2 * 1024 * 1024  # kill sd-cli below 2 GB MemAvailable, before the OOM killer hits the desktop

DRAFT_SYS = """You turn a story idea into a page plan for a text-free illustrated book.
Each page is one full-page image. A scene may span several pages.
Images never contain text: no speech bubbles, captions, signs or labels. Tell the story
through action, expression and setting. The only exception is the cover.
Page 1 is the book cover: the main characters in a striking composition, portrait, with a
Title: line holding the book title. No other page has a Title: line.
Give every character one fixed visual description that starts with who they are
(e.g. "Young woman", "Old man", "Little girl"), then age, hair, face, eyes, build,
and never vary it between pages. Clothing never goes in that description: put the
character's usual clothes on their Outfit: line. When a page needs different clothes
(a formal event, pajamas, swimwear, a uniform), give that page an Outfit: line for the
characters who change; their clothes then stay changed only on that page.
Characters are people only. Animals and objects belong to the scene: describe them in the
page text, never under ## Characters.
If a page has a Model: line, keep it unchanged.
Write the plan in exactly this Markdown format, with nothing before or after it:

# <Title>
Series: <series folder name>
Book: <book folder name>
Style: <one line describing the art style of every page>

## Characters

### <Name>
Pronoun: <she, he or they>
Outfit: <usual clothes>
<fixed visual description, no clothing>

## Pages

### 1. <short title>
<what happens, setting, mood>
Characters: <Name>, <Name>   (or: none)
Pose: <Name>: <position in the frame, body orientation and pose>; <Name>: <...>
Outfit: <Name>: <clothes on this page>; <Name>: <...>   (only for characters whose clothes differ)
Framing: <camera angle, shot size>, <portrait or landscape>
Title: <book title, cover page only>
"""

OBJ = {"she": "her", "he": "him", "they": "them"}  # pronoun -> object form for edit prompts


def load_conf(path=CONF):
    conf = {}
    for line in path.read_text().splitlines():
        if "=" in line and not line.lstrip().startswith("#"):
            k, v = line.split("=", 1)
            conf[k.strip()] = shlex.split(v)[0] if v.strip() else ""
    return conf


def page_path(lib, pj, n):
    """Chapters share one book folder: 01-001.png, 01-002.png, 02-001.png sort in reading order."""
    name = f"{pj['chapter']:02d}-{n:03d}.png" if pj.get("chapter") else f"{n:03d}.png"
    return Path(lib) / pj["series"] / pj["book"] / name


def parse_front(text):
    """story.md frontmatter between --- lines -> ({key: value}, story body)."""
    m = re.match(r"---\n(.*?)\n---\n?(.*)", text, re.S)
    if not m:
        return {}, text
    meta = {}
    for line in m[1].splitlines():
        k, sep, v = line.partition(":")
        if sep and v.strip():
            meta[k.strip().lower()] = v.strip().strip("\"'")
    return meta, m[2]


def apply_front(md, meta):
    """Force the story's frontmatter into plan.md, whatever Gemma wrote there."""
    keys = [k for k in ("Series", "Book", "Chapter", "Style") if meta.get(k.lower())]
    drop = keys + (["Title"] if meta.get("title") else [])
    lines = [x for x in md.splitlines()
             if not (drop and re.match(r"[-*\s]*\**(%s)\**:" % "|".join(drop), x.strip()))]
    top = next((i for i, x in enumerate(lines) if x.startswith("# ")), None)
    if top is None:
        lines.insert(0, "# " + meta.get("title", "Untitled"))
        top = 0
    lines[top + 1:top + 1] = [f"{k}: {meta[k.lower()]}" for k in keys]
    if meta.get("title"):
        p1 = next((i for i, x in enumerate(lines) if re.match(r"###\s+1\.", x.strip())), None)
        if p1 is not None:
            end = next((i for i in range(p1 + 1, len(lines)) if lines[i].startswith("#")), len(lines))
            while end > p1 + 1 and not lines[end - 1].strip():
                end -= 1
            lines.insert(end, f"Title: {meta['title']}")
    return "\n".join(lines) + "\n"


def todo(pj, wd, lib):
    """What is left to do, derived from disk: (characters without a ref, pages not in the library)."""
    refs = [c for c in pj["characters"] if not (wd / "refs" / f"{c['name']}.png").exists()]
    pages = [p for p in pj["pages"] if not page_path(lib, pj, p["n"]).exists()]
    return refs, pages


def keep(src, dst):
    """Copy src to dst atomically, so a half-written file never counts as done."""
    dst.parent.mkdir(parents=True, exist_ok=True)
    tmp = dst.with_name(dst.name + ".tmp")
    shutil.copyfile(src, tmp)
    os.replace(tmp, dst)
    return dst


def parse_plan(text):
    plan = {"title": "", "series": "", "book": "", "chapter": "", "style": "", "characters": {}, "pronouns": {},
            "outfits": {}, "pages": []}
    section = cur = cname = None
    for line in text.splitlines():
        s = line.strip()
        if s.startswith("```") or s == "---":
            continue
        if s.startswith("# ") and not plan["title"]:
            plan["title"] = s[2:].strip()
        elif s.startswith("## "):
            section, cur = s[3:].strip().lower(), None
        elif s.startswith("### "):
            head = s[4:].strip()
            if section == "characters":
                cname, cur = head, plan["characters"].setdefault(head, [])
            elif section == "pages":
                m = re.match(r"(\d+)\.\s*(.*)", head)
                cur = {"n": int(m[1]) if m else 0, "title": m[2] if m else head, "text": [],
                       "heading": "" if m else head, "characters": None, "pose": "", "outfit": "", "framing": "",
                       "model": "", "lettering": ""}
                plan["pages"].append(cur)
        elif section is None and (m := re.match(r"[-*\s]*\**(Series|Book|Chapter|Style)\**:\**\s*(.*)", s)):
            plan[m[1].lower()] = m[2].strip().strip("*").strip()
        elif section == "characters" and cur is not None and \
                (m := re.match(r"[-*\s]*\**(Pronoun|Outfit)\**:\**\s*(.*)", s)):
            val = m[2].strip().strip("*").strip()
            if m[1] == "Pronoun":
                plan["pronouns"][cname] = val.lower()
            else:
                plan["outfits"][cname] = val
        elif section == "pages" and cur is not None and \
                (m := re.match(r"[-*\s]*\**(Characters|Pose|Outfit|Framing|Model|Title)\**:\**\s*(.*)", s)):
            key, val = m[1].lower(), m[2].strip().strip("*").strip()
            if key == "title":  # the page heading already uses "title"
                key, val = "lettering", val.strip('"\u201c\u201d').strip()
            if key == "characters":
                val = [] if val.lower() == "none" else [c.strip() for c in val.split(",") if c.strip()]
            elif key == "model":
                val = val.split()[0].lower() if val else ""
            cur[key] = val
        elif cur is not None and s:
            (cur if isinstance(cur, list) else cur["text"]).append(s)
    plan["characters"] = {k: " ".join(v) for k, v in plan["characters"].items()}
    for p in plan["pages"]:
        p["text"] = " ".join(p["text"])
    return plan


def parse_poses(pose):
    """'Mara: left, facing right; Jon: center' -> {"Mara": "left, facing right", "Jon": "center"}"""
    return {k.strip(): v.strip().rstrip(".") for k, _, v in (x.partition(":") for x in pose.split(";"))
            if v.strip()}


def validate_plan(plan):
    errs = [f"missing {k}" for k in ("title", "series", "book") if not plan[k]]
    if not plan["characters"]:
        errs.append("no characters")
    if not plan["pages"]:
        errs.append("no pages")
    errs += [f'bad {k} name "{plan[k]}"' for k in ("series", "book")
             if plan[k] and (("/" in plan[k]) or plan[k] in (".", ".."))]
    errs += [f'bad character name "{c}"' for c in plan["characters"] if "/" in c or c in (".", "..")]
    if plan["chapter"] and not plan["chapter"].isdigit():
        errs.append(f'Chapter must be a number, not "{plan["chapter"]}"')
    errs += [f"character {c}: Pronoun must be she, he or they" for c in plan["characters"]
             if plan["pronouns"].get(c) not in OBJ]
    ns = [p["n"] for p in plan["pages"]]
    if sorted(ns) != list(range(1, len(ns) + 1)):
        errs.append("page numbers must be unique and start at 1")
    for p in plan["pages"]:
        if p["heading"]:
            errs.append(f'page heading "{p["heading"]}" must look like "### N. title"')
        if p["characters"] is None:
            errs.append(f"page {p['n']}: Characters missing")
        chars = p["characters"] or []
        errs += [f"page {p['n']}: unknown character {c}" for c in chars
                 if c not in plan["characters"]]
        if len(set(chars)) != len(chars):
            errs.append(f"page {p['n']}: duplicate character")
        poses = parse_poses(p["pose"])
        errs += [f"page {p['n']}: Pose missing for {c}" for c in chars if c not in poses]
        errs += [f"page {p['n']}: Outfit for {c}, who is not on the page" for c in parse_poses(p["outfit"])
                 if c not in chars]
        if not re.search(r"portrait|landscape", p["framing"], re.I):
            errs.append(f"page {p['n']}: Framing must say portrait or landscape")
        if p["model"] and p["model"] not in ("zimage", "krea2"):
            errs.append(f"page {p['n']}: unknown model {p['model']}")
    return errs


SAMPLE_PLAN = """# Rooftop
Series: Test Series
Book: Book One
Style: soft watercolor, muted colors

## Characters

### Mara
Pronoun: she
Young woman, long red hair,
green eyes, grey coat.

### Jon
Pronoun: he
Outfit: blue jacket
Tall man, short beard.

### Kit
Pronoun: she
Small girl, black bob, yellow raincoat.

## Pages

### 1. Arrival
They reach the roof at dusk.
Characters: Mara, Jon, Kit
Pose: Mara: left, facing right; Jon: center, facing camera; Kit: right, sitting.
Outfit: Jon: wearing a black tuxedo
Framing: wide shot, eye level, landscape

### 2. Empty roof
Nobody is there any more.
Characters: none
Framing: wide shot, portrait
Model: krea2
"""


def edit_groups(chars):
    """Qwen edit takes the scene plus at most 2 reference sheets, so chain groups of 2."""
    return [chars[i:i + 2] for i in range(0, len(chars), 2)]


def sentence(s):
    s = s.strip()
    return s[:1].upper() + s[1:] + ("" if not s or s[-1] in ".!?" else ".")


def build_json(plan):
    """plan.md -> plan.json. Every image prompt comes from a fixed template, no LLM involved."""
    style = sentence(plan["style"])

    def look(n, outfit=""):
        """Fixed description plus one outfit: the page's when it has one, else the usual one."""
        o = re.sub(r"(?i)^(wearing|wears)\s+", "", outfit or plan["outfits"].get(n, ""))
        return sentence(plan["characters"][n]) + (f" Wearing {o.rstrip('.')}." if o else "")

    chars = [{"name": n, "turnaround_prompt":
              f"Character turnaround reference sheet of one person: {look(n)} The same person shown "
              "three times side by side, full body: front view, side view, back view. Neutral standing pose, "
              f"arms relaxed. Plain light-grey background, even studio lighting. {style} No text, no labels."}
             for n in plan["characters"]]
    pages = []
    for p in plan["pages"]:
        poses, outfits = parse_poses(p["pose"]), parse_poses(p["outfit"])
        parts = [style] + [f"{look(n, outfits.get(n))} {sentence(poses[n])}" for n in p["characters"]]
        parts.append(sentence(p["text"]))
        parts.append(sentence(p["framing"]))
        if p["lettering"]:
            parts.append(f'Title lettering at the top of the image reading exactly "{p["lettering"]}". '
                         "No other text, no speech bubbles.")
        else:
            parts.append("No text, no speech bubbles.")
        keep = "background, lighting" + (", title lettering" if p["lettering"] else "")
        edits = [{"characters": g, "prompt": " ".join(
                     f"In image 1, change only {n}'s head: give {OBJ[plan['pronouns'][n]]} the face, eyes and hair "
                     f"of the person in image {i + 2}. Keep {n}'s exact pose from image 1: {poses[n]}."
                     for i, n in enumerate(g))
                  + f" Keep bodies, clothing, other people, {keep} and art style of image 1 unchanged."}
                 for g in edit_groups(p["characters"])]
        pages.append({"n": p["n"], "orientation": "landscape" if "landscape" in p["framing"].lower() else "portrait",
                      "model": p["model"], "base_prompt": " ".join(x for x in parts if x), "edits": edits})
    return {"title": plan["title"], "series": plan["series"], "book": plan["book"],
            "chapter": int(plan["chapter"]) if plan["chapter"] else None, "characters": chars, "pages": pages}


def sd_args(profile, prompt, out, size, seed, refs=()):
    args = ["sd-cli"]
    it = iter(PROFILES[profile])
    for x in it:
        args.append(x)
        if x in MODEL_FLAGS:
            args.append(str(SD_DIR / next(it)))
    args += ["-p", prompt, "-W", str(size[0]), "-H", str(size[1]), "-s", str(seed),
             "-b", "3", "-o", str(out), "--output-begin-idx", "1"]
    for r in refs:
        args += ["-r", str(r)]
    return args


def mem_available_kb(meminfo):
    return int(re.search(r"MemAvailable:\s+(\d+)", meminfo)[1])


def progress(start, log):
    """One self-overwriting status line: elapsed time and sd-cli's latest progress output."""
    with open(log, "rb") as f:
        f.seek(max(0, os.path.getsize(log) - 4096))
        last = next((x for x in reversed(re.split(r"[\r\n]", f.read().decode(errors="replace"))) if x.strip()), "")
    width = shutil.get_terminal_size().columns - 1
    print(f"\r{int(time.time() - start):4d}s {last.strip()}"[:width].ljust(width), end="", flush=True)


def run_sd(args, log):
    """Run sd-cli with output in log; return an error string or None."""
    with open(log, "w") as f:
        f.write(shlex.join(args) + "\n")
        f.flush()
        try:
            proc = subprocess.Popen(args, stdout=f, stderr=subprocess.STDOUT)
        except OSError as e:
            return f"cannot run {args[0]}: {e}"
        start = time.time()
        try:
            while proc.poll() is None:
                time.sleep(2)
                progress(start, log)
                if proc.poll() is None and \
                        mem_available_kb(Path("/proc/meminfo").read_text()) < MIN_RAM_KB:
                    proc.kill()
                    proc.wait()
                    return "sd-cli killed by RAM watchdog (MemAvailable below 2 GB)"
        except KeyboardInterrupt:
            proc.kill()
            proc.wait()
            raise
        finally:
            print()
    if proc.returncode or "[ERROR" in Path(log).read_text(errors="replace"):
        return f"sd-cli failed (exit {proc.returncode})"
    return None


LOG = None  # wd / "run.log", set by main


def log(msg):
    """Append one timestamped event to run.log, once the work dir exists."""
    if LOG and LOG.parent.is_dir():
        with open(LOG, "a") as f:
            f.write(f"{time.strftime('%Y-%m-%d %H:%M:%S')} {msg}\n")


NOTIFY_ID = None  # notify-send id of the last notification, replaced by the next one


def notify(body):
    """Desktop notification that replaces the previous one instead of stacking."""
    global NOTIFY_ID
    try:
        out = subprocess.run(["notify-send", "-a", "fanfictioner", "-p", *(["-r", NOTIFY_ID] if NOTIFY_ID else []),
                              "fanfictioner", body], capture_output=True, text=True, timeout=5).stdout.strip()
    except (OSError, subprocess.TimeoutExpired):
        return
    NOTIFY_ID = out if out.isdigit() else None


def editor(path):
    subprocess.run([*shlex.split(os.environ.get("EDITOR") or "vi"), str(path)])


def generate(d, profile, prompt, size, refs=()):
    """One sd-cli call producing 3 candidates in d. Returns (files, seed, error)."""
    unload_llm()
    d.mkdir(parents=True, exist_ok=True)
    seed = random.randint(0, 2**31 - 4)
    log = d / f"{seed}.log"
    print(f"sd-cli {profile}, seed {seed} ...", flush=True)
    err = run_sd(sd_args(profile, prompt, d / f"{seed}_%d.png", size, seed, refs), log)
    files = [d / f"{seed}_{i}.png" for i in (1, 2, 3)]
    if not err and not all(f.exists() for f in files):
        err = "sd-cli wrote no output"
    if err:
        print(err, *log.read_text(errors="replace").splitlines()[-15:], sep="\n")
    return files, seed, err


def fit(m):
    """A copy of the montage short enough to leave the key prompt visible below it."""
    try:
        w, h = map(int, subprocess.run(["kitty", "+kitten", "icat", "--print-window-size"],
                                       capture_output=True, text=True).stdout.strip().split("x"))
        rows = shutil.get_terminal_size().lines
        v = m.with_name(f"{m.stem}_view.png")
        subprocess.run(["magick", str(m), "-resize", f"{w}x{h * (rows - 4) // rows}>", str(v)], check=True)
        return v
    except (OSError, ValueError, subprocess.CalledProcessError):
        return m


def show(files, seed, d):
    m = d / f"{seed}_montage.png"
    labelled = [x for i, f in enumerate(files) for x in ("-label", f"{i + 1}   seed {seed + i}", str(f))]
    try:
        failed = subprocess.run(["magick", "montage", "-pointsize", "32", *labelled,
                                 "-geometry", "+8+8", "-tile", "3x1", str(m)]).returncode
    except OSError:
        failed = True
    if failed:
        print("montage failed, candidates:", *files, sep="\n")
        return
    print(m)
    try:  # inline in kitty, else in the desktop's image viewer
        if os.environ.get("KITTY_WINDOW_ID"):
            subprocess.run(["kitty", "+kitten", "icat", str(fit(m))])
        else:
            subprocess.Popen(["xdg-open", str(m)], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    except OSError as e:
        print(f"cannot show image: {e}")


def review(d, profile, prompt, size, refs=(), skip=False):
    """Generate, show, ask. Returns the picked file, or None when the user skips."""
    keys = "1-3 pick, r regen, p edit prompt" + (", s skip edit" if skip else "") + ", q quit"
    stage = f"{d.parent.name}/{d.name}"
    while True:
        files, seed, err = generate(d, profile, prompt, size, refs)
        log(f"{stage} {profile} seed {seed}" + (f" FAILED: {err}" if err else ""))
        if not err:
            show(files, seed, d)
        notify(f"{stage}: " + ("sd-cli failed" if err else "candidates ready, pick one"))
        while True:
            try:
                k = input(f"[{keys}] > ").strip().lower()
            except EOFError:
                k = "q"
            if k in ("1", "2", "3") and not err:
                log(f"{stage} pick {k}, seed {seed + int(k) - 1}, {files[int(k) - 1].name}")
                print(f"picked {k} (seed {seed + int(k) - 1})")
                return files[int(k) - 1]
            if k == "r":
                log(f"{stage} regenerate")
                print("regenerating with a new seed")
                break
            if k == "p":
                (d / "prompt.txt").write_text(prompt)
                editor(d / "prompt.txt")
                prompt = (d / "prompt.txt").read_text().strip() or prompt
                log(f"{stage} prompt edited: {prompt}")
                print(f"regenerating with prompt:\n{prompt}")
                break
            if k == "s" and skip:
                log(f"{stage} skip edit")
                print("edit skipped, scene kept as is")
                return None
            if k == "q":
                log(f"{stage} quit")
                sys.exit("quit, run again to resume")


def llm(path, body=None, timeout=900):
    data = json.dumps(body).encode() if body is not None else None
    req = urllib.request.Request(LLM_URL + path, data=data, headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.load(r)


def check_llm():
    try:
        ids = [m["id"] for m in llm("/v1/models", timeout=10)["data"]]
    except OSError as e:
        sys.exit(f"llama-server unreachable at {LLM_URL}: {e}")
    if LLM_MODEL not in ids:
        sys.exit(f"llama-server does not offer {LLM_MODEL} (has: {', '.join(ids)})")


def unload_llm():
    """Free VRAM for sd-cli. Idempotent: an unloaded or unreachable server is fine."""
    try:
        llm("/models/unload", {"model": LLM_MODEL}, timeout=10)
    except (OSError, ValueError):
        pass


def chat(system, user):
    body = {"model": LLM_MODEL,
            "messages": [{"role": "system", "content": system}, {"role": "user", "content": user}]}
    print("asking Gemma ...", flush=True)
    return llm("/v1/chat/completions", body)["choices"][0]["message"]["content"]


def strip_fences(s):
    return re.sub(r"^```\w*\n|\n```$", "", s.strip())


GEMMA_ERRS = (OSError, ValueError, KeyError, http.client.HTTPException)


def revise(md, meta, wd, ins):
    """Gemma rewrites the plan; returns the new plan.md text, or None when it fails or is invalid."""
    try:
        new = apply_front(strip_fences(chat(
            DRAFT_SYS, f"Current plan:\n\n{md}\n\nRevise it: {ins}\nReturn the complete new plan.")), meta)
    except GEMMA_ERRS as e:
        print(f"Gemma failed: {e}")
        return None
    nerrs = validate_plan(parse_plan(new))
    if nerrs:
        (wd / "plan.rejected.md").write_text(new)
        log(f"Gemma revision rejected ({'; '.join(nerrs)}): {ins}")
        print("revision invalid, plan.md kept, see plan.rejected.md", *nerrs, sep="\n  ")
        return None
    (wd / "plan.md.bak").write_text(md)
    log(f"Gemma revision accepted: {ins}")
    return new


def plan_stage(story, wd):
    md_path, js_path = wd / "plan.md", wd / "plan.json"
    seen = wd / "story.seen"  # the story body plan.md was last drafted or revised from
    if js_path.exists() and md_path.exists() and \
            js_path.stat().st_mtime >= max(md_path.stat().st_mtime, story.stat().st_mtime):
        return json.loads(js_path.read_text())
    meta, body = parse_front(story.read_text())
    if md_path.exists() and (seen.read_text() != body if seen.exists()
                             else story.stat().st_mtime > md_path.stat().st_mtime):
        notify("story changed, plan needs a decision")
        try:
            k = input("\nstory changed since plan.md was drafted: "
                      "[r]evise plan with Gemma [n]ew draft [k]eep plan [q]uit > ").strip().lower()
        except EOFError:
            k = "q"
        if k == "n":
            os.replace(md_path, wd / "plan.md.bak")
            log("story changed, plan.md moved to plan.md.bak for a new draft")
        elif k == "r":
            check_llm()
            new = revise(md_path.read_text(), meta, wd, "the story changed, update the plan to match "
                         f"this new story and keep whatever still fits:\n\n{body}")
            if new:
                md_path.write_text(new)
                seen.write_text(body)
        elif k == "k":
            seen.write_text(body)
            log("story changed, plan.md kept")
        else:
            sys.exit("quit, run again to resume")
    if not md_path.exists():
        check_llm()
        wd.mkdir(exist_ok=True)
        try:
            draft = strip_fences(chat(DRAFT_SYS, body))
        except GEMMA_ERRS as e:
            sys.exit(f"Gemma failed: {e}")
        md_path.write_text(apply_front(draft, meta))
        seen.write_text(body)
        log("plan.md drafted by Gemma")
    while True:
        md = md_path.read_text()
        if apply_front(md, meta) != md:  # story.md frontmatter wins over plan.md
            md = apply_front(md, meta)
            md_path.write_text(md)
        plan = parse_plan(md)
        errs = validate_plan(plan)
        print(f"\n{md}\n--- {md_path}: {len(plan['characters'])} characters, {len(plan['pages'])} pages")
        for e in errs:
            print("  !", e)
        notify("plan ready for review")
        try:
            k = input("[c]onfirm [e]dit [g]emma revise [q]uit > ").strip().lower()
        except EOFError:
            k = "q"
        if k == "c" and errs:
            print("fix the ! lines first")
        elif k == "c":
            pj = build_json(plan)
            tmp = js_path.with_name("plan.json.tmp")
            tmp.write_text(json.dumps(pj, indent=2) + "\n")
            os.replace(tmp, js_path)
            log(f"plan confirmed: {len(pj['characters'])} characters, {len(pj['pages'])} pages")
            return pj
        elif k == "e":
            editor(md_path)
        elif k == "g":
            try:
                ins = input("instruction for Gemma: ")
            except EOFError:
                continue
            if new := revise(md, meta, wd, ins):
                md_path.write_text(new)
        elif k == "q":
            sys.exit(0)


def shrink(src, dst, geom):
    dst.parent.mkdir(parents=True, exist_ok=True)
    subprocess.run(["magick", str(src), "-resize", geom, str(dst)], check=True)
    return dst


def refs_stage(chars, wd):
    for c in chars:
        print(f"\n== reference sheet: {c['name']}")
        pick = review(wd / "cand" / f"ref_{c['name']}", "zimage", c["turnaround_prompt"], SIZES["turnaround"])
        keep(pick, wd / "refs" / f"{c['name']}.png")
        log(f"reference sheet {c['name']} kept")


def page_stage(p, wd, out, base):
    """Base scene, then chained head edits; each pick kept in cand/pNNN so a quit resumes mid-page."""
    key = json.dumps([p["base_prompt"], p["orientation"], p["model"], p["edits"]])
    d = wd / "cand" / f"p{p['n']:03d}_{hashlib.sha1(key.encode()).hexdigest()[:8]}"
    size = SIZES[p["orientation"]]
    scene = d / "base.png"
    if not scene.exists():
        print(f"\n== page {p['n']}: base scene")
        keep(review(d / "base", p["model"] or base, p["base_prompt"], size), scene)
    else:
        print(f"page {p['n']}: base scene reused from earlier pick")
    for i, e in enumerate(p["edits"], 1):
        nxt = d / f"edit{i}.png"
        if not nxt.exists():
            print(f"\n== page {p['n']}: edit {i} ({', '.join(e['characters'])})")
            refs = [shrink(scene, d / f"edit{i}_in.png", SHRINK[p["orientation"]])]
            refs += [shrink(wd / "refs" / f"{n}.png", wd / "cand" / "refs_small" / f"{n}.png", SHRINK["ref"])
                     for n in e["characters"]]
            pick = review(d / f"edit{i}", "qwen_edit", e["prompt"], size, refs, skip=True)
            keep(pick or scene, nxt)
        else:
            print(f"page {p['n']}: edit {i} ({', '.join(e['characters'])}) reused from earlier pick")
        scene = nxt
    keep(scene, out)
    log(f"page {p['n']} -> {out}")
    print(f"page {p['n']} -> {out}")


def selftest():
    plan = parse_plan(SAMPLE_PLAN)
    assert (plan["title"], plan["series"], plan["book"]) == ("Rooftop", "Test Series", "Book One"), plan
    assert plan["style"] == "soft watercolor, muted colors"
    assert list(plan["characters"]) == ["Mara", "Jon", "Kit"]
    assert plan["characters"]["Mara"] == "Young woman, long red hair, green eyes, grey coat."
    p1, p2 = plan["pages"]
    assert (p1["n"], p1["title"], p1["characters"]) == (1, "Arrival", ["Mara", "Jon", "Kit"])
    assert p1["text"] == "They reach the roof at dusk." and p1["pose"].startswith("Mara: left")
    assert p1["model"] == "" and p2["model"] == "krea2" and p2["characters"] == []
    assert validate_plan(plan) == []

    bad = parse_plan(SAMPLE_PLAN.replace("Book: Book One\n", "")
                     .replace("Characters: none", "Characters: Zed")
                     .replace("Model: krea2", "Model: sdxl")
                     .replace("### 2.", "### 1."))
    errs = validate_plan(bad)
    assert "missing book" in errs, errs
    assert "page 1: unknown character Zed" in errs, errs
    assert "page 1: Pose missing for Zed" in errs, errs
    assert "page 1: unknown model sdxl" in errs, errs
    assert "page numbers must be unique and start at 1" in errs, errs
    assert validate_plan(parse_plan("")) == ["missing title", "missing series", "missing book",
                                            "no characters", "no pages"]

    md = parse_plan(SAMPLE_PLAN.replace("Characters: Mara, Jon, Kit", "**Characters:** Mara, Jon, Kit")
                    .replace("Pose:", "- **Pose**:").replace("Framing:", "* Framing:")
                    .replace("Series:", "**Series:**").replace("Style:", "- Style:"))
    assert md["series"] == "Test Series" and md["style"] == "soft watercolor, muted colors", md
    assert md["pages"][0]["characters"] == ["Mara", "Jon", "Kit"], md
    assert md["pages"][0]["pose"].startswith("Mara: left") and md["pages"][0]["framing"].startswith("wide"), md
    assert validate_plan(md) == []
    assert "page 2: Characters missing" in validate_plan(parse_plan(SAMPLE_PLAN.replace("Characters: none\n", "")))
    assert 'bad series name "../x"' in validate_plan(parse_plan(SAMPLE_PLAN.replace("Test Series", "../x")))
    assert 'bad book name ".."' in validate_plan(parse_plan(SAMPLE_PLAN.replace("Book One", "..")))
    assert 'bad character name "Jon/Kit"' in validate_plan(parse_plan(SAMPLE_PLAN.replace("### Jon", "### Jon/Kit")))
    assert "page numbers must be unique and start at 1" in validate_plan(
        parse_plan(SAMPLE_PLAN.replace("### 1.", "### 2.").replace("### 2. Empty", "### 3. Empty")))
    assert 'page heading "Page 2: Empty roof" must look like "### N. title"' in validate_plan(
        parse_plan(SAMPLE_PLAN.replace("### 2. Empty", "### Page 2: Empty")))
    assert "page 1: duplicate character" in validate_plan(
        parse_plan(SAMPLE_PLAN.replace("Characters: Mara, Jon, Kit", "Characters: Mara, Mara")))
    fenced = parse_plan(SAMPLE_PLAN.replace("They reach the roof at dusk.", "```\nThey reach the roof at dusk.\n---\n```"))
    assert fenced["pages"][0]["text"] == "They reach the roof at dusk.", fenced
    assert parse_plan(SAMPLE_PLAN.replace("Series: Test Series", "**Series: Test Series**"))["series"] == "Test Series"
    assert parse_plan(SAMPLE_PLAN.replace("Characters: Mara, Jon, Kit", "**Characters: Mara, Jon**"))["pages"][0]["characters"] == ["Mara", "Jon"]
    assert edit_groups([]) == []
    assert edit_groups(["Mara", "Jon", "Kit"]) == [["Mara", "Jon"], ["Kit"]]
    assert edit_groups(list("abcd")) == [["a", "b"], ["c", "d"]]
    assert parse_poses("Mara: left, facing right; Jon: center.") == {"Mara": "left, facing right", "Jon": "center"}
    assert "character Jon: Pronoun must be she, he or they" in validate_plan(
        parse_plan(SAMPLE_PLAN.replace("Pronoun: he\n", "")))
    assert "page 1: Pose missing for Kit" in validate_plan(
        parse_plan(SAMPLE_PLAN.replace("; Kit: right, sitting.", "")))
    assert "page 2: Framing must say portrait or landscape" in validate_plan(
        parse_plan(SAMPLE_PLAN.replace("wide shot, portrait", "wide shot")))

    meta, body = parse_front('---\nseries: A new life\nbook: "Life as Daniel"\nchapter: 1\n'
                             'title: Meet Daniel\n---\nThis is the story.\n')
    assert meta == {"series": "A new life", "book": "Life as Daniel", "chapter": "1",
                    "title": "Meet Daniel"} and body == "This is the story.\n", meta
    assert parse_front("no frontmatter") == ({}, "no frontmatter")
    forced = apply_front(SAMPLE_PLAN.replace("Framing: wide shot, eye level, landscape",
                                             "Framing: wide shot, eye level, landscape\nTitle: Gemma Title"), meta)
    assert apply_front(forced, meta) == forced
    assert apply_front(SAMPLE_PLAN, {}) == SAMPLE_PLAN
    styled = parse_plan(apply_front(SAMPLE_PLAN, {"style": "ink manga"}))
    assert styled["style"] == "ink manga" and validate_plan(styled) == [], styled
    fp = parse_plan(forced)
    assert (fp["title"], fp["series"], fp["book"], fp["chapter"]) == ("Rooftop", "A new life", "Life as Daniel", "1"), fp
    assert fp["pages"][0]["lettering"] == "Meet Daniel" and fp["pages"][1]["lettering"] == "", fp["pages"]
    assert fp["pages"][0]["framing"] == "wide shot, eye level, landscape" and validate_plan(fp) == []
    assert forced.count("Title:") == 1 and forced.startswith("# Rooftop\nSeries: A new life\nBook: Life as Daniel\n"
                                                             "Chapter: 1\nStyle:"), forced
    assert build_json(fp)["chapter"] == 1 and build_json(plan)["chapter"] is None
    assert 'Chapter must be a number, not "one"' in validate_plan(
        parse_plan(SAMPLE_PLAN.replace("Book: Book One", "Book: Book One\nChapter: one")))

    pj = build_json(plan)
    assert (pj["title"], pj["series"], pj["book"]) == ("Rooftop", "Test Series", "Book One")
    assert [c["name"] for c in pj["characters"]] == ["Mara", "Jon", "Kit"]
    assert "of one person: Young woman, long red hair, green eyes, grey coat. The same person" in \
        pj["characters"][0]["turnaround_prompt"]
    assert pj["characters"][0]["turnaround_prompt"].endswith("Soft watercolor, muted colors. No text, no labels.")
    assert "one person: Tall man, short beard. Wearing blue jacket. The same" in pj["characters"][1]["turnaround_prompt"]
    assert "page 1: Outfit for Zed, who is not on the page" in validate_plan(
        parse_plan(SAMPLE_PLAN.replace("Outfit: Jon:", "Outfit: Zed:")))
    j1, j2 = pj["pages"]
    assert (j1["n"], j1["orientation"], j1["model"]) == (1, "landscape", "")
    assert (j2["n"], j2["orientation"], j2["model"], j2["edits"]) == (2, "portrait", "krea2", [])
    assert j1["base_prompt"] == (
        "Soft watercolor, muted colors. "
        "Young woman, long red hair, green eyes, grey coat. Left, facing right. "
        "Tall man, short beard. Wearing a black tuxedo. Center, facing camera. "
        "Small girl, black bob, yellow raincoat. Right, sitting. They reach the roof at dusk. "
        "Wide shot, eye level, landscape. No text, no speech bubbles."), j1["base_prompt"]
    assert j2["base_prompt"] == ("Soft watercolor, muted colors. Nobody is there any more. "
                                 "Wide shot, portrait. No text, no speech bubbles."), j2["base_prompt"]
    assert [e["characters"] for e in j1["edits"]] == [["Mara", "Jon"], ["Kit"]]
    assert j1["edits"][0]["prompt"] == (
        "In image 1, change only Mara's head: give her the face, eyes and hair of the person in image 2. "
        "Keep Mara's exact pose from image 1: left, facing right. "
        "In image 1, change only Jon's head: give him the face, eyes and hair of the person in image 3. "
        "Keep Jon's exact pose from image 1: center, facing camera. "
        "Keep bodies, clothing, other people, background, lighting and art style of image 1 unchanged.")
    cover = build_json(parse_plan(SAMPLE_PLAN.replace("Framing: wide shot, eye level, landscape",
                                                      'Framing: wide shot, eye level, landscape\nTitle: "Rooftop"')))
    c1, c2 = cover["pages"]
    assert c1["base_prompt"].endswith('landscape. Title lettering at the top of the image reading exactly '
                                      '"Rooftop". No other text, no speech bubbles.'), c1["base_prompt"]
    assert c1["edits"][1]["prompt"].endswith("background, lighting, title lettering and art style of image 1 unchanged.")
    assert c2["base_prompt"].endswith("No text, no speech bubbles.") and c1["n"] == 1
    assert "Kit's head: give her the face, eyes and hair of the person in image 2. " \
        "Keep Kit's exact pose from image 1: right, sitting." in j1["edits"][1]["prompt"]
    q = sd_args("qwen_edit", "edit it", Path("c/7_%d.png"), (832, 1216), 7,
                [Path("in1.png"), Path("mara.png")])
    assert q[0] == "sd-cli"
    assert q[q.index("--diffusion-model") + 1] == str(SD_DIR / "Qwen-Image-2.1-viggle-turbo-v0.3-6step-Q8_0.gguf")
    assert q[q.index("--llm_vision") + 1] == str(SD_DIR / "mmproj-Qwen3VL-8B-Instruct-F16.gguf")
    assert q[q.index("--backend") + 1] == "te=cpu,diffusion=vulkan0,vae=vulkan0"
    assert "--offload-to-cpu" not in q
    assert q[q.index("-r") + 1] == "in1.png" and q.count("-r") == 2
    for flag, val in (("-p", "edit it"), ("-W", "832"), ("-H", "1216"), ("-s", "7"), ("-b", "3"),
                      ("-o", "c/7_%d.png"), ("--output-begin-idx", "1"), ("--steps", "6")):
        assert q[q.index(flag) + 1] == val, (flag, q)
    z = sd_args("zimage", "x", Path("o_%d.png"), SIZES["turnaround"], 1)
    assert "--offload-to-cpu" in z and "-r" not in z
    assert z[z.index("--vae") + 1] == str(SD_DIR / "vae/flux1-ae.safetensors")
    k = sd_args("krea2", "x", Path("o_%d.png"), SIZES["portrait"], 1)
    assert k[k.index("--llm") + 1] == str(SD_DIR / "Qwen3VL-4B-Instruct-Q8_0.gguf")
    assert k[k.index("--steps") + 1] == "8" and k[k.index("--cfg-scale") + 1] == "1.0"
    assert strip_fences("```markdown\n# T\nx\n```\n") == "# T\nx"
    assert strip_fences("# T\nx") == "# T\nx"
    with tempfile.TemporaryDirectory() as t:
        t = Path(t)
        (t / "c.conf").write_text('# comment\nLIBRARY="/data/my lib"\nPORT=8642\n')
        conf = load_conf(t / "c.conf")
        assert conf == {"LIBRARY": "/data/my lib", "PORT": "8642"}, conf

        pj = {"series": "S", "book": "B",
              "characters": [{"name": "Mara"}, {"name": "Jon"}],
              "pages": [{"n": 1}, {"n": 2}]}
        lib = t / "lib"
        assert page_path(lib, pj, 7) == lib / "S" / "B" / "007.png"
        assert page_path(lib, dict(pj, chapter=2), 7) == lib / "S" / "B" / "02-007.png"
        (t / "refs").mkdir()
        (t / "refs" / "Mara.png").touch()
        page_path(lib, pj, 1).parent.mkdir(parents=True)
        page_path(lib, pj, 1).touch()
        refs, pages = todo(pj, t, lib)
        assert [c["name"] for c in refs] == ["Jon"], refs
        assert [p["n"] for p in pages] == [2], pages

        keep(t / "refs" / "Mara.png", t / "deep" / "x.png")
        assert (t / "deep" / "x.png").read_bytes() == (t / "refs" / "Mara.png").read_bytes()
        assert not list((t / "deep").glob("*.tmp"))
        assert mem_available_kb("MemTotal: 30000000 kB\nMemAvailable:   1500000 kB\n") == 1500000
        assert run_sd(["true"], t / "ok.log") is None
        assert (t / "ok.log").read_text().startswith("true")
        assert run_sd(["sh", "-c", "echo '[ERROR] boom'"], t / "e.log") == "sd-cli failed (exit 0)"
        assert run_sd(["false"], t / "f.log") == "sd-cli failed (exit 1)"
        assert run_sd(["/nonexistent-sd-cli"], t / "n.log").startswith("cannot run /nonexistent-sd-cli")
        global LOG
        LOG = t / "nowd" / "run.log"
        log("before the work dir exists")
        LOG = t / "run.log"
        log("one")
        log("two")
        lines = LOG.read_text().splitlines()
        LOG = None
        assert [x[20:] for x in lines] == ["one", "two"] and not (t / "nowd").exists(), lines
    print("selftest ok")


def main():
    ap = argparse.ArgumentParser(description="Turn story.md into a text-free illustrated book.")
    ap.add_argument("story", nargs="?", type=Path, help="story idea in free prose")
    ap.add_argument("--base", choices=("zimage", "krea2"), default="zimage",
                    help="base model for pages without a Model: line")
    ap.add_argument("--selftest", action="store_true", help="run built-in checks and exit")
    a = ap.parse_args()
    if a.selftest:
        return selftest()
    if not a.story:
        ap.error("story file required")
    if not a.story.is_file():
        ap.error(f"{a.story}: not a file")
    if a.story.suffix == "":
        ap.error("story file needs an extension, e.g. story.md")
    try:
        try:
            lib = load_conf().get("LIBRARY")
        except FileNotFoundError:
            lib = None
        if not lib:
            sys.exit(f"set LIBRARY in {CONF}")
        global LOG
        wd = a.story.with_suffix("")
        LOG = wd / "run.log"
        log(f"start, base model {a.base}")
        pj = plan_stage(a.story, wd)
        chars, pages = todo(pj, wd, lib)
        refs_stage(chars, wd)
        for p in pages:
            page_stage(p, wd, page_path(lib, pj, p["n"]), a.base)
        log("done")
        print(f"\ndone: {page_path(lib, pj, 1).parent}")
        notify("book done")
    except (KeyboardInterrupt, EOFError):
        log("interrupted")
        sys.exit("\ninterrupted, run again to resume")


if __name__ == "__main__":
    main()
