#!/usr/bin/env python3
"""Poll live field-map layer buffers under Wine.

This is a narrow runtime fallback for the map-animation producer problem.  The
static reports prove the consumer side, but not the writer that changes the
visible animated layer0 cells.  This probe does not need WineDbg; it starts the
EXE under Xvfb, samples relocated globals through /proc/<pid>/mem, and records
map dimensions/checksums plus selected animated cell values.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
import struct
import subprocess
import time
from pathlib import Path
from typing import Any

from probe_runtime_input_path import (
    DEFAULT_PREFIX,
    IMAGE_BASE,
    OUT,
    ROOT,
    env_for,
    find_free_display,
    find_hwanse_pid,
    hex32,
    loaded_base,
    parse_windows,
    run,
    truncate,
    write_key_buffer,
)
from probe_runtime_key_sequences import KEY_OFFSETS


LIVE_LAYER0 = 0x00595AF0
LIVE_LAYER1 = 0x0058D7D0
MAP_WIDTH = 0x00595ADA
MAP_HEIGHT = 0x00595ADC
CAMERA_X = 0x004576DC
CAMERA_Y = 0x004576DE
DIRTY_GRID = 0x005957D0

KNOWN_ANIMATED_MAPS = {
    (37, 48): "map1_01a-or-same-size",
    (100, 96): "map1_02b-or-same-size",
}


def static_to_runtime(static_va: int, base: int) -> int:
    return base + (static_va - IMAGE_BASE)


def process_maps(pid: int) -> list[tuple[int, int]]:
    rows: list[tuple[int, int]] = []
    try:
        lines = (Path("/proc") / str(pid) / "maps").read_text(encoding="utf-8", errors="replace").splitlines()
    except OSError:
        return rows
    for line in lines:
        parts = line.split(maxsplit=5)
        if not parts:
            continue
        try:
            start_text, end_text = parts[0].split("-", 1)
            rows.append((int(start_text, 16), int(end_text, 16)))
        except ValueError:
            continue
    return rows


def is_mapped(ranges: list[tuple[int, int]], address: int, size: int = 1) -> bool:
    return any(start <= address and address + size <= end for start, end in ranges)


def read_bytes(mem: Any, ranges: list[tuple[int, int]], address: int, size: int) -> bytes | None:
    if size <= 0 or not is_mapped(ranges, address, size):
        return None
    try:
        mem.seek(address)
        data = mem.read(size)
    except OSError:
        return None
    return data if len(data) == size else None


def read_u16(mem: Any, ranges: list[tuple[int, int]], address: int) -> int | None:
    data = read_bytes(mem, ranges, address, 2)
    return struct.unpack("<H", data)[0] if data else None


def checksum(data: bytes | None) -> str | None:
    if data is None:
        return None
    return hashlib.sha1(data).hexdigest()[:16]


def focus_windows(env: dict[str, str]) -> dict[str, Any]:
    if not shutil.which("xdotool"):
        return {"xdotoolPath": None, "focusedWindows": []}
    class_search = run(["xdotool", "search", "--class", "hwanse2"], env)
    name_search = run(["xdotool", "search", "--name", "hwanse"], env)
    windows = parse_windows([class_search.get("output", ""), name_search.get("output", "")])
    focused = []
    for window in windows[:8]:
        run(["xdotool", "windowactivate", "--sync", window], env)
        run(["xdotool", "windowfocus", window], env)
        run(["xdotool", "mousemove", "--window", window, "320", "240", "click", "1"], env)
        focused.append(window)
        time.sleep(0.04)
    return {
        "xdotoolPath": shutil.which("xdotool"),
        "classSearch": class_search,
        "nameSearch": name_search,
        "focusedWindows": focused,
    }


def capture_screenshot(env: dict[str, str], path: Path) -> dict[str, Any]:
    if not shutil.which("import"):
        return {"status": "unavailable", "importPath": None}
    path.parent.mkdir(parents=True, exist_ok=True)
    result = run(["import", "-window", "root", str(path)], env)
    result["path"] = str(path)
    result["exists"] = path.exists()
    return result


def kill_wineserver(env: dict[str, str]) -> dict[str, Any]:
    try:
        result = subprocess.run(
            ["wineserver", "-k"],
            cwd=ROOT,
            env=env,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            errors="replace",
            timeout=8,
        )
        return {"status": result.returncode, "output": truncate(result.stdout or "")}
    except subprocess.TimeoutExpired as exc:
        return {"status": "timeout", "output": truncate((exc.stdout or "") if isinstance(exc.stdout, str) else "")}


def parse_key_sequence(raw: str) -> list[tuple[str, float]]:
    if not raw:
        return []
    rows: list[tuple[str, float]] = []
    for token in raw.split(","):
        token = token.strip()
        if not token:
            continue
        if ":" in token:
            key, delay_text = token.split(":", 1)
            delay = float(delay_text)
        else:
            key, delay = token, 0.35
        if key not in KEY_OFFSETS:
            raise SystemExit(f"unknown key {key!r}")
        rows.append((key, delay))
    return rows


def maybe_send_keys(
    *,
    pid: int | None,
    base: int | None,
    key_rows: list[tuple[str, float]],
    next_key_index: int,
    elapsed: float,
    key_times: list[float],
) -> tuple[int, list[dict[str, Any]]]:
    events: list[dict[str, Any]] = []
    while next_key_index < len(key_rows) and elapsed >= key_times[next_key_index]:
        key, hold = key_rows[next_key_index]
        if pid and base:
            write = write_key_buffer(pid, base, KEY_OFFSETS[key], duration=hold)
        else:
            write = {"writeOk": False, "error": "pid/base unavailable"}
        events.append(
            {
                "timeMs": int(elapsed * 1000),
                "key": key,
                "holdSeconds": hold,
                "write": write,
            }
        )
        next_key_index += 1
    return next_key_index, events


def words_from_bytes(data: bytes | None, limit: int | None = None) -> list[int] | None:
    if data is None:
        return None
    count = len(data) // 2
    if limit is not None:
        count = min(count, limit)
    return list(struct.unpack_from(f"<{count}H", data, 0))


def sample(pid: int, base: int, t0: float, *, capture_words: bool = False) -> dict[str, Any]:
    ranges = process_maps(pid)
    row: dict[str, Any] = {
        "timeMs": int((time.time() - t0) * 1000),
        "readOk": False,
    }
    try:
        with (Path("/proc") / str(pid) / "mem").open("rb", buffering=0) as mem:
            width = read_u16(mem, ranges, static_to_runtime(MAP_WIDTH, base))
            height = read_u16(mem, ranges, static_to_runtime(MAP_HEIGHT, base))
            camera_x = read_u16(mem, ranges, static_to_runtime(CAMERA_X, base))
            camera_y = read_u16(mem, ranges, static_to_runtime(CAMERA_Y, base))
            tile_count = int(width or 0) * int(height or 0)
            sane = 0 < tile_count <= 24000
            layer0 = read_bytes(mem, ranges, static_to_runtime(LIVE_LAYER0, base), tile_count * 2) if sane else None
            layer1 = read_bytes(mem, ranges, static_to_runtime(LIVE_LAYER1, base), tile_count * 2) if sane else None
            dirty = read_bytes(mem, ranges, static_to_runtime(DIRTY_GRID, base), 37 * 21) if sane else None
            animated_count = 0
            animated_samples: list[dict[str, Any]] = []
            if layer0 and layer1:
                for index in range(tile_count):
                    flag = struct.unpack_from("<H", layer1, index * 2)[0]
                    if flag & 0x40:
                        animated_count += 1
                        if len(animated_samples) < 24:
                            tile = struct.unpack_from("<H", layer0, index * 2)[0]
                            animated_samples.append(
                                {
                                    "index": index,
                                    "x": index % int(width),
                                    "y": index // int(width),
                                    "tile": tile,
                                    "flagHex": f"0x{flag:04x}",
                                }
                            )
            row.update(
                {
                    "readOk": width is not None and height is not None,
                    "width": width,
                    "height": height,
                    "mapSizeLabel": KNOWN_ANIMATED_MAPS.get((int(width or 0), int(height or 0))),
                    "cameraX": camera_x,
                    "cameraY": camera_y,
                    "tileCount": tile_count,
                    "saneMapSize": sane,
                    "layer0Sha1": checksum(layer0),
                    "layer1Sha1": checksum(layer1),
                    "dirtySha1": checksum(dirty),
                    "animatedFlagCellCount": animated_count,
                    "animatedSamples": animated_samples,
                }
            )
            if capture_words and sane:
                row["layer0Words"] = words_from_bytes(layer0)
                row["layer1Words"] = words_from_bytes(layer1)
    except OSError as exc:
        row["error"] = str(exc)
    return row


def build_summary(args: argparse.Namespace) -> dict[str, Any]:
    prefix: Path = args.prefix
    prefix.mkdir(parents=True, exist_ok=True)
    display = find_free_display()
    env = env_for(prefix, display)
    initial_kill = kill_wineserver(env)
    xvfb = subprocess.Popen(
        ["Xvfb", display, "-screen", "0", "1280x1024x24"],
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        errors="replace",
    )
    time.sleep(1.0)
    wine = subprocess.Popen(
        ["wine", "explorer", "/desktop=hwanse,640x480", "Hwanse2.exe"],
        cwd=ROOT,
        env=env,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        errors="replace",
    )
    t0 = time.time()
    key_rows = parse_key_sequence(args.keys)
    next_key_index = 0
    key_times: list[float] = []
    key_schedule_base: float | None = None
    pid = None
    base = None
    focus = {}
    samples: list[dict[str, Any]] = []
    key_events: list[dict[str, Any]] = []
    try:
        deadline = t0 + args.duration
        while time.time() < deadline:
            if pid is None:
                pid = find_hwanse_pid()
                base = loaded_base(pid) if pid else None
                if pid and base and not focus:
                    focus = focus_windows(env)
                    key_schedule_base = time.time() if args.keys_after_ready else t0
                    cursor = args.startup_wait
                    for _key, hold in key_rows:
                        key_times.append((key_schedule_base - t0) + cursor)
                        cursor += hold + args.key_gap
            elapsed = time.time() - t0
            if key_rows and key_times:
                next_key_index, events = maybe_send_keys(
                    pid=pid,
                    base=base,
                    key_rows=key_rows,
                    next_key_index=next_key_index,
                    elapsed=elapsed,
                    key_times=key_times,
                )
                key_events.extend(events)
            if pid and base:
                samples.append(sample(pid, base, t0, capture_words=args.capture_words))
            time.sleep(args.interval)
    finally:
        screenshot = capture_screenshot(env, args.screenshot) if args.screenshot else {}
        final_kill = kill_wineserver(env)
        try:
            wine_output, _ = wine.communicate(timeout=5)
        except subprocess.TimeoutExpired:
            wine.kill()
            wine_output, _ = wine.communicate(timeout=5)
        xvfb.terminate()
        try:
            xvfb_output, _ = xvfb.communicate(timeout=5)
        except subprocess.TimeoutExpired:
            xvfb.kill()
            xvfb_output, _ = xvfb.communicate(timeout=5)

    unique_sizes = []
    seen_sizes = set()
    unique_layer0 = []
    seen_layer0 = set()
    animated_samples_by_hash: dict[str, list[dict[str, Any]]] = {}
    for row in samples:
        size = (row.get("width"), row.get("height"))
        if size not in seen_sizes:
            seen_sizes.add(size)
            unique_sizes.append(
                {
                    "width": row.get("width"),
                    "height": row.get("height"),
                    "firstTimeMs": row.get("timeMs"),
                    "label": row.get("mapSizeLabel"),
                }
            )
        layer0_hash = row.get("layer0Sha1")
        if layer0_hash and layer0_hash not in seen_layer0:
            seen_layer0.add(layer0_hash)
            unique_layer0.append(
                {
                    "sha1": layer0_hash,
                    "timeMs": row.get("timeMs"),
                    "width": row.get("width"),
                    "height": row.get("height"),
                    "animatedFlagCellCount": row.get("animatedFlagCellCount"),
                }
            )
        if row.get("animatedFlagCellCount"):
            animated_samples_by_hash.setdefault(str(layer0_hash), row.get("animatedSamples") or [])

    layer0_changed = len(seen_layer0) > 1
    reached_known_animated_size = any(row.get("label") for row in unique_sizes)
    return {
        "kind": "hwanse-runtime-map-layer0-poll",
        "status": (
            "layer0-changing-runtime-sample"
            if layer0_changed
            else "runtime-readable-no-layer0-change"
            if samples
            else "runtime-process-not-sampled"
        ),
        "winePrefix": str(prefix),
        "display": display,
        "durationSeconds": args.duration,
        "intervalSeconds": args.interval,
        "startupWaitSeconds": args.startup_wait,
        "keysAfterReady": args.keys_after_ready,
        "keys": args.keys,
        "linuxPid": pid,
        "loadedBaseHex": hex32(base),
        "initialWineserverKill": initial_kill,
        "finalWineserverKill": final_kill,
        "focus": focus,
        "screenshot": screenshot,
        "keyEvents": key_events,
        "sampleCount": len(samples),
        "readOkCount": sum(1 for row in samples if row.get("readOk")),
        "uniqueMapSizes": unique_sizes,
        "uniqueLayer0States": unique_layer0[:40],
        "uniqueLayer0StateCount": len(seen_layer0),
        "layer0Changed": layer0_changed,
        "reachedKnownAnimatedSize": reached_known_animated_size,
        "animatedSamplesByLayer0Hash": animated_samples_by_hash,
        "samples": samples,
        "startupOutput": truncate(wine_output or ""),
        "xvfbOutput": truncate(xvfb_output or ""),
        "conclusion": (
            "Runtime map buffers were sampled and layer0 changed. Inspect samples to bind the writer/time window."
            if layer0_changed
            else "Runtime map buffers were sampled, but no layer0 change was observed in the bounded window."
            if samples
            else "Hwanse2.exe was not sampled in this bounded runtime run."
        ),
    }


def write_outputs(report: dict[str, Any], out_dir: Path) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "runtime_map_layer0_poll.json").write_text(
        json.dumps(report, ensure_ascii=False, separators=(",", ":")) + "\n",
        encoding="utf-8",
    )


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--prefix", type=Path, default=DEFAULT_PREFIX)
    parser.add_argument("--duration", type=float, default=24.0)
    parser.add_argument("--interval", type=float, default=0.05)
    parser.add_argument("--startup-wait", type=float, default=4.0)
    parser.add_argument("--keys", default="", help="comma-separated key[:holdSeconds] entries")
    parser.add_argument("--keys-after-ready", action="store_true", help="start key schedule after pid/base is readable")
    parser.add_argument("--key-gap", type=float, default=0.4)
    parser.add_argument("--screenshot", type=Path, default=None, help="optional final X root screenshot path")
    parser.add_argument("--capture-words", action="store_true", help="include live layer0/layer1 word arrays in sane samples")
    parser.add_argument("--out-dir", type=Path, default=OUT)
    args = parser.parse_args()
    report = build_summary(args)
    write_outputs(report, args.out_dir)
    print(f"runtime_map_layer0_poll -> {args.out_dir / 'runtime_map_layer0_poll.json'}")
    print(json.dumps({k: report.get(k) for k in ["status", "sampleCount", "readOkCount", "uniqueLayer0StateCount", "layer0Changed", "uniqueMapSizes", "loadedBaseHex"]}, ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())