"""Pure-Python GDB Remote Serial Protocol (RSP) client for QEMU's gdbstub.

risc-v-gsp-rm-debugger deliberately avoids GDB: the whole instrumentation stack is a Python
library, so sessions, probes and plugins are scriptable end to end.  This
module implements the RSP subset QEMU needs: packet framing/acks, no-ack mode,
registers (with layout read from ``target.xml``), memory, software/hardware
breakpoints, watchpoints, continue/step/interrupt and monitor commands.

Tested against QEMU 8.2 (``qemu-system-riscv64``).
"""
from __future__ import annotations

import socket
import time
from dataclasses import dataclass, field as _field

from .errors import RspProtocolError, RspTimeout

# ---------------------------------------------------------------------------
# Packet framing
# ---------------------------------------------------------------------------


def checksum(data: bytes) -> int:
    """RSP 8-bit checksum: sum of bytes mod 256."""
    return sum(data) & 0xFF


def encode_packet(data: bytes) -> bytes:
    """Frame a payload as ``$<data>#<checksum>``."""
    return b"$" + data + b"#" + f"{checksum(data):02x}".encode("ascii")


def _hex_value(payload: str) -> int:
    try:
        return int(payload, 16)
    except ValueError as exc:  # pragma: no cover - defensive
        raise RspProtocolError(f"malformed hex value: {payload!r}") from exc


# ---------------------------------------------------------------------------
# Stop replies
# ---------------------------------------------------------------------------


@dataclass
class StopReply:
    """A parsed stop reply (``T``/``S``/``W``/``X`` packet)."""

    raw: str
    signal: int | None = None
    thread: str | None = None
    regs: dict[int, int] = _field(default_factory=dict)
    watch_addr: int | None = None
    watch_kind: str | None = None
    exit_status: int | None = None

    @property
    def is_exit(self) -> bool:
        return self.raw[:1] in ("W", "X")

    @property
    def is_stop(self) -> bool:
        return self.raw[:1] in ("T", "S")

    def __str__(self) -> str:  # pragma: no cover - cosmetic
        if self.is_exit:
            return f"exit(status={self.exit_status})"
        bits = [f"signal={self.signal}"]
        if self.watch_addr is not None:
            bits.append(f"{self.watch_kind}@{self.watch_addr:#x}")
        if self.thread:
            bits.append(f"thread={self.thread}")
        return "stop(" + ", ".join(bits) + ")"


def parse_stop_reply(raw: str) -> StopReply:
    """Parse a stop packet body (without the leading channel)."""
    reply = StopReply(raw=raw)
    if not raw:
        return reply
    head = raw[0]
    if head in ("T", "S"):
        if len(raw) >= 3:
            try:
                reply.signal = int(raw[1:3], 16)
            except ValueError:
                pass
        rest = raw[3:]
        for item in rest.split(";"):
            if not item or ":" not in item:
                continue
            key, value = item.split(":", 1)
            if key == "thread":
                reply.thread = value
            elif key in ("watch", "rwatch", "awatch"):
                reply.watch_kind = key
                reply.watch_addr = _hex_value(value)
            else:
                try:
                    reply.regs[int(key, 16)] = _hex_value(value)
                except ValueError:
                    continue
    elif head in ("W", "X"):
        if len(raw) >= 3:
            try:
                reply.exit_status = int(raw[1:3], 16)
            except ValueError:
                pass
    return reply


# ---------------------------------------------------------------------------
# Register layout
# ---------------------------------------------------------------------------


@dataclass
class Register:
    number: int
    name: str
    bitsize: int
    offset: int = 0

    @property
    def nbytes(self) -> int:
        return (self.bitsize + 7) // 8


#: Fallback layout for RV64 when target.xml is unavailable: x0..x31 + pc.
RV64_FALLBACK_REGS = [
    Register(n, name, 64, n * 8)
    for n, name in enumerate(
        ["x0", "ra", "sp", "gp", "tp", "t0", "t1", "t2"]
        + [f"s{i}" for i in range(0, 2)]
        + [f"a{i}" for i in range(0, 8)]
        + [f"s{i}" for i in range(2, 12)]
        + [f"t{i}" for i in range(3, 7)]
        + ["pc"]
    )
]


