"""Drive the live OS through CDP and verify Windows-like chrome from the
client's OWN framebuffer (exact pixels, no screenshot letterbox maths).

Opens the Start menu and two windows through the real click/Text2OS interface,
then asserts: active + inactive Luna titlebars, red close button, Start menu
panel, taskbar band and the desktop icon column.
"""
import asyncio
import base64
import io
import json
import os
import subprocess
import sys
import tempfile
import time
import urllib.request

import numpy as np
from PIL import Image
import websockets

CHROME = r"C:\Program Files\Google\Chrome\Application\chrome.exe"
PORT = 7867
CDP_PORT = 9223
BASE = f"http://127.0.0.1:{PORT}"
HERE = os.path.dirname(os.path.abspath(__file__))
FB = 1024

# window geometry must match app.py DEFAULT_SIZE / open() placement
NOTEPAD = (120, 90, 430, 330)
EXPLORER = (162, 126, 560, 420)


def wait_up(url=BASE + "/healthz"):
    for _ in range(80):
        try:
            with urllib.request.urlopen(url, timeout=2) as r:
                json.loads(r.read())
                return True
        except Exception:
            time.sleep(0.25)
    return False


class CDP:
    def __init__(self, ws):
        self.ws = ws
        self.i = 0

    async def call(self, method, **params):
        self.i += 1
        mid = self.i
        await self.ws.send(json.dumps({"id": mid, "method": method, "params": params}))
        while True:
            msg = json.loads(await self.ws.recv())
            if msg.get("id") == mid:
                if "error" in msg:
                    raise RuntimeError(f"{method}: {msg['error']}")
                return msg.get("result", {})

    async def js(self, expr, await_promise=False):
        r = await self.call("Runtime.evaluate", expression=expr,
                            returnByValue=True, awaitPromise=await_promise)
        return r.get("result", {}).get("value")


async def drive():
    pages = []
    last_err = None
    for _ in range(80):
        try:
            with urllib.request.urlopen(f"http://127.0.0.1:{CDP_PORT}/json/list", timeout=2) as r:
                raw = r.read()
            targets = json.loads(raw)
            pages = [t for t in targets if t.get("webSocketDebuggerUrl")
                     and t.get("type") in ("page", "webview")]
            if pages:
                break
            last_err = f"targets={[ (t.get('type'), t.get('url')) for t in targets ]}"
        except Exception as e:
            last_err = f"{type(e).__name__}: {e}"
        time.sleep(0.35)
    if not pages:
        raise SystemExit(f"no CDP page target ({last_err})")

    async with websockets.connect(pages[0]["webSocketDebuggerUrl"], max_size=None) as ws:
        c = CDP(ws)
        await c.call("Runtime.enable")
        for _ in range(120):
            if await c.js("!!window.NeuralOS && window.NeuralOS.stats().frames > 0"):
                break
            await asyncio.sleep(0.25)
        else:
            raise SystemExit("page never received a frame")

        await c.js("window.NeuralOS.send({type:'click',x:40,y:999})")   # Start menu
        await c.js("window.NeuralOS.text2os('open notepad')")
        await c.js("window.NeuralOS.text2os('open explorer')")
        await asyncio.sleep(3.0)

        stats = await c.js("window.NeuralOS.stats()")
        data = await c.js("window.NeuralOS.dump()", await_promise=True)
        assert data and data.startswith("data:image/png;base64,"), "client framebuffer dump failed"
        png = base64.b64decode(data.split(",", 1)[1])
        with open(os.path.join(HERE, "ui_framebuffer.png"), "wb") as f:
            f.write(png)
        img = Image.open(io.BytesIO(png)).convert("RGB")

        # Force a SMALL, patch-only update, then read the VISIBLE <canvas> (not
        # the offscreen buffer). Regression: the partial blit used to clear the
        # whole canvas, blacking out everything except the patched tiles.
        await c.js("window.NeuralOS.send({type:'move',x:300,y:300})")
        await asyncio.sleep(0.8)
        vis = await c.js("document.getElementById('fb').toDataURL('image/png')")
        assert vis and vis.startswith("data:image/png;base64,"), \
            "visible canvas dump failed"
        return img, stats, vis


def reg(a, x, y, w, h):
    return a[y:y + h, x:x + w].astype(int)


