#!/usr/bin/env python3
"""Try bounded key-buffer sequences against the original Hwanse2.exe runtime."""
from __future__ import annotations

import argparse
import json
import subprocess
import time
from pathlib import Path
from typing import Any

from probe_runtime_input_path import (
    DEFAULT_PREFIX,
    OUT,
    ROOT,
    env_for,
    find_free_display,
    find_hwanse_pid,
    hex32,
    load_selector_contexts,
    loaded_base,
    parse_windows,
    run,
    sample_process,
    selected_pointer_context,
    truncate,
    write_key_buffer,
)


KEY_OFFSETS = {
    "Return": 0x1C,
    "space": 0x39,
    "z": 0x2C,
    "x": 0x2D,
    "c": 0x2E,
    "Backspace": 0x0E,
    "Escape": 0x01,
    "PageUp": 0xC9,
    "PageDown": 0xD1,
    "Up": 0xC8,
    "Down": 0xD0,
    "Left": 0xCB,
    "Right": 0xCD,
}


DEFAULT_SEQUENCES = [
    ("accept", ["Return"]),
    ("accept-accept", ["Return", "Return"]),
    ("accept-z", ["Return", "z"]),
    ("accept-down-accept", ["Return", "Down", "Return"]),
    ("accept-up-accept", ["Return", "Up", "Return"]),
    ("accept-escape", ["Return", "Escape"]),
]


EXPANDED_SEQUENCES = DEFAULT_SEQUENCES + [
    ("space", ["space"]),
    ("z", ["z"]),
    ("x", ["x"]),
    ("c", ["c"]),
    ("accept-space", ["Return", "space"]),
    ("accept-x", ["Return", "x"]),
    ("accept-c", ["Return", "c"]),
    ("accept-left-accept", ["Return", "Left", "Return"]),
    ("accept-right-accept", ["Return", "Right", "Return"]),
    ("accept-down-down-accept", ["Return", "Down", "Down", "Return"]),
    ("accept-accept-accept", ["Return", "Return", "Return"]),
    ("accept-accept-space", ["Return", "Return", "space"]),
    ("accept-accept-z", ["Return", "Return", "z"]),
    ("accept-accept-x", ["Return", "Return", "x"]),
    ("accept-accept-down-accept", ["Return", "Return", "Down", "Return"]),
    ("accept-accept-up-accept", ["Return", "Return", "Up", "Return"]),
    ("accept-accept-left-accept", ["Return", "Return", "Left", "Return"]),
    ("accept-accept-right-accept", ["Return", "Return", "Right", "Return"]),
    ("accept-accept-escape", ["Return", "Return", "Escape"]),
    ("accept-page-down-accept", ["Return", "PageDown", "Return"]),
    ("accept-page-up-accept", ["Return", "PageUp", "Return"]),
    ("accept-backspace", ["Return", "Backspace"]),
]


SEQUENCE_SETS = {
    "default": DEFAULT_SEQUENCES,
    "expanded": EXPANDED_SEQUENCES,
}


def context_for_sample(sample: dict[str, Any], roots: list[int], contexts: dict[int, dict]) -> dict[str, Any] | None:
    return selected_pointer_context(sample.get("selectedPointerStaticHex"), roots, contexts)


def unique_selectors(rows: list[dict[str, Any]]) -> list[str]:
    values = []
    seen = set()
    for row in rows:
        selector = (row.get("selectorContext") or {}).get("selector")
        if selector and selector not in seen:
            values.append(selector)
            seen.add(selector)
    return values


def run_input_path_prelude(windows: list[str], env: dict[str, str], pid: int | None, base: int | None) -> list[dict[str, Any]]:
    rows = []
    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)
        for key_name in ["Return", "space", "z", "Down"]:
            run(["xdotool", "keydown", "--window", window, "--clearmodifiers", key_name], env)
            time.sleep(0.25)
            held_sample = sample_process(pid, base) if pid and base else {"readOk": False}
            run(["xdotool", "keyup", "--window", window, "--clearmodifiers", key_name], env)
            run(["xdotool", "key", "--window", window, "--clearmodifiers", key_name], env)
            time.sleep(0.4)
            event_sample = sample_process(pid, base) if pid and base else {"readOk": False}
            rows.append({
                "window": window,
                "key": key_name,
                "heldSample": held_sample,
                "eventSample": event_sample,
            })
    return rows


