#!/usr/bin/env python3
"""Sample the live fixed-point state of the original Win16 game under Wine."""

from __future__ import annotations

import argparse
import csv
import hashlib
import os
import struct
import sys
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Self

ORIGINAL_SHA256 = "a9022f1894e3e6e21fc42e8f6c932f7c549ca77f63aaa0c488bb9d55d9d0174c"
DATA_SIGNATURE = b"IWIKTDK Highscores"
DATA_SIGNATURE_OFFSET = 0x10
DATA_READ_SIZE = 0x900

# Recovered offsets in the original DGROUP segment.
PLAYER_COUNT = 0x07D0
CURRENT_PLAYER = 0x07D1
BALL_X = 0x07D3
BALL_Y = 0x07D7
NEXT_BALL_X = 0x07DB
NEXT_BALL_Y = 0x07DF
CLAW_ACTIVE = 0x07CA
CLAW_FRAME = 0x086D
CLAW_TARGET_FRAME = 0x086F
CLAW_SPRITE_BANK = 0x0873
CLAW_ANIMATION_BLOCKED = 0x0874
PHYSICS_SUBSTEPS = 0x0871

# Recovered offsets in the live game-window object.
MAXIMUM_SPEED = 0x0056
GRAVITY_PER_SUBSTEP = 0x005A
ACTIVE_BALL_COUNT = 0x0061
VELOCITY_X = 0x0BAA
VELOCITY_Y = 0x0BAE
SPIN_REAL48 = 0x0BCE
PREVIOUS_BALL_X = 0x0BBA
PREVIOUS_BALL_Y = 0x0BBE
OBJECT_BALL_X = 0x0BC2
OBJECT_BALL_Y = 0x0BC6
LEFT_FLIPPER_POSITION = 0x0BD4
RIGHT_FLIPPER_POSITION = 0x0BD6
LEFT_FLIPPER_DOWN = 0x0BDA
RIGHT_FLIPPER_DOWN = 0x0BDB
BALL_SUSPENDED = 0x0BDF
BALL_ACTIVE = 0x0BDD
COLLISION_DISABLED = 0x0BD8
OBJECT_READ_SIZE = 0x0BE0


@dataclass(frozen=True)
class MemoryMap:
    start: int
    end: int
    permissions: str
    path: str

    @property
    def size(self) -> int:
        return self.end - self.start


class ProcessMemory:
    def __init__(self, pid: int, *, writable: bool = False) -> None:
        self.pid = pid
        self.maps = self._read_maps()
        try:
            self.stream = open(  # noqa: SIM115 - closed by the context protocol
                f"/proc/{pid}/mem", "r+b" if writable else "rb", buffering=0
            )
        except PermissionError as error:
            raise RuntimeError(
                "cannot read the Wine process; launch it with "
                "original/tools/run_traceable_original.sh"
            ) from error

    def __enter__(self) -> Self:
        return self

    def __exit__(self, *_args: object) -> None:
        self.stream.close()

    def _read_maps(self) -> list[MemoryMap]:
        result = []
        with open(f"/proc/{self.pid}/maps", encoding="ascii") as stream:
            for line in stream:
                fields = line.rstrip().split(maxsplit=5)
                start_text, end_text = fields[0].split("-")
                result.append(
                    MemoryMap(
                        start=int(start_text, 16),
                        end=int(end_text, 16),
                        permissions=fields[1],
                        path=fields[5] if len(fields) == 6 else "",
                    )
                )
        return result

    def read(self, address: int, size: int) -> bytes:
        data = os.pread(self.stream.fileno(), size, address)
        if len(data) != size:
            raise RuntimeError(f"short process read at {address:#x}: {len(data)}/{size}")
        return data

    def write(self, address: int, data: bytes) -> None:
        written = os.pwrite(self.stream.fileno(), data, address)
        if written != len(data):
            raise RuntimeError(f"short process write at {address:#x}: {written}/{len(data)}")

    def search_live_writable_memory(self, needle: bytes) -> list[int]:
        matches = []
        for mapping in self.maps:
            if (
                mapping.start >= 0x68000000
                or "r" not in mapping.permissions
                or "w" not in mapping.permissions
                or mapping.path
            ):
                continue
            try:
                data = self.read(mapping.start, mapping.size)
            except (OSError, RuntimeError):
                continue
            offset = data.find(needle)
            while offset >= 0:
                matches.append(mapping.start + offset)
                offset = data.find(needle, offset + 1)
        return matches


