"""
NeuralOS Hypervisor v3  --  TEXT2OS / TEXT2IMG image operating system.

ZERO DOM POLICY
---------------
The browser is a *dumb framebuffer terminal*: it owns exactly one <canvas> and
nothing else. Every pixel of the OS shell (taskbar, start menu, window chrome,
icons, cursor, overlays) is produced by the backend layer compositor. There is
no HTML/CSS UI to speak of, and no client-side framework.

ARCHITECTURE
------------
    Text2OS  : free text -> intent ops -> kernel mutations
    Text2IMG : prompts   -> raster layers (generated, then cached forever)
    Layers   : LayerStack composites RGBA surfaces into a 1024x1024 framebuffer

SPEED MODEL (what replaces the old "regenerate the whole world each frame")
--------------------------------------------------------------------------
The old build ran the UNet + TAE decode for the *entire* canvas on *every*
interaction and shipped a base64 data-URL over JSON. This build instead:

  1. generates each chrome sprite / wallpaper / icon ONCE (disk-cached),
  2. composites in pixel space (O(area), no diffusion in the hot path),
  3. sends only DIRTY TILES as raw WebP over a binary WebSocket (no base64),
  4. refines window content asynchronously, off the input critical path.

That is the concrete answer to "bypass the sequential per-frame pipeline":
generation is amortised, compositing is parallel/tiled, and the wire format is
binary deltas instead of full-frame JSON.
"""

import io
import json
import os
import queue
import struct
import threading
import time
import warnings

warnings.filterwarnings("ignore")

import numpy as np
from PIL import Image, ImageDraw, ImageFont

# --------------------------------------------------------------------------- #
# Config
# --------------------------------------------------------------------------- #

FB = int(os.environ.get("NEURALOS_FB", "1024"))            # framebuffer edge
TASKBAR_H = 40
TITLEBAR_H = 26
CHROME_BORDER = 3
TILE = 64                                                   # dirty-tile granularity (16x16 grid)
# Full-vs-patch decision constants (measured, quality 82): every patch pays a
# fresh WebP header, so many small tiles can cost more than one full frame. We
# therefore compare estimated shipped BYTES, not tile count.
PATCH_OVERHEAD = 700        # bytes of per-image WebP overhead
PX_RATE = 0.019            # encoded bytes per pixel for UI content

# auto | mock | cpu | cuda
ENGINE_MODE = os.environ.get("NEURALOS_ENGINE", "auto").lower()
# 1 => generate chrome sprites with the diffusion model (pure Text2IMG),
# 0 => draw authentic Windows "Luna" chrome procedurally (raster, still not DOM)
GEN_CHROME = os.environ.get("NEURALOS_GEN_CHROME", "0") == "1"
# 0 => when chrome = generative, still verify and fall back on a bad sample
STRICT = os.environ.get("NEURALOS_STRICT", "0") == "1"

SD_MODEL = os.environ.get("NEURALOS_SD_MODEL", "stable-diffusion-v1-5/stable-diffusion-v1-5")
TAE_MODEL = os.environ.get("NEURALOS_TAE_MODEL", "madebyollin/taesd")
LCM_LORA = os.environ.get("NEURALOS_LCM_LORA", "latent-consistency/lcm-lora-sdv1-5")
LLM_MODEL = os.environ.get("NEURALOS_LLM_MODEL", "Qwen/Qwen2.5-Coder-0.5B-Instruct")
CACHE_DIR = os.environ.get("NEURALOS_CACHE", "/tmp/neural_cache")

# Windows "Luna" palette
LUNA_TITLE_A = (0, 84, 227)
LUNA_TITLE_B = (58, 147, 255)
LUNA_TASKBAR_A = (36, 94, 219)
LUNA_TASKBAR_B = (25, 65, 165)
LUNA_FRAME = (0, 84, 227)
LUNA_CLOSE = (226, 69, 54)
LUNA_DESK = (58, 110, 165)


# --------------------------------------------------------------------------- #
# Text / drawing helpers  (all raster -> still image-based, never DOM)
# --------------------------------------------------------------------------- #

_FONT_CACHE = {}


def font(size, bold=False):
    key = (size, bold)
    if key in _FONT_CACHE:
        return _FONT_CACHE[key]
    f = None
    try:
        f = ImageFont.load_default(size=size)          # Pillow >= 10.1 bundled TTF
    except Exception:
        try:
            f = ImageFont.truetype("DejaVuSans-Bold.ttf" if bold else "DejaVuSans.ttf", size)
        except Exception:
            f = ImageFont.load_default()
    _FONT_CACHE[key] = f
    return f


def text_size(draw, s, f):
    try:
        l, t, r, b = draw.textbbox((0, 0), s, font=f)
        return r - l, b - t
    except Exception:
        return len(s) * f.size // 2, f.size


def vgrad(size, top, bottom):
    """Vertical gradient RGB image."""
    w, h = size
    col = np.linspace(np.array(top), np.array(bottom), max(h, 1))
    arr = np.repeat(col[:, None, :], w, axis=1).astype(np.uint8)
    return Image.fromarray(arr, "RGB")


def hgrad(size, left, right):
    w, h = size
    col = np.linspace(np.array(left), np.array(right), max(w, 1))
    arr = np.repeat(col[None, :, :], h, axis=0).astype(np.uint8)
    return Image.fromarray(arr, "RGB")


def round_mask(size, radius):
    m = Image.new("L", size, 0)
    ImageDraw.Draw(m).rounded_rectangle([0, 0, size[0] - 1, size[1] - 1], radius, fill=255)
    return m


def label_text(draw, xy, s, f, fill, shadow=(0, 0, 0)):
    x, y = xy
    draw.text((x + 1, y + 1), s, font=f, fill=shadow)
    draw.text((x, y), s, font=f, fill=fill)


# --------------------------------------------------------------------------- #
# Layer model
# --------------------------------------------------------------------------- #

class Layer:
    """A positioned RGBA raster surface in the compositor."""

    __slots__ = ("name", "surface", "x", "y", "z", "opacity", "dirty")

    def __init__(self, name, surface, x, y, z, opacity=255):
        self.name = name
        self.surface = surface          # PIL.Image RGBA
        self.x = x
        self.y = y
        self.z = z
        self.opacity = opacity
        self.dirty = True