def run_sequence(
    name: str,
    keys: list[str],
    *,
    startup_wait: float,
    hold: float,
    gap: float,
    prelude: str,
    prefix: Path,
    roots: list[int],
    contexts: dict[int, dict],
) -> dict[str, Any]:
    display = find_free_display()
    env = env_for(prefix, display)
    prefix.mkdir(parents=True, exist_ok=True)
    xvfb = subprocess.Popen(
        ["Xvfb", display, "-screen", "0", "1280x1024x24"],
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        errors="replace",
    )
    time.sleep(1)
    wine = subprocess.Popen(
        ["wine", "explorer", "/desktop=hwanse,640x480", "Hwanse2.exe"],
        cwd=ROOT,
        env=env,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        errors="replace",
    )
    time.sleep(startup_wait)
    pid = find_hwanse_pid()
    base = loaded_base(pid) if pid else None
    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", "")])
    focus_rows = []
    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)
        focus_rows.append(window)
        time.sleep(0.05)
    prelude_rows = run_input_path_prelude(windows, env, pid, base) if prelude == "input-path" else []
    initial = sample_process(pid, base) if pid and base else {"readOk": False}
    rows: list[dict[str, Any]] = [
        {
            "step": 0,
            "key": "initial",
            "sample": initial,
            "selectorContext": context_for_sample(initial, roots, contexts),
        }
    ]
    for index, key_name in enumerate(keys, start=1):
        offset = KEY_OFFSETS[key_name]
        write = write_key_buffer(pid, base, offset, duration=hold) if pid and base else {"writeOk": False}
        time.sleep(gap)
        sample = sample_process(pid, base) if pid and base else {"readOk": False}
        rows.append({
            "step": index,
            "key": key_name,
            "keyOffsetHex": hex32(offset),
            "write": write,
            "sample": sample,
            "selectorContext": context_for_sample(sample, roots, contexts),
        })
    final = sample_process(pid, base) if pid and base else {"readOk": False}
    rows.append({
        "step": len(keys) + 1,
        "key": "final",
        "sample": final,
        "selectorContext": context_for_sample(final, roots, contexts),
    })
    subprocess.run(["wineserver", "-k"], cwd=ROOT, env=env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, timeout=5)
    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)
    reached_current = any((row.get("sample") or {}).get("selectedPointerEqualsCurrentRoot") for row in rows)
    reached_route_selector = any((row.get("selectorContext") or {}).get("selector") == "2:0" for row in rows)
    return {
        "name": name,
        "keys": keys,
        "display": display,
        "linuxPid": pid,
        "loadedBaseHex": hex32(base),
        "classSearch": class_search,
        "nameSearch": name_search,
        "focusedWindows": focus_rows,
        "prelude": prelude,
        "preludeRowCount": len(prelude_rows),
        "preludeRows": prelude_rows,
        "initialSelectedPointerHex": initial.get("selectedPointerStaticHex"),
        "finalSelectedPointerHex": final.get("selectedPointerStaticHex"),
        "finalSelectorContext": context_for_sample(final, roots, contexts),
        "uniqueSelectorContexts": unique_selectors(rows),
        "reachedCurrentRoot": reached_current,
        "reachedRouteSelectorContext": reached_route_selector,
        "rows": rows,
        "startupOutput": truncate(wine_output or ""),
        "xvfbOutput": truncate(xvfb_output or ""),
    }


