"""Browser-side generation driver for PXG-Tiny (Pyodide, no PIL).

Mirrors the shipped PXGPipeline escalation schedule exactly, but without the
render/PIL dependency:
  attempt 0     raw sample
  attempts 1-3  palette-logit self-guidance, strength ramps
  attempts 4-5  + spatial priors (face box / vial margins)
  attempts 6-7  + canonical anchor caption
Structural prefixes (deer antlers) engage from attempt 3, as shipped.
"""
import numpy as np

from pxg_tiny import quality as Q
from pxg_tiny.config import encode_caption
from pxg_tiny.runtime_pipeline import OfflinePipeline


class WebPipeline:
    def __init__(self, bundle_dir):
        self.pipe = OfflinePipeline(bundle_dir)

    def _anchor_ids(self, spec):
        cls = spec.get("cls")
        cap = Q.ANCHOR_CAPTIONS.get(cls)
        if cap is None:
            return None
        if (spec.get("material") and cls not in ("wizard", "archer", "zombie")
                and f" {spec['material']}" not in cap):
            art = "an" if spec["material"][0] in "aeiou" else "a"
            cap = cap.replace("a ", f"{art} {spec['material']} ", 1)
        return np.array(encode_caption(cap), dtype=np.int64)

    def generate(self, text, seed=0, retries=8, enforce_quality=True):
        label, msg = Q.should_ask(text)
        if label != "accept":
            return None, {"gate": label, "message": msg}
        spec = Q.parse_prompt(text)
        anchor_ids = self._anchor_ids(spec)
        struct_prefix = Q.STRUCTURAL_PREFIX.get(spec.get("cls"))
        rejects = {}
        grid = None
        for k in range(max(1, retries)):
            lb = None
            ids = None
            ptoks = struct_prefix if (struct_prefix is not None and k >= 3) else None
            if k >= 1 and enforce_quality:
                lb = Q.bias_from_spec(spec, strength=1.6 + 0.5 * min(k, 5))
            if k >= 4 and enforce_quality:
                combined = np.zeros((256, 32), dtype=np.float64)
                if lb is not None:
                    combined[:] = np.asarray(lb, dtype=np.float64)[None, :]
                combined += Q.positional_bias_from_spec(
                    spec, strength=1.4 + 0.3 * (k - 4))
                lb = combined
            if k >= 6 and anchor_ids is not None:
                ids = anchor_ids
            grid = self.pipe.generate_grid(text, seed=seed + 1013 * k,
                                           logit_bias=lb,
                                           ids_override=ids,
                                           prefix_tokens=ptoks)
            if not enforce_quality:
                return grid, {"gate": "accept", "seed": seed, "attempt": k,
                              "rejects": {}}
            ok, reasons = Q.check_sprite(grid, spec)
            if ok:
                return grid, {"gate": "accept", "seed": seed + 1013 * k,
                              "attempt": k, "rejects": rejects}
            for r in reasons:
                rejects[r] = rejects.get(r, 0) + 1
        return grid, {"gate": "accept_degraded", "seed": seed,
                      "attempt": retries, "rejects": rejects}

    def grid_rgba(self, grid):
        rgba = np.zeros((16, 16, 4), dtype=np.uint8)
        for idx in range(1, len(self.pipe.palette)):
            m = grid == idx
            rgba[m] = self.pipe.palette[idx]
        return rgba