class LayerStack:
    """Painter's-algorithm compositor with dirty-rect tracking.

    Composite is pure pixel maths -> constant cost, no diffusion.
    """

    def __init__(self, size):
        self.size = size
        self.layers = {}                # name -> Layer
        self._dirty = []
        self.mark_all_dirty()

    def put(self, name, surface, x=0, y=0, z=0, opacity=255):
        """Insert/update a layer. Only marks dirty when something actually changed,
        so an unchanged cached sprite costs nothing and patches stay small."""
        lay = self.layers.get(name)
        if lay is None:
            lay = Layer(name, surface, x, y, z, opacity)
            self.layers[name] = lay
            self.mark_dirty(lay)
            return lay
        if (lay.surface is not surface or lay.x != x or lay.y != y
                or lay.z != z or lay.opacity != opacity):
            # mark BOTH footprints: the new one to paint, the old one to erase
            self._dirty.append((lay.x, lay.y, lay.surface.width, lay.surface.height))
            lay.surface, lay.x, lay.y, lay.z, lay.opacity = surface, x, y, z, opacity
            self.mark_dirty(lay)
        return lay

    def remove(self, name):
        lay = self.layers.pop(name, None)
        if lay is not None:
            self.mark_dirty(lay)
        return lay

    def mark_dirty(self, lay):
        self._dirty.append((lay.x, lay.y, lay.surface.width, lay.surface.height))

    def mark_rect(self, x, y, w, h):
        self._dirty.append((x, y, w, h))

    def mark_all_dirty(self):
        self._dirty.append((0, 0, self.size[0], self.size[1]))

    def take_dirty_rects(self):
        rects = self._dirty
        self._dirty = []
        return rects

    def composite(self, clip=None):
        """Compose layers into an RGB image, optionally clipped to a rect."""
        w, h = self.size
        if clip is None:
            base = Image.new("RGB", (w, h), LUNA_DESK)
            ox, oy, cw, ch = 0, 0, w, h
        else:
            ox, oy, cw, ch = clip
            base = Image.new("RGB", (cw, ch), LUNA_DESK)
        for lay in sorted(self.layers.values(), key=lambda l: l.z):
            s = lay.surface
            # intersect layer rect with clip
            lx0, ly0 = lay.x, lay.y
            lx1, ly1 = lx0 + s.width, ly0 + s.height
            cx0, cy0 = max(lx0, ox), max(ly0, oy)
            cx1, cy1 = min(lx1, ox + cw), min(ly1, oy + ch)
            if cx0 >= cx1 or cy0 >= cy1:
                continue
            crop = s.crop((cx0 - lx0, cy0 - ly0, cx1 - lx0, cy1 - ly0))
            if lay.opacity < 255:
                a = crop.getchannel("A").point(lambda v: v * lay.opacity // 255)
                crop = crop.copy()
                crop.putalpha(a)
            base.paste(crop, (cx0 - ox, cy0 - oy), crop)
        return base


def merge_rects(rects, size):
    """Coalesce dirty rects into tile-aligned, non-overlapping-ish rects."""
    w, h = size
    tiles = set()
    for (x, y, rw, rh) in rects:
        x0 = max(0, x // TILE) * TILE
        y0 = max(0, y // TILE) * TILE
        x1 = min(w, ((x + rw + TILE - 1) // TILE) * TILE)
        y1 = min(h, ((y + rh + TILE - 1) // TILE) * TILE)
        for ty in range(y0, y1, TILE):
            for tx in range(x0, x1, TILE):
                tiles.add((tx, ty))
    return [(tx, ty, min(TILE, w - tx), min(TILE, h - ty)) for (tx, ty) in sorted(tiles)]


# --------------------------------------------------------------------------- #
# Persistent asset store  (generate once, reuse forever)
# --------------------------------------------------------------------------- #

class AssetStore:
    def __init__(self):
        self.mem = {}
        self.hits = 0
        self.misses = 0
        try:
            import diskcache
            self.disk = diskcache.Cache(CACHE_DIR)
        except Exception:
            self.disk = None

    def get(self, key):
        if key in self.mem:
            self.hits += 1
            return self.mem[key]
        if self.disk is not None:
            v = self.disk.get(key)
            if v is not None:
                self.hits += 1
                img = Image.open(io.BytesIO(v))
                img.load()
                img = img.convert("RGBA")
                self.mem[key] = img
                return img
        self.misses += 1
        return None

    def put(self, key, img):
        img = img.convert("RGBA")
        self.mem[key] = img
        if self.disk is not None:
            buf = io.BytesIO()
            img.save(buf, format="PNG")
            self.disk.set(key, buf.getvalue())
        return img

    def get_or_make(self, key, maker):
        got = self.get(key)
        return got if got is not None else self.put(key, maker())


# --------------------------------------------------------------------------- #
# Neural engines
# --------------------------------------------------------------------------- #

class MockNeuralEngine:
    """Deterministic procedural stand-in. No model, no network.

    Guarantees the fallback path stays deterministic (required for a chat/runtime
    system): same key -> same image, every run.
    """

    name = "mock"
    device = "cpu"

    def text2img(self, prompt, size, seed=0, steps=1):
        w, h = size
        rng = np.random.default_rng(abs(hash((prompt, seed))) % (2 ** 32))
        y = np.linspace(0, 255, max(h, 1))[:, None]
        x = np.linspace(0, 255, max(w, 1))[None, :]
        r = ((x * 0.5 + y * 0.2) % 255)
        g = ((y * 0.7 + seed) % 255)
        b = ((x * 0.3 + 128) % 255)
        arr = np.stack(np.broadcast_arrays(r, g, b), axis=-1) + rng.integers(0, 24, (h, w, 3))
        img = Image.fromarray(np.clip(arr, 0, 255).astype(np.uint8), "RGB").convert("RGBA")
        d = ImageDraw.Draw(img)
        d.text((4, 4), f"IMG {w}x{h}", font=font(11), fill=(255, 255, 255, 255))
        return img

    def text2text(self, prompt, max_tokens=64):
        return ""

    def stats(self):
        return {"engine": self.name, "device": self.device}


class DiffusionEngine:
    """SD1.5 + LCM-LoRA (1-step) + Tiny AutoEncoder, CPU/GPU aware."""

    name = "diffusion"

    def __init__(self):
        import torch
        from diffusers import StableDiffusionPipeline, AutoencoderTiny, LCMScheduler

        self.torch = torch
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.dtype = torch.float16 if self.device == "cuda" else torch.float32
        torch.set_num_threads(max(1, (os.cpu_count() or 4)))
        if self.device == "cuda":
            torch.backends.cudnn.benchmark = True

        t0 = time.time()
        pipe = StableDiffusionPipeline.from_pretrained(
            SD_MODEL, torch_dtype=self.dtype, safety_checker=None,
            requires_safety_checker=False, low_cpu_mem_usage=False,
        )
        try:
            pipe.load_lora_weights(LCM_LORA)
            pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
            self.lcm = True
        except Exception:
            self.lcm = False
        try:
            pipe.vae = AutoencoderTiny.from_pretrained(TAE_MODEL, torch_dtype=self.dtype)
        except Exception:
            pass
        pipe = pipe.to(self.device)
        if self.device == "cuda":
            pipe.enable_attention_slicing()
        self.pipe = pipe
        self.load_s = time.time() - t0

        self.llm = self.tok = None
        try:
            from transformers import AutoModelForCausalLM, AutoTokenizer
            self.tok = AutoTokenizer.from_pretrained(LLM_MODEL)
            if self.tok.pad_token_id is None:
                self.tok.pad_token_id = self.tok.eos_token_id
            self.llm = AutoModelForCausalLM.from_pretrained(LLM_MODEL, torch_dtype=self.dtype).to(self.device)
        except Exception:
            pass

    def text2img(self, prompt, size, seed=0, steps=1):
        torch = self.torch
        w, h = size
        w = max(64, int(round(w / 64.0)) * 64)
        h = max(64, int(round(h / 64.0)) * 64)
        g = torch.Generator(device="cpu").manual_seed(int(seed) & 0x7FFFFFFF)
        with torch.inference_mode():
            out = self.pipe(
                prompt, width=w, height=h, num_inference_steps=max(1, steps),
                guidance_scale=1.0, generator=g, output_type="pil",
            ).images[0]
        return out.convert("RGBA")

    def text2text(self, prompt, max_tokens=64):
        if self.llm is None:
            return ""
        torch = self.torch
        enc = self.tok(prompt, return_tensors="pt", truncation=True).to(self.device)
        with torch.inference_mode():
            out = self.llm.generate(
                enc.input_ids, attention_mask=enc.attention_mask,
                max_new_tokens=max_tokens, do_sample=False,
                pad_token_id=self.tok.eos_token_id,
            )
        return self.tok.decode(out[0][enc.input_ids.shape[1]:], skip_special_tokens=True).strip()

    def stats(self):
        return {"engine": self.name, "device": self.device, "lcm": self.lcm,
                "load_s": round(getattr(self, "load_s", 0), 2)}


ENGINE = None
ENGINE_LOCK = threading.Lock()


def get_engine():
    global ENGINE
    with ENGINE_LOCK:
        if ENGINE is not None:
            return ENGINE
        if ENGINE_MODE == "mock":
            ENGINE = MockNeuralEngine()
            return ENGINE
        try:
            ENGINE = DiffusionEngine()
        except Exception as e:                                  # deterministic fallback
            print(f"[!] Diffusion engine unavailable ({e}); using deterministic mock.")
            ENGINE = MockNeuralEngine()
        return ENGINE


# --------------------------------------------------------------------------- #
# Text2IMG: sprite + texture generation
# --------------------------------------------------------------------------- #

class SpriteLibrary:
    """Generative-first raster assets, memoised through AssetStore."""

    ICON_PROMPTS = {
        "notepad": "windows xp notepad application icon, yellow notepad with blue pen, 64x64 pixel art, crisp",
        "paint": "windows xp ms paint icon, colorful palette and brush, 64x64 pixel art, crisp",
        "cmd": "windows xp command prompt icon, black screen with white prompt, 64x64 pixel art, crisp",
        "explorer": "windows xp file explorer icon, manila folder with magnifier, 64x64 pixel art, crisp",
        "browser": "windows xp internet explorer icon, blue lowercase e with orbit, 64x64 pixel art, crisp",
    }

    def __init__(self, store):
        self.store = store

    # -- chrome (taskbar / frame / menu / cursor) --------------------------- #
    def taskbar(self):
        return self.store.get_or_make("chrome.taskbar.v3", self._taskbar)

    def _taskbar(self):
        if GEN_CHROME:
            try:
                img = get_engine().text2img(
                    "windows xp luna taskbar bar, glossy blue gradient, start orb left, "
                    "system tray right, seamless horizontal UI strip, flat, 1024x40",
                    (FB, TASKBAR_H), seed=101, steps=2)
                if img.width >= FB // 2:
                    return self._ensure_taskbar(img)
            except Exception:
                pass
        bar = vgrad((FB, TASKBAR_H), LUNA_TASKBAR_B, LUNA_TASKBAR_A).convert("RGBA")
        d = ImageDraw.Draw(bar)
        d.line([(0, 0), (FB, 0)], fill=(92, 156, 255, 255), width=2)
        d.line([(0, 1), (FB, 1)], fill=(20, 50, 140, 255))
        return bar

    def _ensure_taskbar(self, img):
        return img.resize((FB, TASKBAR_H))

    def start_button(self):
        return self.store.get_or_make("chrome.start.v3", self._start_button)

    def _start_button(self):
        w, h = 92, TASKBAR_H - 6
        img = vgrad((w, h), (63, 162, 63), (37, 115, 37)).convert("RGBA")
        img.putalpha(round_mask((w, h), 12))
        d = ImageDraw.Draw(img)
        for i in range(6):
            d.arc([6 + i, 4 + i, w - 6 - i, h - 4 - i], 250, 290, fill=(120, 220, 120, 120))
        d.text((14, h // 2 - 9), "\u229e", font=font(20, True), fill=(255, 255, 255, 255))
        label_text(d, (38, h // 2 - 9), "start", font(19, True), (255, 255, 255, 255))
        return img

    def start_menu(self):
        return self.store.get_or_make("chrome.menu.v3", self._start_menu)

    def _start_menu(self):
        w, h = 380, 470
        img = Image.new("RGBA", (w, h), (0, 0, 0, 0))
        body = vgrad((w, h), (214, 223, 247), (236, 240, 250)).convert("RGBA")
        img.paste(body, (0, 0), round_mask((w, h), 10))
        d = ImageDraw.Draw(img)
        d.rounded_rectangle([0, 0, w - 1, h - 1], 10, outline=(0, 84, 227, 255), width=3)
        hdr = vgrad((w - 6, 54), (92, 156, 255), (36, 94, 219)).convert("RGBA")
        img.paste(hdr, (3, 3))
        label_text(d, (66, 12), "NeuralOS", font(20, True), (255, 255, 255, 255))
        d.text((66, 34), "Administrator", font=font(12), fill=(230, 238, 255, 255))
        av = vgrad((44, 44), (250, 220, 170), (200, 160, 110)).convert("RGBA")
        av.putalpha(round_mask((44, 44), 6))
        img.paste(av, (14, 8), av)
        d.rectangle([0, 57, 58, h - 1], fill=(36, 94, 219, 255))
        side = "NeuralOS"
        sf = font(13, True)
        sy = h - 30
        for ch in reversed(side):
            d.text((20, sy), ch, font=sf, fill=(255, 255, 255, 255))
            sy -= 15
        y = 70
        for key, name in PROGRAM_ORDER:
            d.rectangle([62, y + 4, w - 8, y + 44], fill=(255, 255, 255, 0))
            label_text(d, (70, y + 14), name, font(15), (0, 0, 0, 255))
            y += 52
        return img

    @staticmethod
    def menu_item_rects():
        """Deterministic menu hit geometry (must not depend on cache state)."""
        return [(key, 62, 70 + i * 52 + 4, 380 - 70, 40)
                for i, (key, _) in enumerate(PROGRAM_ORDER)]

    def window_frame(self, w, h, title, active=True):
        key = f"chrome.frame.{w}x{h}.{int(active)}"
        return self.store.get_or_make(key, lambda: self._draw_frame(w, h, title, active))

    def _draw_frame(self, w, h, title, active):
        img = Image.new("RGBA", (w, h), (0, 0, 0, 0))
        if GEN_CHROME:
            try:
                gen = get_engine().text2img(
                    "windows xp luna window frame chrome, blue glossy title bar, "
                    "rounded top corners, close minimize maximize buttons, empty white client area, "
                    "flat UI, sharp", (w, h), seed=200 + (w * 7 + h), steps=2)
                if gen.width == w and gen.height == h:
                    return gen
            except Exception:
                pass
        tA, tB = (LUNA_TITLE_B, LUNA_TITLE_A) if active else ((160, 180, 210), (120, 140, 175))
        d = ImageDraw.Draw(img)
        d.rounded_rectangle([0, 0, w - 1, h - 1], 8, outline=(20, 60, 170, 255), width=1)
        d.rectangle([2, 2, w - 3, h - 3], fill=(236, 236, 236, 255))
        bar = vgrad((w - 2 * CHROME_BORDER, TITLEBAR_H), tA, tB).convert("RGBA")
        bar.putalpha(round_mask((bar.width, TITLEBAR_H), 6))
        img.paste(bar, (CHROME_BORDER, CHROME_BORDER), bar)
        f = font(13, True)
        tw, th = text_size(d, title, f)
        label_text(d, (9, CHROME_BORDER + (TITLEBAR_H - th) // 2 - 1), title, f, (255, 255, 255, 255))
        # window controls
        by = CHROME_BORDER + 5
        bx = w - CHROME_BORDER - 18
        for kind, col in (("close", LUNA_CLOSE), ("max", (60, 120, 220)), ("min", (60, 120, 220))):
            img.paste(self._ctrl(kind, col), (bx, by), self._ctrl(kind, col))
            bx -= 20
        return img

    def _ctrl(self, kind, col):
        s = 16
        img = Image.new("RGBA", (s, s), (0, 0, 0, 0))
        d = ImageDraw.Draw(img)
        base = vgrad((s, s), tuple(min(255, c + 60) for c in col), col).convert("RGBA")
        base.putalpha(round_mask((s, s), 4))
        img.paste(base, (0, 0), base)
        d.rounded_rectangle([0, 0, s - 1, s - 1], 4, outline=(255, 255, 255, 90), width=1)
        if kind == "close":
            d.line([4, 4, s - 5, s - 5], fill=(255, 255, 255, 255), width=2)
            d.line([s - 5, 4, 4, s - 5], fill=(255, 255, 255, 255), width=2)
        elif kind == "max":
            d.rectangle([4, 4, s - 5, s - 5], outline=(255, 255, 255, 255), width=2)
        else:
            d.rectangle([4, s - 6, s - 5, s - 5], fill=(255, 255, 255, 255))
        return img

    def icon(self, app):
        key = f"chrome.icon.{app}.v3"
        return self.store.get_or_make(key, lambda: self._icon(app))

    def _icon(self, app):
        if GEN_CHROME or True:
            try:
                img = get_engine().text2img(self.ICON_PROMPTS.get(app, app), (64, 64), seed=300 + len(app), steps=2)
                img = img.resize((64, 64))
                # alpha: keep it square-opaque, drop near-black corners for a sprite feel
                return img
            except Exception:
                pass
        return self._icon_proc(app)

    def _icon_proc(self, app):
        col = {"notepad": (60, 110, 220), "paint": (230, 150, 40),
               "cmd": (30, 30, 30), "explorer": (240, 200, 70),
               "browser": (40, 150, 220)}.get(app, (120, 120, 120))
        img = Image.new("RGBA", (64, 64), (0, 0, 0, 0))
        d = ImageDraw.Draw(img)
        b = vgrad((56, 56), tuple(min(255, c + 50) for c in col), col).convert("RGBA")
        b.putalpha(round_mask((56, 56), 8))
        img.paste(b, (4, 4), b)
        d.rounded_rectangle([4, 4, 59, 59], 8, outline=(255, 255, 255, 120), width=2)
        glyph = {"notepad": "\u2261", "paint": "\U0001f58c", "cmd": ">_",
                 "explorer": "\U0001f4c1", "browser": "\U0001f310"}.get(app, "?")
        d.text((20, 18), glyph, font=font(18, True), fill=(255, 255, 255, 255))
        return img

    def wallpaper(self):
        return self.store.get_or_make("chrome.wallpaper.v3", self._wallpaper)

    def _wallpaper(self):
        if GEN_CHROME:
            try:
                img = get_engine().text2img(
                    "windows xp bliss wallpaper, rolling green hill, blue sky, white clouds, "
                    "photographic, 1024x1024", (FB, FB), seed=7, steps=2)
                return img.resize((FB, FB))
            except Exception:
                pass
        # restrained procedural "bliss": sky gradient + green hill
        img = vgrad((FB, FB), (80, 150, 235), (170, 210, 245)).convert("RGB")
        d = ImageDraw.Draw(img)
        pts = [(0, FB)]
        for x in range(0, FB + 1, 32):
            y = int(FB * 0.62 + 60 * np.sin(x / 160.0) - x * 0.05)
            pts.append((x, y))
        pts.append((FB, FB))
        d.polygon(pts, fill=(90, 160, 60))
        d.polygon([(p[0], p[1] + 40) for p in pts] + [(FB, FB), (0, FB)], fill=(60, 120, 40))
        return img.convert("RGBA")

    def cursor(self):
        return self.store.get_or_make("chrome.cursor.v3", self._cursor)

    def _cursor(self):
        s = 24
        img = Image.new("RGBA", (s, s), (0, 0, 0, 0))
        d = ImageDraw.Draw(img)
        d.polygon([(0, 0), (0, 16), (4, 12), (7, 19), (10, 17), (7, 11), (13, 11)],
                  fill=(255, 255, 255, 255), outline=(0, 0, 0, 255))
        return img


PROGRAM_ORDER = [("notepad", "Notepad"), ("paint", "Paint"), ("cmd", "Command Prompt"),
                 ("explorer", "File Explorer"), ("browser", "Internet Browser")]


# --------------------------------------------------------------------------- #
# LASM  --  Latent Assembly VM (image opcodes over layers)
# --------------------------------------------------------------------------- #

class LatentVM:
    """Text2IMG-oriented bytecode: LSET/LCOPY/LADD/LNOISE/LFILL.

    Operates on the LayerStack so 'programs' are literally image edits.
    """

    def __init__(self, stack, sprites):
        self.stack = stack
        self.sprites = sprites

    def run(self, program):
        log = []
        for op in program:
            code = op[0]
            try:
                if code == "LSET":
                    _, layer, x, y, prompt, w, h = op
                    img = get_engine().text2img(prompt, (int(w), int(h)), seed=hash(prompt) & 0xFFFF, steps=2)
                    self.stack.put(layer, img.convert("RGBA"), int(x), int(y), 20)
                    log.append(f"LSET {layer} {w}x{h}")
                elif code == "LCOPY":
                    _, src, dst, x, y = op
                    s = self.stack.layers[src].surface
                    self.stack.put(dst, s.copy(), int(x), int(y), 21)
                    log.append(f"LCOPY {src}->{dst}")
                elif code == "LADD":
                    _, layer, prompt = op
                    add = get_engine().text2img(prompt, (128, 128), seed=hash(prompt) & 0xFFFF, steps=2)
                    base = self.stack.layers[layer].surface
                    blend = Image.blend(base.convert("RGB"), add.convert("RGB").resize(base.size), 0.5)
                    self.stack.put(layer, blend.convert("RGBA"), self.stack.layers[layer].x,
                                   self.stack.layers[layer].y, self.stack.layers[layer].z)
                    log.append(f"LADD {layer}")
                elif code == "LNOISE":
                    _, layer, amount = op
                    lay = self.stack.layers[layer]
                    a = np.asarray(lay.surface.convert("RGB")).astype(np.int16)
                    n = np.random.randint(-int(amount), int(amount) + 1, a.shape)
                    out = Image.fromarray(np.clip(a + n, 0, 255).astype(np.uint8)).convert("RGBA")
                    self.stack.put(layer, out, lay.x, lay.y, lay.z)
                    log.append(f"LNOISE {layer} {amount}")
                elif code == "LFILL":
                    _, layer, rgba = op
                    self.stack.put(layer, Image.new("RGBA", (FB, FB), tuple(rgba)), 0, 0, 1)
                    log.append(f"LFILL {layer}")
                else:
                    log.append(f"?{code}")
            except Exception as e:
                log.append(f"ERR {code}: {e}")
        return log


# --------------------------------------------------------------------------- #
# OS kernel
# --------------------------------------------------------------------------- #

DEFAULT_SIZE = {"notepad": (430, 330), "paint": (520, 400), "cmd": (470, 320),
                "explorer": (560, 420), "browser": (620, 460)}

APP_PROMPTS = {
    "notepad": "windows notepad document, white page, black monospace text, blue menu bar",
    "paint": "ms paint canvas, white area, color palette, tool box, brushes",
    "cmd": "command prompt, black background, white monospace text, C:\\> prompt",
    "explorer": "windows explorer, folder tree, file icons, toolbar, address bar",
    "browser": "web browser window, address bar, tabs, page content",
}


class Window:
    __slots__ = ("pid", "app", "name", "x", "y", "w", "h", "z", "refine", "content", "dirty")

    def __init__(self, pid, app, name, x, y, w, h, z):
        self.pid, self.app, self.name = pid, app, name
        self.x, self.y, self.w, self.h, self.z = x, y, w, h, z
        self.refine = 0
        self.content = None          # RGBA client-area raster
        self.dirty = True

    def rect(self):
        return (self.x, self.y, self.w, self.h)

    def client(self):
        return (self.x + CHROME_BORDER, self.y + TITLEBAR_H + CHROME_BORDER,
                self.w - 2 * CHROME_BORDER, self.h - TITLEBAR_H - 2 * CHROME_BORDER)

    def ctrl_rects(self):
        by = self.y + CHROME_BORDER + 5
        bx = self.x + self.w - CHROME_BORDER - 18
        out = {}
        for kind in ("close", "max", "min"):
            out[kind] = (bx, by, 16, 16)
            bx -= 20
        return out


class OSKernel:
    """Text2OS state machine. No DOM knowledge whatsoever."""

    def __init__(self, sprites):
        self.sprites = sprites
        self.windows = {}
        self.next_pid = 1
        self.focused = None
        self.menu_open = False
        self.selected_icon = None
        self.text_buffer = ""          # Text2OS command line
        self.text_cursor_on = True
        self.overlays = []             # list of (text, expiry); capped below
        self.desk_icons = [(k, n) for k, n in PROGRAM_ORDER]
        self.refine_q = queue.Queue()
        self.mouse = (FB // 2, FB // 2)

    MAX_OVERLAYS = 4

    def log(self, text, ttl=4.0):
        """Append a transient pixel log line, bounding the stack so it can
        never grow past the corner it is drawn in."""
        self.overlays.append((text, time.time() + ttl))
        del self.overlays[:-self.MAX_OVERLAYS]

    # -- window ops --------------------------------------------------------- #
    def open(self, app, x=None, y=None):
        if app not in DEFAULT_SIZE:
            return None
        w, h = DEFAULT_SIZE[app]
        n = len(self.windows)
        x = x if x is not None else 120 + (n % 5) * 42
        y = y if y is not None else 90 + (n % 5) * 36
        x = max(0, min(x, FB - w))
        y = max(0, min(y, FB - TASKBAR_H - h))
        pid = self.next_pid
        self.next_pid += 1
        win = Window(pid, app, app.capitalize(), x, y, w, h, pid)
        self.windows[pid] = win
        self.focus(pid)
        self.refine_q.put(pid)
        return win

    def close(self, pid):
        self.windows.pop(pid, None)
        if self.focused == pid:
            self.focused = max(self.windows, key=lambda p: self.windows[p].z) if self.windows else None

    def focus(self, pid):
        if pid in self.windows:
            self.focused = pid
            mx = max((w.z for w in self.windows.values()), default=0)
            self.windows[pid].z = mx + 1

    def move(self, pid, x, y):
        win = self.windows.get(pid)
        if not win:
            return
        win.x = max(-win.w + 60, min(x, FB - 60))
        win.y = max(0, min(y, FB - TASKBAR_H - TITLEBAR_H))
        win.dirty = True

    # -- hit testing (pixel space) ----------------------------------------- #
    def start_button_rect(self):
        return (0, FB - TASKBAR_H + 3, 92, TASKBAR_H - 6)

    def taskbar_order(self):
        return sorted(self.windows.values(), key=lambda w: w.pid)

    def taskbar_button_rect(self, i):
        return (100 + i * 172, FB - TASKBAR_H + 4, 168, TASKBAR_H - 8)

    def menu_rect(self):
        w, h = 380, 470
        return (0, FB - TASKBAR_H - h + 2, w, h)

    def icon_rect(self, idx):
        return (16, 20 + idx * 82, 64, 78)

    def hit(self, x, y):
        # topmost window first
        for win in sorted(self.windows.values(), key=lambda w: -w.z):
            wx, wy, ww, wh = win.rect()
            if wx <= x < wx + ww and wy <= y < wy + wh:
                for kind, (cx, cy, cw, ch) in win.ctrl_rects().items():
                    if cx <= x < cx + cw and cy <= y < cy + ch:
                        return ("ctrl", win, kind)
                if wy <= y < wy + TITLEBAR_H + CHROME_BORDER:
                    return ("titlebar", win, "drag")
                return ("client", win, None)
        if self.menu_open:
            mx, my, mw, mh = self.menu_rect()
            if mx <= x < mx + mw and my <= y < my + mh:
                for key, ix, iy, iw, ih in self.sprites.menu_item_rects():
                    if ix <= x - mx < ix + iw and iy <= y - my < iy + ih:
                        return ("menu_item", None, key)
                return ("menu", None, None)
            self.menu_open = False
        sx, sy, sw, sh = self.start_button_rect()
        if sx <= x < sx + sw and sy <= y < sy + sh:
            return ("start", None, None)
        if y >= FB - TASKBAR_H:
            for i, win in enumerate(self.taskbar_order()):
                bx, by, bw, bh = self.taskbar_button_rect(i)
                if bx <= x < bx + bw and by <= y < by + bh:
                    return ("taskbar_btn", win, None)
            return ("taskbar", None, None)
        for idx in range(len(self.desk_icons)):
            ix, iy, iw, ih = self.icon_rect(idx)
            if ix <= x < ix + iw and iy <= y < iy + ih:
                return ("icon", None, self.desk_icons[idx][0])
        return ("desktop", None, None)


# --------------------------------------------------------------------------- #
# Text2OS  --  free text -> intent ops
# --------------------------------------------------------------------------- #

INTENT_SCHEMA = (
    'Return ONLY JSON: {"ops":[{"op":"open|close|focus|move|wallpaper|lasm|none",'
    '"app":"notepad|paint|cmd|explorer|browser","pid":0,"x":0,"y":0}]}'
)


def keyword_intents(text):
    t = text.lower()
    ops = []
    for key, _ in PROGRAM_ORDER:
        if key in t or ("command" in t and key == "cmd"):
            ops.append({"op": "open", "app": key})
    if "close" in t:
        ops.append({"op": "close_all"})
    if not ops:
        ops.append({"op": "none", "text": text})
    return ops


def parse_intents(text):
    """Text2OS: deterministic keywords first (instant), LLM only for the
    ambiguous remainder. Keeps common commands off the model entirely."""
    kw = keyword_intents(text)
    if not any(o.get("op") == "none" for o in kw):
        return kw
    eng = get_engine()
    try:
        raw = eng.text2text(
            f'You are a Text2OS intent parser. {INTENT_SCHEMA}\nUser: "{text}"\nJSON:', max_tokens=96)
        if raw:
            s = raw[raw.find("{"): raw.rfind("}") + 1]
            data = json.loads(s)
            if isinstance(data, dict) and data.get("ops"):
                return data["ops"]
    except Exception:
        pass
    return kw


# --------------------------------------------------------------------------- #
# Compositor
# --------------------------------------------------------------------------- #

class Compositor:
    def __init__(self, kernel, sprites, store):
        self.k = kernel
        self.s = sprites
        self.store = store
        self.stack = LayerStack((FB, FB))
        self.last = None
        self._memo = {}          # signature -> raster, avoids per-frame realloc
        self.metrics = {"frames": 0, "full": 0, "patch": 0, "idle": 0,
                        "build_ms": 0.0, "frame_ms": 0.0, "tiles": 0, "bytes": 0}

    def _cached(self, sig, maker):
        """Return a stable raster for a content signature, so unchanged UI
        layers keep object identity and never dirty the tile grid."""
        got = self._memo.get(sig)
        if got is None:
            got = self._memo[sig] = maker()
            if len(self._memo) > 512:
                self._memo.clear()
                self._memo[sig] = got
        return got

    # -- build the layer stack from state ---------------------------------- #
    def build(self):
        k = self.k
        desired = set()

        def P(name, *a, **kw):
            desired.add(name)
            return self.stack.put(name, *a, **kw)

        P("wallpaper", self.s.wallpaper(), 0, 0, 1)
        for idx, (app, name) in enumerate(k.desk_icons):
            ix, iy, iw, ih = k.icon_rect(idx)
            sel = k.selected_icon == app
            P(f"icon.{app}", self._cached(("icon", app, sel), lambda a=app, n=name, s=sel: self._icon_tile(a, n, s)), ix, iy, 5)
        for win in sorted(k.windows.values(), key=lambda w: w.z):
            P(f"win.{win.pid}", self.s.window_frame(
                win.w, win.h, win.name, active=(k.focused == win.pid)), win.x, win.y, 10 + win.z)
            cx, cy, cw, ch = win.client()
            if win.content is not None:
                content = win.content
            else:
                content = self._cached(("placeholder", win.pid, cw, ch), lambda w=cw, h=ch: self._placeholder(w, h))
            P(f"content.{win.pid}", content, cx, cy, 10 + win.z + 1)
        P("taskbar", self.s.taskbar(), 0, FB - TASKBAR_H, 60)
        P("startbtn", self.s.start_button(), 4, FB - TASKBAR_H + 3, 61)
        for i, win in enumerate(k.taskbar_order()):
            bx, by, bw, bh = k.taskbar_button_rect(i)
            act = k.focused == win.pid
            P(f"taskbtn.{win.pid}", self._cached(
                ("taskbtn", win.pid, win.name, win.app, act),
                lambda w=win, a=act: self._task_button(w, a)), bx, by, 62)
        clock = time.strftime("%I:%M %p")
        P("tray", self._cached(("tray", clock), lambda: self._tray(clock)), FB - 260, FB - TASKBAR_H + 2, 63)
        if k.menu_open:
            mx, my, mw, mh = k.menu_rect()
            P("menu", self.s.start_menu(), mx, my, 70)
        P("cmdbar", self._cached(("cmdbar",), self._cmdbar), 0, self.CMD_Y, 80)
        if k.text_buffer:
            P("cmdtext", self._cached(("cmdtext", k.text_buffer),
                                      lambda t=k.text_buffer: self._cmdtext(t)),
              self.CMD_X, self.CMD_Y, 81)
        if k.text_cursor_on:
            caret_x = min(FB - 14, self.CMD_X + self.text_w(k.text_buffer) + 2)
            P("caret", self._cached(("caret",), self._caret), caret_x, self.CMD_Y + 4, 82)
        for i, (text, ttl) in enumerate(k.overlays):
            P(f"overlay.{i}", self._cached(("overlay", text),
                                            lambda t=text: self._overlay(t)),
              FB - 380, 16 + i * 30, 90)
        P("cursor", self.s.cursor(), k.mouse[0], k.mouse[1], 99)

        # reconcile: drop layers whose window/overlay no longer exists
        for name in list(self.stack.layers):
            if name not in desired:
                self.stack.remove(name)

    def _placeholder(self, w, h):
        img = Image.new("RGBA", (w, h), (245, 245, 245, 255))
        ImageDraw.Draw(img).text((10, 10), "generating content...", font=font(12),
                                 fill=(120, 120, 120, 255))
        return img

    def _task_button(self, win, active):
        w, h = 168, TASKBAR_H - 8
        img = vgrad((w, h), (70, 130, 220), (40, 85, 175)) if active else \
            vgrad((w, h), (60, 110, 190), (35, 70, 150))
        img = img.convert("RGBA")
        img.putalpha(round_mask((w, h), 5))
        d = ImageDraw.Draw(img)
        ic = self.s.icon(win.app).resize((18, 18))
        img.paste(ic, (6, 6), ic)
        label_text(d, (30, h // 2 - 8), win.name[:14], font(13, True), (255, 255, 255, 255))
        return img

    def _icon_tile(self, app, name, selected):
        w, h = 64, 78
        tile = Image.new("RGBA", (w, h), (0, 0, 0, 0))
        if selected:
            hl = Image.new("RGBA", (w, h), (60, 120, 220, 90))
            tile.paste(hl, (0, 0), hl)
        ic = self.s.icon(app).resize((48, 48))
        tile.paste(ic, (8, 2), ic)
        d = ImageDraw.Draw(tile)
        f = font(11)
        tw, _ = text_size(d, name, f)
        label_text(d, ((w - tw) // 2, 54), name, f, (255, 255, 255, 255))
        return tile

    def _tray(self, clock=None):
        clock = clock or time.strftime("%I:%M %p")
        img = Image.new("RGBA", (256, TASKBAR_H - 4), (0, 0, 0, 0))
        d = ImageDraw.Draw(img)
        f = font(13)
        tw, th = text_size(d, clock, f)
        d.text((256 - tw - 10, 10), clock, font=f, fill=(255, 255, 255, 255))
        return img

    CMD_Y = FB - TASKBAR_H - 26
    CMD_X = 78

    def _cmdbar(self):
        """Static command-bar base: never dirties again once cached."""
        img = Image.new("RGBA", (FB, 26), (12, 12, 12, 235))
        ImageDraw.Draw(img).text((8, 6), "Text2OS>", font=font(13, True), fill=(120, 255, 120, 255))
        return img

    @staticmethod
    def text_w(s):
        d = ImageDraw.Draw(Image.new("RGB", (1, 1)))
        return text_size(d, s, font(13))[0]

    def _cmdtext(self, s):
        """Only as wide as the typed text, so a keystroke dirties 1-2 tiles
        instead of the whole 1024px bar."""
        w = max(8, self.text_w(s) + 6)
        img = Image.new("RGBA", (w, 26), (12, 12, 12, 0))
        ImageDraw.Draw(img).text((0, 6), s, font=font(13), fill=(220, 255, 220, 255))
        return img

    def _caret(self):
        """The blinking caret is its own 10x18 layer so a blink dirties ONE
        tile instead of the full-width command bar."""
        return Image.new("RGBA", (10, 18), (170, 255, 170, 255))

    def _overlay(self, text):
        d0 = ImageDraw.Draw(Image.new("RGB", (1, 1)))
        img = Image.new("RGBA", (360, 26), (0, 0, 0, 200))
        ImageDraw.Draw(img).text((8, 6), text[:48], font=font(13), fill=(0, 255, 120, 255))
        return img

    # -- render ------------------------------------------------------------- #
    def render(self):
        t0 = time.time()
        self.build()
        rects = merge_rects(self.stack.take_dirty_rects(), (FB, FB))
        t1 = time.time()
        if self.last is not None and not rects:
            self.metrics["idle"] += 1
            self.metrics["tiles"] = 0
            return ("none", None, None)
        # A full frame is one image; a patch set is N images, each with its own
        # WebP header. Ship whichever is estimated to cost fewer bytes.
        dirty_px = sum(w * h for (_, _, w, h) in rects)
        est_patch = PATCH_OVERHEAD * len(rects) + PX_RATE * dirty_px
        est_full = PATCH_OVERHEAD + PX_RATE * FB * FB
        full = self.last is None or est_patch >= 0.9 * est_full
        if full:
            frame = self.stack.composite()
            payload = self._encode(frame)
            self.last = frame
            self.metrics["full"] += 1
            self.metrics["bytes"] = len(payload)
            self.metrics["px"] = FB * FB
            out = ("full", None, payload)
        else:
            patches = []
            nbytes = 0
            for (x, y, w, h) in rects:
                enc = self._encode(self.stack.composite(clip=(x, y, w, h)))
                nbytes += len(enc)
                patches.append((x, y, w, h, enc))
            self.metrics["patch"] += 1
            self.metrics["bytes"] = nbytes
            self.metrics["px"] = sum(w * h for (_, _, w, h) in rects)
            out = ("patch", patches, None)
        t2 = time.time()
        self.metrics["frames"] += 1
        self.metrics["build_ms"] = round((t1 - t0) * 1000, 2)
        self.metrics["frame_ms"] = round((t2 - t1) * 1000, 2)
        self.metrics["tiles"] = 0 if full else len(rects)
        return out

    def _encode(self, img):
        buf = io.BytesIO()
        img.convert("RGB").save(buf, format="WebP", quality=82, method=0)
        return buf.getvalue()


# --------------------------------------------------------------------------- #
# Binary wire protocol
# --------------------------------------------------------------------------- #

MAGIC = 0x4E
T_FULL, T_PATCH, T_LOG = 0, 1, 2


def pack(msg_type, header, payload=b""):
    h = json.dumps(header).encode("utf-8")
    return struct.pack(">BB I", MAGIC, msg_type, len(h)) + h + payload


def send_log(ws, text):
    ws.send(pack(T_LOG, {"text": text}))


# --------------------------------------------------------------------------- #
# Runtime
# --------------------------------------------------------------------------- #

STORE = AssetStore()
SPRITES = SpriteLibrary(STORE)
KERNEL = OSKernel(SPRITES)
COMPOSITOR = Compositor(KERNEL, SPRITES, STORE)
STATE_LOCK = threading.RLock()
ENGINE_READY = threading.Event()


def refinement_worker():
    """Async content generation + progressive refinement, off the input path."""
    eng = get_engine()
    ENGINE_READY.set()
    while True:
        pid = KERNEL.refine_q.get()                 # blocking, no busy-wait
        with STATE_LOCK:
            win = KERNEL.windows.get(pid)
        if win is None:
            continue
        try:
            cx, cy, cw, ch = win.client()
            steps = 2 if win.refine == 0 else 3
            prompt = f"windows xp {APP_PROMPTS.get(win.app, win.app)}, crisp ui, detailed"
            img = eng.text2img(prompt, (cw, ch), seed=win.pid * 31 + win.refine, steps=steps)
            with STATE_LOCK:
                win = KERNEL.windows.get(pid)
                if win is None:
                    continue
                win.content = img.resize((cw, ch))
                win.dirty = True
                win.refine += 1
                COMPOSITOR.stack.mark_rect(win.x, win.y, win.w, win.h)
                KERNEL.log(f"{win.name} refined (step {win.refine})")
                if win.refine < 3:
                    KERNEL.refine_q.put(pid)
        except Exception as e:
            print("[refine]", e)


def _frame_stats():
    """Metrics + kernel state shipped in every frame header.

    Deliberately does NOT call get_engine(): during a cold boot that blocks on
    model load, and the first frame must not wait behind it."""
    stats = dict(COMPOSITOR.metrics)
    stats["cache"] = {"hits": STORE.hits, "misses": STORE.misses}
    stats["engine"] = ENGINE.stats() if ENGINE is not None else \
        {"engine": "loading", "device": "-", "lcm": False}
    stats["windows"] = len(KERNEL.windows)
    stats["menu"] = KERNEL.menu_open
    stats["buffer"] = KERNEL.text_buffer
    return stats


def send_full_to(ws):
    """Push a complete frame to ONE just-connected client.

    The compositor streams a single delta stream shared by every client, so a
    client that joins after the first frame would otherwise only ever receive
    the tiles the cursor/caret happens to dirty -- most of its canvas stays
    blank ("black tiles until mouseover"). Each new socket therefore starts
    from a full snapshot. Chrome is procedural, so this is cheap and never
    touches the model."""
    with STATE_LOCK:
        COMPOSITOR.build()                    # layers must exist before compositing
        COMPOSITOR.stack.take_dirty_rects()   # we are shipping the whole state
        frame = COMPOSITOR.stack.composite()
        blob = COMPOSITOR._encode(frame)
        COMPOSITOR.last = frame
        COMPOSITOR.metrics["full"] += 1
        COMPOSITOR.metrics["bytes"] = len(blob)
        COMPOSITOR.metrics["px"] = FB * FB
        COMPOSITOR.metrics["tiles"] = 0
        stats = _frame_stats()
    ws.send(pack(T_FULL, {"w": FB, "h": FB, **stats}, blob))


def broadcast_dirty(clients):
    with STATE_LOCK:
        for (text, ttl) in list(KERNEL.overlays):
            if ttl < time.time():
                KERNEL.overlays.remove((text, ttl))
        result = COMPOSITOR.render()
        stats = _frame_stats()
    kind, payload, blob = result
    if kind == "none":
        return result
    if kind == "full":
        data = pack(T_FULL, {"w": FB, "h": FB, **stats}, blob)
    else:
        data = [pack(T_PATCH, {"x": x, "y": y, "w": w, "h": h, **stats}, enc)
                for (x, y, w, h, enc) in payload]

    def _deliver(ws):
        if isinstance(data, list):
            for p in data:
                ws.send(p)
        else:
            ws.send(data)
    for ws in list(clients):
        try:
            _deliver(ws)
        except Exception:
            clients.discard(ws)
    return result


def handle_client_message(msg, clients):
    t = msg.get("type")
    with STATE_LOCK:
        if t in ("click", "dblclick"):
            x, y = int(msg.get("x", 0)), int(msg.get("y", 0))
            KERNEL.mouse = (x, y)
            target = KERNEL.hit(x, y)
            kind, obj, extra = target
            # All state below is reflected in layers, and LayerStack.put marks
            # both old and new footprints, so no blanket full-frame invalidation
            # is needed: every one of these interactions ships only its own tiles.
            if kind == "ctrl":
                if extra == "close":
                    KERNEL.close(obj.pid)
                elif extra == "min":
                    obj.dirty = True
                    KERNEL.focus(max(KERNEL.windows) if KERNEL.windows else None)
                elif extra == "max":
                    obj.x, obj.y = 0, 0
                    obj.w, obj.h = FB, FB - TASKBAR_H
            elif kind == "titlebar":
                KERNEL.focus(obj.pid)
                KERNEL.drag = (obj.pid, x - obj.x, y - obj.y)
            elif kind == "icon":
                KERNEL.selected_icon = extra
                if t == "dblclick":
                    KERNEL.open(extra)
            elif kind == "menu_item":
                KERNEL.menu_open = False
                KERNEL.open(extra)
            elif kind == "menu":
                pass
            elif kind == "start":
                KERNEL.menu_open = not KERNEL.menu_open
            elif kind == "desktop":
                KERNEL.selected_icon = None
                KERNEL.menu_open = False
            elif kind == "taskbar_btn":
                KERNEL.focus(obj.pid)
        elif t == "move":
            x, y = int(msg.get("x", 0)), int(msg.get("y", 0))
            KERNEL.mouse = (x, y)
            drag = getattr(KERNEL, "drag", None)
            if drag:
                pid, ox, oy = drag
                KERNEL.move(pid, x - ox, y - oy)
        elif t == "release":
            KERNEL.drag = None
        elif t == "key":
            key = msg.get("key", "")
            if key == "Enter":
                # The intent parse may call the LLM for tens of seconds on CPU.
                # Never do that under STATE_LOCK: hand it to a worker thread so
                # the render loop and the heartbeat keep flowing.
                text = KERNEL.text_buffer
                KERNEL.text_buffer = ""
                KERNEL.log(f"Text2OS: {text[:40]}", ttl=6)
                _spawn(apply_text2os, text)
            elif key == "Backspace":
                KERNEL.text_buffer = KERNEL.text_buffer[:-1]
            elif len(key) == 1:
                KERNEL.text_buffer = (KERNEL.text_buffer + key)[-120:]
        elif t == "lasm":
            _spawn(apply_lasm, msg.get("program", []))
        elif t == "wallpaper":
            _spawn(apply_wallpaper, msg.get("prompt", "abstract"))
    broadcast_dirty(clients)


def _spawn(fn, *args):
    """Run slow, generation-bound work off the input/render critical path."""
    threading.Thread(target=fn, args=args, daemon=True).start()


def apply_text2os(text):
    ops = parse_intents(text)          # slow (LLM) -> done without the lock
    with STATE_LOCK:
        for op in ops:
            if op.get("op") == "open" and op.get("app") in DEFAULT_SIZE:
                KERNEL.open(op["app"])
            elif op.get("op") == "close_all":
                for pid in list(KERNEL.windows):
                    KERNEL.close(pid)
    broadcast_dirty(CLIENTS)


def apply_lasm(program):
    vm = LatentVM(COMPOSITOR.stack, SPRITES)
    logs = vm.run(program)             # may run the model per opcode
    with STATE_LOCK:
        for l in logs:
            KERNEL.log(l, ttl=5)
    broadcast_dirty(CLIENTS)


def apply_wallpaper(prompt):
    img = get_engine().text2img(prompt, (FB, FB), seed=int(time.time()), steps=2)
    STORE.put("chrome.wallpaper.v3", img)
    with STATE_LOCK:
        KERNEL.log("wallpaper regenerated", ttl=5)
    broadcast_dirty(CLIENTS)


def heartbeat(clients):
    """Cursor/clock tick: small dirty rect, proves patch path."""
    while True:
        time.sleep(1.0)
        if not CLIENTS:
            continue
        with STATE_LOCK:
            KERNEL.text_cursor_on = not KERNEL.text_cursor_on
            KERNEL.overlays = [(t, e) for (t, e) in KERNEL.overlays if e > time.time()]
            # The caret is its own layer, so flipping it on/off already dirties
            # exactly its footprint. The tray clock invalidates itself through
            # the memo key when the displayed minute actually changes. No
            # blanket rect marking is required here.
        try:
            broadcast_dirty(clients)
        except Exception as e:
            print("[beat]", e)


# --------------------------------------------------------------------------- #
# Flask app
# --------------------------------------------------------------------------- #

from flask import Flask, jsonify, send_file  # noqa: E402

app = Flask(__name__)
CLIENTS = set()

try:
    from flask_sock import Sock
    sock = Sock(app)

    @sock.route("/kernel")
    def kernel_ws(ws):
        CLIENTS.add(ws)
        try:
            send_log(ws, "Kernel attached")
            # Paint the whole desktop for this client FIRST: chrome is
            # procedural, so the first frame need not wait behind engine load.
            send_full_to(ws)
            get_engine()
            broadcast_dirty(CLIENTS)          # then stream deltas as usual
            while True:
                data = ws.receive()
                if not data:
                    break
                try:
                    msg = json.loads(data)
                except Exception:
                    continue
                handle_client_message(msg, CLIENTS)
        except Exception as e:
            print("[ws]", e)
        finally:
            CLIENTS.discard(ws)
except Exception as e:                                       # pragma: no cover
    print("[!] flask_sock unavailable:", e)


@app.route("/")
def index():
    here = os.path.dirname(os.path.abspath(__file__))
    return send_file(os.path.join(here, "index.html"))


@app.route("/frame.png")
def frame_png():
    """Debug/verification: the exact current composited framebuffer."""
    with STATE_LOCK:
        img = COMPOSITOR.stack.composite()
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    buf.seek(0)
    return send_file(buf, mimetype="image/png")


@app.route("/metrics")
def metrics():
    with STATE_LOCK:
        return jsonify(COMPOSITOR.metrics)


@app.route("/healthz")
def healthz():
    return jsonify({"ok": True, "engine": get_engine().stats(), "fb": FB})


def main():
    threading.Thread(target=refinement_worker, daemon=True).start()
    threading.Thread(target=heartbeat, args=(CLIENTS,), daemon=True).start()
    port = int(os.environ.get("PORT", "7860"))
    app.run(host="0.0.0.0", port=port, threaded=True)


if __name__ == "__main__":
    main()