def build_summary(args: argparse.Namespace) -> dict[str, Any]:
    roots, contexts = load_selector_contexts()
    sequence_set = SEQUENCE_SETS[args.sequence_set]
    if args.sequence_name:
        by_name = dict(sequence_set)
        missing = [name for name in args.sequence_name if name not in by_name]
        if missing:
            raise SystemExit(f"unknown sequence name(s): {', '.join(missing)}")
        selected = [(name, by_name[name]) for name in args.sequence_name]
    else:
        selected = sequence_set[: args.max_sequences]
    rows = [
        run_sequence(
            name,
            keys,
            startup_wait=args.startup_wait,
            hold=args.hold,
            gap=args.gap,
            prelude=args.prelude,
            prefix=args.prefix,
            roots=roots,
            contexts=contexts,
        )
        for name, keys in selected
    ]
    any_route = any(row.get("reachedRouteSelectorContext") for row in rows)
    any_current = any(row.get("reachedCurrentRoot") for row in rows)
    return {
        "objective": "bounded runtime key-buffer sequence probe for original Hwanse2.exe",
        "winePrefix": str(args.prefix),
        "startupWaitSeconds": args.startup_wait,
        "holdSeconds": args.hold,
        "gapSeconds": args.gap,
        "prelude": args.prelude,
        "sequenceSet": args.sequence_set,
        "availableSequenceCount": len(sequence_set),
        "selectedSequenceNames": [name for name, _keys in selected],
        "sequenceCount": len(rows),
        "anyReachedCurrentRoot": any_current,
        "anyReachedRouteSelectorContext": any_route,
        "sequences": rows,
        "promotionStatus": "blocked",
        "conclusion": (
            "These direct key-buffer sequences are runtime diagnostics only. They can show whether a bounded "
            "input path reaches selector 2:0/current root, but they do not replace a strict map hotspot or "
            "watchpoint trace unless that route selector is actually observed."
        ),
    }


def markdown(summary: dict[str, Any]) -> str:
    lines = [
        "# Runtime Key Sequence Probe",
        "",
        f"- sequence count: {summary.get('sequenceCount')}",
        f"- sequence set: `{summary.get('sequenceSet') or '-'}`",
        f"- available sequences: {summary.get('availableSequenceCount')}",
        f"- selected sequences: {', '.join(summary.get('selectedSequenceNames') or []) or '-'}",
        f"- prelude: `{summary.get('prelude') or '-'}`",
        f"- any reached current root: {summary.get('anyReachedCurrentRoot')}",
        f"- any reached route selector 2:0: {summary.get('anyReachedRouteSelectorContext')}",
        f"- promotion status: `{summary.get('promotionStatus')}`",
        "",
        summary.get("conclusion") or "",
        "",
        "| sequence | keys | initial | final | final selector | selectors seen | reached 2:0 |",
        "| --- | --- | --- | --- | --- | --- | --- |",
    ]
    for row in summary.get("sequences") or []:
        lines.append(
            "| "
            + " | ".join([
                row.get("name") or "-",
                ", ".join(row.get("keys") or []),
                f"`{row.get('initialSelectedPointerHex') or '-'}`",
                f"`{row.get('finalSelectedPointerHex') or '-'}`",
                f"`{(row.get('finalSelectorContext') or {}).get('selector') or '-'}`",
                ", ".join(row.get("uniqueSelectorContexts") or []) or "-",
                str(row.get("reachedRouteSelectorContext")),
            ])
            + " |"
        )
    lines.append("")
    return "\n".join(lines)


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


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--startup-wait", type=float, default=14.0)
    parser.add_argument("--hold", type=float, default=0.75)
    parser.add_argument("--gap", type=float, default=0.75)
    parser.add_argument("--sequence-set", choices=sorted(SEQUENCE_SETS), default="default")
    parser.add_argument("--sequence-name", action="append", help="Run one named sequence from the selected sequence set. Can be repeated.")
    parser.add_argument("--max-sequences", type=int, default=len(DEFAULT_SEQUENCES))
    parser.add_argument("--prelude", choices=["focus-only", "input-path"], default="focus-only")
    parser.add_argument("--prefix", type=Path, default=DEFAULT_PREFIX)
    parser.add_argument("--out-dir", type=Path, default=OUT)
    parser.add_argument(
        "--output-prefix",
        default="runtime_key_sequence_probe",
        help="Output filename stem. Use a different value for supplemental probes.",
    )
    args = parser.parse_args()
    summary = build_summary(args)
    write_outputs(summary, args.out_dir, args.output_prefix)
    print(f"wrote runtime key sequence probe -> {args.out_dir / (args.output_prefix + '.json')}")


if __name__ == "__main__":
    main()