#!/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 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"}


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):
    return Path(lib) / pj["series"] / pj["book"] / f"{n:03d}.png"


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": "", "style": "", "characters": {}, "pages": []}
    section = cur = 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":
                cur = 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": "", "framing": "", "model": ""}
                plan["pages"].append(cur)
        elif section is None and (m := re.match(r"[-*\s]*\**(Series|Book|Style)\**:\**\s*(.*)", s)):
            plan[m[1].lower()] = m[2].strip()
        elif section == "pages" and cur is not None and \
                (m := re.match(r"[-*\s]*\**(Characters|Pose|Framing|Model)\**:\**\s*(.*)", s)):
            key, val = m[1].lower(), m[2].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 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 (".", "..")]
    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")
        if chars and not p["pose"]:
            errs.append(f"page {p['n']}: Pose missing")
        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
Young woman, long red hair,
green eyes, grey coat.

### Jon
Tall man, short beard, blue jacket.

### Kit
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.
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 edit_hint(plan):
    return "\n".join(f"Page {p['n']}: " + (" then ".join("[" + ", ".join(g) + "]"
                                                        for g in edit_groups(p["characters"] or []))
                                            or "no edits")
                     for p in plan["pages"])


def check_json(pj, plan):
    errs = []
    if [c["name"] for c in pj["characters"]] != list(plan["characters"]):
        errs.append("characters differ from plan.md")
    want = {p["n"]: p["characters"] or [] for p in plan["pages"]}
    if sorted(p["n"] for p in pj["pages"]) != sorted(want):
        errs.append("pages differ from plan.md")
    for p in pj["pages"]:
        got, exp = [e["characters"] for e in p["edits"]], edit_groups(want.get(p["n"], []))
        if got != exp:
            errs.append(f"page {p['n']}: edit groups {got}, expected {exp}")
    return errs


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 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" 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 edit_hint(parse_plan(SAMPLE_PLAN.replace("Characters: none\n", ""))).endswith("Page 2: no edits")
    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 edit_groups([]) == []
    assert edit_groups(["Mara", "Jon", "Kit"]) == [["Mara", "Jon"], ["Kit"]]
    assert edit_groups(list("abcd")) == [["a", "b"], ["c", "d"]]
    assert edit_hint(plan) == "Page 1: [Mara, Jon] then [Kit]\nPage 2: no edits"

    good = {"characters": [{"name": n, "turnaround_prompt": "x"} for n in ("Mara", "Jon", "Kit")],
            "pages": [{"n": 1, "orientation": "landscape", "base_prompt": "x",
                       "edits": [{"characters": ["Mara", "Jon"], "prompt": "x"},
                                 {"characters": ["Kit"], "prompt": "x"}]},
                      {"n": 2, "orientation": "portrait", "base_prompt": "x", "edits": []}]}
    assert check_json(good, plan) == []
    regrouped = json.loads(json.dumps(good))
    regrouped["pages"][0]["edits"] = [{"characters": ["Mara"], "prompt": "x"},
                                      {"characters": ["Jon", "Kit"], "prompt": "x"}]
    assert check_json(regrouped, plan) == [
        "page 1: edit groups [['Mara'], ['Jon', 'Kit']], expected [['Mara', 'Jon'], ['Kit']]"]
    short = json.loads(json.dumps(good))
    del short["characters"][2], short["pages"][1]
    assert check_json(short, plan) == ["characters differ from plan.md", "pages differ from plan.md"]
    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"
    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"
        (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"))
    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 __name__ == "__main__":
    main()