def verify_original() -> None:
    executable = Path(__file__).resolve().parent.parent / "TDKPIN.EXE"
    digest = hashlib.sha256(executable.read_bytes()).hexdigest()
    if digest != ORIGINAL_SHA256:
        raise RuntimeError(f"unexpected TDKPIN.EXE SHA-256: {digest}")


def find_winevdm_pid() -> int:
    matches = []
    for process in Path("/proc").iterdir():
        if not process.name.isdigit():
            continue
        try:
            command = (process / "cmdline").read_bytes().replace(b"\0", b" ").lower()
        except (FileNotFoundError, PermissionError, ProcessLookupError):
            continue
        if b"winevdm.exe" in command and b"tdkpin.exe" in command:
            matches.append(int(process.name))
    if len(matches) != 1:
        raise RuntimeError(f"expected one TDKPIN winevdm.exe process, found {matches}")
    return matches[0]


def unpack_i32(data: bytes, offset: int) -> int:
    return struct.unpack_from("<i", data, offset)[0]


def unpack_u16(data: bytes, offset: int) -> int:
    return struct.unpack_from("<H", data, offset)[0]


def locate_data_segment(memory: ProcessMemory) -> int:
    matches = memory.search_live_writable_memory(DATA_SIGNATURE)
    if len(matches) != 1:
        raise RuntimeError(f"expected one live DGROUP signature, found {[hex(x) for x in matches]}")
    return matches[0] - DATA_SIGNATURE_OFFSET


def locate_game_object(memory: ProcessMemory, data: bytes) -> int:
    current = data[BALL_X : BALL_Y + 4]
    current_x = unpack_i32(data, BALL_X)
    current_y = unpack_i32(data, BALL_Y)
    next_x = unpack_i32(data, NEXT_BALL_X)
    next_y = unpack_i32(data, NEXT_BALL_Y)
    candidates = []

    for address in memory.search_live_writable_memory(current):
        candidate = address - OBJECT_BALL_X
        try:
            obj = memory.read(candidate, OBJECT_READ_SIZE)
        except (OSError, RuntimeError):
            continue
        if obj[OBJECT_BALL_X : OBJECT_BALL_Y + 4] != current:
            continue
        velocity_x = unpack_i32(obj, VELOCITY_X)
        velocity_y = unpack_i32(obj, VELOCITY_Y)
        previous_x = unpack_i32(obj, PREVIOUS_BALL_X)
        previous_y = unpack_i32(obj, PREVIOUS_BALL_Y)
        waiting_at_current_position = (
            velocity_x == 0
            and velocity_y == 0
            and previous_x == current_x
            and previous_y == current_y
        )
        prediction_matches = waiting_at_current_position or (
            next_x - current_x == velocity_x and next_y - current_y == velocity_y
        )
        if prediction_matches:
            candidates.append(candidate)

    unique = sorted(set(candidates))
    if len(unique) != 1:
        raise RuntimeError(f"expected one live game object, found {[hex(x) for x in unique]}")
    return unique[0]


CSV_FIELDS = [
    "sequence",
    "elapsed_us",
    "player_count",
    "current_player",
    "physics_substeps",
    "active_ball_count",
    "maximum_speed_milli_per_tick",
    "gravity_milli_per_substep",
    "ball_x_milli",
    "ball_y_milli",
    "next_ball_x_milli",
    "next_ball_y_milli",
    "velocity_x_milli_per_tick",
    "velocity_y_milli_per_tick",
    "spin_real48_bits",
    "previous_ball_x_milli",
    "previous_ball_y_milli",
    "object_ball_x_milli",
    "object_ball_y_milli",
    "left_flipper_position",
    "right_flipper_position",
    "left_flipper_down",
    "right_flipper_down",
    "ball_active",
    "collision_disabled",
    "ball_suspended",
    "claw_active",
    "claw_frame",
    "claw_target_frame",
    "claw_sprite_bank",
    "claw_animation_blocked",
]