# ---------------------------------------------------------------------------
# Client
# ---------------------------------------------------------------------------


class RspClient:
    """A blocking RSP client over TCP."""

    def __init__(self, host: str = "127.0.0.1", port: int = 4444, timeout: float = 10.0):
        self.host = host
        self.port = port
        self.timeout = timeout
        self.sock: socket.socket | None = None
        self.no_ack = False
        self.packet_size = 0x1000
        self.registers: list[Register] = []
        self._by_name: dict[str, Register] = {}
        self._by_number: dict[int, Register] = {}
        self.features: str = ""

    # -- lifecycle ---------------------------------------------------------

    def connect(self, retries: int = 1, retry_delay: float = 0.1) -> None:
        last: Exception | None = None
        for _ in range(max(1, retries)):
            try:
                self.sock = socket.create_connection((self.host, self.port), timeout=self.timeout)
                self.sock.settimeout(self.timeout)
                return
            except OSError as exc:
                last = exc
                time.sleep(retry_delay)
        raise RspTimeout(f"cannot connect to gdbstub {self.host}:{self.port}: {last}")

    def close(self) -> None:
        if self.sock is not None:
            try:
                self.sock.close()
            except OSError:
                pass
            self.sock = None

    @property
    def connected(self) -> bool:
        return self.sock is not None

    def __enter__(self) -> "RspClient":
        if not self.connected:
            self.connect()
        return self

    def __exit__(self, *exc) -> None:
        self.close()

    # -- low level I/O -----------------------------------------------------

    def _send_raw(self, data: bytes) -> None:
        if self.sock is None:
            raise RspTimeout("not connected")
        self.sock.sendall(data)

    def _read_exact(self, count: int, deadline: float) -> bytes:
        if self.sock is None:
            raise RspTimeout("not connected")
        chunks = bytearray()
        while len(chunks) < count:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                raise RspTimeout(f"timeout waiting for {count} bytes")
            self.sock.settimeout(remaining)
            try:
                part = self.sock.recv(count - len(chunks))
            except socket.timeout as exc:
                raise RspTimeout(f"timeout waiting for {count} bytes") from exc
            if not part:
                raise RspProtocolError("gdbstub closed the connection")
            chunks += part
        return bytes(chunks)

    def _send_packet(self, payload: bytes) -> None:
        packet = encode_packet(payload)
        self._send_raw(packet)
        if self.no_ack:
            return
        deadline = time.monotonic() + self.timeout
        while True:
            ack = self._read_exact(1, deadline)
            if ack == b"+":
                return
            if ack == b"-":
                self._send_raw(packet)
                continue
            # Some stubs answer with a stop packet instead of an ack when a
            # run-control command stops immediately; push it back for the
            # packet reader by buffering is complex, so treat as protocol
            # error only if it is not a start-of-packet marker.
            if ack == b"$":
                raise RspProtocolError("unexpected packet while waiting for ack")

    def _read_packet(self, timeout: float | None = None) -> str:
        eff = self.timeout if timeout is None else timeout
        deadline = time.monotonic() + eff
        # Skip stale acks and interrupt notifications until a packet starts.
        while True:
            byte = self._read_exact(1, deadline)
            if byte == b"$":
                break
            if byte == b"+":
                continue
            if byte == b"\x03":
                return "S02"  # async interrupt notification
        body = bytearray()
        while True:
            byte = self._read_exact(1, deadline)
            if byte == b"#":
                break
            body += byte
        got = self._read_exact(2, deadline)
        try:
            expected = int(got.decode("ascii"), 16)
        except ValueError as exc:
            raise RspProtocolError(f"malformed checksum bytes {got!r}") from exc
        if checksum(bytes(body)) != expected:
            if not self.no_ack:
                self._send_raw(b"-")
            raise RspProtocolError(f"bad checksum for {bytes(body)!r}")
        if not self.no_ack:
            self._send_raw(b"+")
        return bytes(body).decode("latin1")

    def command(self, payload: str | bytes, timeout: float | None = None) -> str:
        """Send one packet and return the reply payload."""
        data = payload.encode("ascii") if isinstance(payload, str) else payload
        self._send_packet(data)
        return self._read_packet(timeout)

    # -- handshake ---------------------------------------------------------

    def sync(self) -> None:
        """Negotiate features and register layout; consume the initial stop."""
        # QEMU acknowledges QStartNoAckMode with an EMPTY packet ($#00), then
        # switches to no-ack mode once the client acks that reply.
        try:
            self._send_packet(b"QStartNoAckMode")
            reply = self._read_packet()
            if reply in ("", "OK"):
                self.no_ack = True
        except RspProtocolError:
            pass  # older stub: stay in ack mode

        self.packet_size = 0x1000
        try:
            self.features = self.command("qSupported:multiprocess+;swbreak+;hwbreak+;qXfer:features:read+")
            for token in self.features.split(";"):
                if token.startswith("PacketSize="):
                    try:
                        self.packet_size = int(token.split("=", 1)[1], 16)
                    except ValueError:
                        pass
        except RspProtocolError:
            self.features = ""

        # Consume the initial stop reply (target is paused).
        try:
            self.command("?")
        except RspProtocolError:
            pass
        # Select CPU 0 for register access in system mode.
        for selection in ("Hg0", "Hc-1"):
            try:
                self.command(selection)
            except RspProtocolError:
                pass
        self._load_registers()

    def _load_registers(self) -> None:
        self.registers = []
        try:
            xml = self._read_target_xml()
        except RspProtocolError:
            xml = ""
        if xml:
            self.registers = _parse_target_xml(xml)
        if not self.registers:
            self.registers = list(RV64_FALLBACK_REGS)
        self._by_name = {r.name: r for r in self.registers}
        self._by_number = {r.number: r for r in self.registers}

    def _read_target_xml(self, depth: int = 0) -> str:
        """Fetch target.xml and recursively inline its xi:include files."""
        return self._collect_annex("target.xml", depth)

    def _collect_annex(self, annex: str, depth: int) -> str:
        import re

        try:
            xml = self._read_xfer("features", annex)
        except RspProtocolError:
            return ""
        if depth >= 4:
            return xml
        parts = [xml]
        for href in re.findall(r'xi:include[^>]*href="([^"]+)"', xml):
            if href.startswith(("http:", "https:", "file:")):
                continue
            child = self._collect_annex(href, depth + 1)
            if child:
                parts.append(child)
        return "\n".join(parts)

    def _read_xfer(self, kind: str, annex: str) -> str:
        data = bytearray()
        offset = 0
        while True:
            reply = self.command(f"qXfer:{kind}:read:{annex}:{offset:x},fff")
            if reply.startswith("E"):
                raise RspProtocolError(f"qXfer {kind}:{annex} failed: {reply}")
            if not reply or reply[0] not in ("l", "m"):
                raise RspProtocolError(f"qXfer {kind}:{annex} malformed: {reply[:16]!r}")
            chunk = reply[1:].encode("latin1")
            data += chunk
            offset += len(chunk)
            if reply[0] == "l":
                break
        return data.decode("latin1", "replace")

    # -- registers ---------------------------------------------------------

    def read_registers(self) -> dict[str, int]:
        payload = self.command("g")
        raw = _unhex(payload)
        out: dict[str, int] = {}
        for reg in self.registers:
            start = reg.offset
            chunk = raw[start : start + reg.nbytes]
            if len(chunk) < reg.nbytes:
                break
            out[reg.name] = int.from_bytes(chunk, "little", signed=False)
        return out

    def read_reg(self, name_or_number: str | int) -> int:
        reg = self._reg(name_or_number)
        reply = self.command(f"p{reg.number:x}")
        if reply.startswith("E"):
            raise RspProtocolError(f"read register {reg.name} failed: {reply}")
        return int.from_bytes(_unhex(reply), "little", signed=False)

    def write_reg(self, name_or_number: str | int, value: int) -> None:
        reg = self._reg(name_or_number)
        payload = value.to_bytes(reg.nbytes, "little", signed=False).hex()
        reply = self.command(f"P{reg.number:x}={payload}")
        if reply != "OK":
            raise RspProtocolError(f"write register {reg.name} failed: {reply}")

    def _reg(self, name_or_number: str | int) -> Register:
        if isinstance(name_or_number, int):
            if name_or_number in self._by_number:
                return self._by_number[name_or_number]
            # allow ABI-numbered convenience: 0..31 map to x0..x31
            if name_or_number < len(self.registers):
                return self.registers[name_or_number]
            raise RspProtocolError(f"unknown register number {name_or_number}")
        if name_or_number in self._by_name:
            return self._by_name[name_or_number]
        aliases = {"pc": "pc", "a0": "a0", "return": "ra"}
        name_or_number = aliases.get(name_or_number, name_or_number)
        if name_or_number in self._by_name:
            return self._by_name[name_or_number]
        raise RspProtocolError(f"unknown register {name_or_number!r}")

    @property
    def pc(self) -> int:
        return self.read_reg("pc")

    # -- memory ------------------------------------------------------------

    def _mem_chunk(self) -> int:
        # Each data byte costs two hex chars; leave room for the header.
        return max(64, (self.packet_size - 128) // 2)

    def read_mem(self, addr: int, size: int, chunk: int = 0) -> bytes:
        if chunk <= 0:
            chunk = self._mem_chunk()
        out = bytearray()
        while len(out) < size:
            want = min(chunk, size - len(out))
            here = addr + len(out)
            reply = self.command(f"m{here:x},{want:x}")
            if reply.startswith("E"):
                raise RspProtocolError(f"memory read {here:#x}+{want:#x} failed: {reply}")
            data = _unhex(reply)
            if not data:
                raise RspProtocolError(f"short memory read at {here:#x}")
            out += data
        return bytes(out)

    def write_mem(self, addr: int, data: bytes, chunk: int = 0) -> None:
        if chunk <= 0:
            chunk = self._mem_chunk()
        view = memoryview(data)
        offset = 0
        while offset < len(view):
            part = bytes(view[offset : offset + chunk])
            here = addr + offset
            reply = self.command(f"M{here:x},{len(part):x}:{part.hex()}")
            if reply != "OK":
                raise RspProtocolError(f"memory write {here:#x} failed: {reply}")
            offset += len(part)

    def read_u32(self, addr: int) -> int:
        return int.from_bytes(self.read_mem(addr, 4), "little")

    def read_u64(self, addr: int) -> int:
        return int.from_bytes(self.read_mem(addr, 8), "little")

    def write_u32(self, addr: int, value: int) -> None:
        self.write_mem(addr, value.to_bytes(4, "little"))

    def write_u64(self, addr: int, value: int) -> None:
        self.write_mem(addr, value.to_bytes(8, "little"))

    def read_cstring(self, addr: int, max_len: int = 256) -> str:
        raw = self.read_mem(addr, max_len)
        return raw.split(b"\x00", 1)[0].decode("utf-8", "replace")

    # -- breakpoints / watchpoints ----------------------------------------

    def set_sw_break(self, addr: int, kind: int = 4) -> None:
        reply = self.command(f"Z0,{addr:x},{kind:x}")
        if reply != "OK":
            raise RspProtocolError(f"set software breakpoint {addr:#x} failed: {reply}")

    def clear_sw_break(self, addr: int, kind: int = 4) -> None:
        reply = self.command(f"z0,{addr:x},{kind:x}")
        if reply != "OK":
            raise RspProtocolError(f"clear software breakpoint {addr:#x} failed: {reply}")

    def set_hw_break(self, addr: int, kind: int = 4) -> None:
        reply = self.command(f"Z1,{addr:x},{kind:x}")
        if reply != "OK":
            raise RspProtocolError(f"set hardware breakpoint {addr:#x} failed: {reply}")

    def clear_hw_break(self, addr: int, kind: int = 4) -> None:
        reply = self.command(f"z1,{addr:x},{kind:x}")
        if reply != "OK":
            raise RspProtocolError(f"clear hardware breakpoint {addr:#x} failed: {reply}")

    _WATCH_KINDS = {"w": 2, "r": 3, "a": 4}

    def set_watch(self, addr: int, size: int = 8, kind: str = "w") -> None:
        code = self._WATCH_KINDS[kind]
        reply = self.command(f"Z{code},{addr:x},{size:x}")
        if reply != "OK":
            raise RspProtocolError(f"set watchpoint {kind}@{addr:#x} failed: {reply}")

    def clear_watch(self, addr: int, size: int = 8, kind: str = "w") -> None:
        code = self._WATCH_KINDS[kind]
        reply = self.command(f"z{code},{addr:x},{size:x}")
        if reply != "OK":
            raise RspProtocolError(f"clear watchpoint {kind}@{addr:#x} failed: {reply}")

    # -- run control -------------------------------------------------------

    def continue_(self, timeout: float | None = None) -> StopReply:
        self._send_packet(b"c")
        raw = self._read_packet(timeout)
        return parse_stop_reply(raw)

    def step(self, timeout: float | None = None) -> StopReply:
        self._send_packet(b"s")
        raw = self._read_packet(timeout)
        return parse_stop_reply(raw)

    def interrupt(self, timeout: float | None = None) -> StopReply:
        self._send_raw(b"\x03")
        raw = self._read_packet(timeout)
        return parse_stop_reply(raw)

    # -- monitor -----------------------------------------------------------

    def monitor(self, cmd: str) -> str:
        payload = ("qRcmd," + cmd.encode("ascii").hex()).encode("ascii")
        self._send_packet(payload)
        out = bytearray()
        while True:
            reply = self._read_packet()
            if reply == "OK":
                break
            if reply.startswith("E"):
                raise RspProtocolError(f"monitor {cmd!r} failed: {reply}")
            if reply.startswith("O"):
                out += bytes.fromhex(reply[1:])
            else:  # pragma: no cover - defensive
                raise RspProtocolError(f"unexpected monitor reply: {reply[:32]!r}")
        return out.decode("utf-8", "replace")


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _unhex(payload: str) -> bytes:
    if len(payload) % 2:
        payload = payload[:-1]  # tolerate odd trailing nibble
    try:
        return bytes.fromhex(payload)
    except ValueError as exc:
        raise RspProtocolError(f"malformed hex payload: {payload[:32]!r}") from exc


_RE_REG = None


def _parse_target_xml(xml: str) -> list[Register]:
    import re

    global _RE_REG
    if _RE_REG is None:
        _RE_REG = re.compile(r"<reg\b[^>]*?/?>", re.IGNORECASE)
    attrs_re = re.compile(r'(\w+)\s*=\s*"([^"]*)"')
    regs: list[Register] = []
    next_number = 0
    offset = 0
    for match in _RE_REG.finditer(xml):
        attrs = {k: v for k, v in attrs_re.findall(match.group(0))}
        name = attrs.get("name")
        if not name:
            continue
        bitsize = int(attrs.get("bitsize", "64"))
        number = int(attrs["regnum"], 0) if "regnum" in attrs else next_number
        if "offset" in attrs:
            try:
                offset = int(attrs["offset"], 0)
            except ValueError:
                pass
        regs.append(Register(number=number, name=name, bitsize=bitsize, offset=offset))
        offset += (bitsize + 7) // 8
        next_number = number + 1
    return regs