def main():
    env = dict(os.environ, NEURALOS_ENGINE="mock", PORT=str(PORT),
               NEURALOS_CACHE="/tmp/neuralos_cdp_cache")
    srv = subprocess.Popen([sys.executable, "app.py"], env=env,
                           stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    chrome = None
    try:
        assert wait_up(), "server did not start"
        profile = tempfile.mkdtemp(prefix="nos-cdp-")
        chromelog = os.path.join(profile, "chrome.log")
        chrome = subprocess.Popen([
            CHROME, "--headless=new", "--no-first-run", "--disable-gpu",
            f"--remote-debugging-port={CDP_PORT}", "--remote-allow-origins=*",
            f"--user-data-dir={profile}",
            "--window-size=1100,1100", BASE],
            stdout=open(chromelog, "wb"), stderr=subprocess.STDOUT)

        try:
            img, stats, vis = asyncio.run(drive())
        except SystemExit as e:
            with open(chromelog, "r", errors="replace") as f:
                print("chrome log:\n" + f.read()[-3000:])
            raise
        client_img = img
        print("client framebuffer:", client_img.size,
              "| windows:", stats.get("windows"),
              "| build_ms:", stats.get("build_ms"), "| frame_ms:", stats.get("frame_ms"),
              "| tiles:", stats.get("tiles"), "| cache:", stats.get("cache"))
        assert client_img.size == (FB, FB), "client framebuffer is not 1024x1024"
        assert stats.get("windows", 0) >= 2, "Text2OS did not open both windows"

        # authoritative server composite
        with urllib.request.urlopen(f"{BASE}/frame.png", timeout=10) as r:
            srv_img = Image.open(io.BytesIO(r.read())).convert("RGB")
        sa = np.asarray(srv_img)
        ca = np.asarray(client_img)
        # WebP q82 is lossy, so compare structure: mean abs error + big-pixel delta
        d = np.abs(sa.astype(int) - ca.astype(int))
        mae = d.mean()
        coarse = (d.max(axis=2) > 64).mean()
        print(f"client vs server: MAE={mae:.2f}/255  coarse-mismatch={coarse:.3%}")
        assert mae < 14, f"client framebuffer diverged from server (MAE {mae:.1f})"
        assert coarse < 0.05, f"structural divergence from server ({coarse:.2%})"

        # --- visible on-screen canvas (what the user actually sees) --------- #
        vpng = base64.b64decode(vis.split(",", 1)[1])
        with open(os.path.join(HERE, "ui_visible.png"), "wb") as f:
            f.write(vpng)
        vimg = Image.open(io.BytesIO(vpng)).convert("RGB")
        va = np.asarray(vimg)
        nonblack = (va.sum(axis=2) > 30).mean()
        print(f"visible canvas {vimg.size}, non-black {nonblack:.1%}")
        # The backing store must be sized to the viewport, not left at the
        # 300x150 <canvas> default and stretched by CSS (that is the "starts
        # zoomed in / resolution too big" bug).
        assert min(vimg.size) >= 900, \
            f"visible canvas backing store is {vimg.size}, not the viewport " \
            "(layout() did not run)"
        assert nonblack > 0.5, \
            "visible canvas went blank after a patch-only update (blit cleared it)"
        vband = va[int(va.shape[0] * 0.90):, :, :]
        blueish = ((vband[:, :, 2].astype(int) > vband[:, :, 0].astype(int) + 25)).mean()
        print(f"visible bottom-blue: {blueish:.1%}")
        assert blueish > 0.2, \
            "visible canvas lost the taskbar (patch blit wiped the screen)"

        a = sa          # assert fidelity against the authoritative composite

        ex, ey, ew, eh = EXPLORER
        nx, ny, nw, nh = NOTEPAD

        # 1. focused window: active Luna titlebar (between title text and buttons)
        tb = reg(a, ex + 90, ey + 6, 400, 18)
        active = ((tb[:, :, 2] > tb[:, :, 0] + 40) & (tb[:, :, 2] > 150)).mean()
        print(f"active titlebar blue: {active:.1%}  mean={tb.reshape(-1,3).mean(0).round(1)}")
        assert active > 0.85, "focused window titlebar is not Luna blue"

        # 2. unfocused window: greyed titlebar, less saturated than the active one
        ntb = reg(a, nx + 250, ny + 6, 150, 18)
        inactive = ntb.reshape(-1, 3).mean(0)
        print(f"inactive titlebar mean: {inactive.round(1)}")
        assert inactive[2] - inactive[0] < 70, "inactive titlebar is not greyed out"

        # 3. close button is red
        cb = reg(a, ex + ew - 3 - 18, ey + 3 + 5, 16, 16)
        red = ((cb[:, :, 0] > cb[:, :, 2] + 60)).mean()
        print(f"close button red: {red:.1%}  mean={cb.reshape(-1,3).mean(0).round(1)}")
        assert red > 0.4, "close button is not red"

        # 4. Start menu panel rendered (open at fb 0,516 380x470)
        menu = reg(a, 120, 560, 220, 160)
        light = (menu.sum(axis=2) > 560).mean()
        print(f"start-menu light: {light:.1%}")
        assert light > 0.5, "start menu panel missing"

        # 5. taskbar band
        bar = reg(a, 200, FB - 36, 700, 30)
        barb = ((bar[:, :, 2] > bar[:, :, 0] + 20)).mean()
        print(f"taskbar blue: {barb:.1%}")
        assert barb > 0.5, "taskbar missing"

        # 6. desktop icon column (glyph pixels vary more than the flat wallpaper)
        icons = reg(a, 16, 20, 64, 400)
        print(f"icon column std: {icons.std(axis=(0, 1)).mean():.1f}")
        assert icons.std(axis=(0, 1)).mean() > 12, "desktop icons missing"

        print("UI CDP OK: titlebars, close button, start menu, taskbar and icons all render")
    finally:
        for p in (chrome, srv):
            if p:
                p.terminate()
                try:
                    p.wait(10)
                except Exception:
                    p.kill()


if __name__ == "__main__":
    main()