def sample_state(data: bytes, obj: bytes) -> tuple[int, ...]:
    return (
        data[PLAYER_COUNT],
        data[CURRENT_PLAYER],
        unpack_u16(data, PHYSICS_SUBSTEPS),
        obj[ACTIVE_BALL_COUNT],
        unpack_i32(obj, MAXIMUM_SPEED),
        unpack_i32(obj, GRAVITY_PER_SUBSTEP),
        unpack_i32(data, BALL_X),
        unpack_i32(data, BALL_Y),
        unpack_i32(data, NEXT_BALL_X),
        unpack_i32(data, NEXT_BALL_Y),
        unpack_i32(obj, VELOCITY_X),
        unpack_i32(obj, VELOCITY_Y),
        int.from_bytes(obj[SPIN_REAL48 : SPIN_REAL48 + 6], "little"),
        unpack_i32(obj, PREVIOUS_BALL_X),
        unpack_i32(obj, PREVIOUS_BALL_Y),
        unpack_i32(obj, OBJECT_BALL_X),
        unpack_i32(obj, OBJECT_BALL_Y),
        unpack_u16(obj, LEFT_FLIPPER_POSITION),
        unpack_u16(obj, RIGHT_FLIPPER_POSITION),
        obj[LEFT_FLIPPER_DOWN],
        obj[RIGHT_FLIPPER_DOWN],
        obj[BALL_ACTIVE],
        obj[COLLISION_DISABLED],
        obj[BALL_SUSPENDED],
        data[CLAW_ACTIVE],
        unpack_u16(data, CLAW_FRAME),
        unpack_u16(data, CLAW_TARGET_FRAME),
        data[CLAW_SPRITE_BANK],
        data[CLAW_ANIMATION_BLOCKED],
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--pid", type=int, help="Linux PID of the TDKPIN winevdm.exe process")
    parser.add_argument(
        "--object-base",
        type=lambda value: int(value, 0),
        help="known live object base; otherwise locate it from the waiting ball",
    )
    parser.add_argument("--duration", type=float, default=10.0, help="capture duration in seconds")
    parser.add_argument("--interval-ms", type=float, default=1.0, help="sampling interval in milliseconds")
    parser.add_argument("--output", type=Path, help="CSV path; stdout when omitted")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.duration <= 0 or args.interval_ms <= 0:
        raise RuntimeError("duration and interval must be positive")
    verify_original()
    pid = args.pid if args.pid is not None else find_winevdm_pid()
    output = args.output.open("w", newline="", encoding="utf-8") if args.output else sys.stdout

    try:
        writer = csv.writer(output)
        writer.writerow(CSV_FIELDS)
        with ProcessMemory(pid) as memory:
            data_base = locate_data_segment(memory)
            initial_data = memory.read(data_base, DATA_READ_SIZE)
            object_base = (
                args.object_base
                if args.object_base is not None
                else locate_game_object(memory, initial_data)
            )
            print(
                f"TDKPIN pid={pid} data={data_base:#x} object={object_base:#x}",
                file=sys.stderr,
            )

            started_ns = time.monotonic_ns()
            deadline_ns = started_ns + int(args.duration * 1_000_000_000)
            interval_seconds = args.interval_ms / 1000.0
            last_state = None
            sequence = 0
            while time.monotonic_ns() < deadline_ns:
                data = memory.read(data_base, DATA_READ_SIZE)
                obj = memory.read(object_base, OBJECT_READ_SIZE)
                state = sample_state(data, obj)
                if state != last_state:
                    sequence += 1
                    elapsed_us = (time.monotonic_ns() - started_ns) // 1000
                    writer.writerow((sequence, elapsed_us, *state))
                    output.flush()
                    last_state = state
                time.sleep(interval_seconds)
    finally:
        if output is not sys.stdout:
            output.close()


if __name__ == "__main__":
    try:
        main()
    except (OSError, RuntimeError) as error:
        raise SystemExit(f"trace_original_state.py: {error}") from error