"""Hardware-only origin trace for the exact SetTeam team context.

Matches the typed integer context pointer selected by SetTeam to the constructor
invocation that produced it. No client memory writes.
"""
from __future__ import annotations

from collections import deque
import json
import os
import struct
import time
import traceback

import gdb

CONTEXT_REUSE = 0x1477C17FC
CONTEXT_ALLOCATED = 0x1477C18C1
SET_TEAM_STUB = 0x147060A80
LOCKED_SETTER_RETURN = 0x1477C2415
CONTEXT_STACK_COUNT = 0x144BCEDA0
CONTEXT_STACK_ARRAY = 0x144BCEDA8
INTERESTING = {73, 130000, 130001}
_STATE = None


def _reg(name):
    return int(gdb.parse_and_eval(f"${name}"))


def _read(address, size):
    if not address or address < 0:
        return None
    try:
        return bytes(gdb.selected_inferior().read_memory(address, size))
    except gdb.error:
        return None


def _u64(address):
    data = _read(address, 8)
    return struct.unpack("<Q", data)[0] if data else None


def _i32(address):
    data = _read(address, 4)
    return struct.unpack("<i", data)[0] if data else None


def _thread():
    thread = gdb.selected_thread()
    if thread is None:
        return {}
    return {"name": thread.name, "ptid": list(thread.ptid), "global_num": thread.global_num}


class State:
    def __init__(self, path):
        self.path = path
        self.index = 0
        self.total_constructor_hits = 0
        self.interesting_constructor_hits = 0
        self.pending_allocations = {}
        self.origins = deque(maxlen=4096)

    def log(self, kind, **payload):
        self.index += 1
        event = {
            "event": kind,
            "event_index": self.index,
            "time_unix": time.time(),
            "thread": _thread(),
            **payload,
        }
        with open(self.path, "a", encoding="utf-8") as handle:
            handle.write(json.dumps(event, sort_keys=True) + "\n")
            handle.flush()
            os.fsync(handle.fileno())

    def thread_key(self):
        return tuple(_thread().get("ptid", ()))

    def remember_origin(self, context, origin):
        if context:
            self.origins.append({**origin, "context": context})

    def find_origin(self, context):
        return next((origin for origin in reversed(self.origins)
                     if origin["context"] == context), None)


class HardwareBreakpoint(gdb.Breakpoint):
    def __init__(self, state, address):
        self.state = state
        self.address = address
        super().__init__(f"*0x{address:x}", type=gdb.BP_HARDWARE_BREAKPOINT, internal=False)
        self.silent = True


class ContextReuseBreakpoint(HardwareBreakpoint):
    def stop(self):
        self.state.total_constructor_hits += 1
        try:
            value = _reg("rcx") & 0xFFFFFFFF
            if value not in INTERESTING:
                return False
            self.state.interesting_constructor_hits += 1
            rsp = _reg("rsp")
            direct_return = _u64(rsp + 0x28)
            origin = {
                "value": value,
                "direct_return_address": direct_return,
                "upstream_return_address": (
                    _u64(rsp + 0x68)
                    if direct_return == LOCKED_SETTER_RETURN
                    else direct_return
                ),
                "constructor_stack_hex": (_read(rsp, 0x100) or b"").hex(),
                "constructor_hit": self.state.total_constructor_hits,
            }
            context = _reg("rax")
            if context:
                self.state.remember_origin(context, origin)
            else:
                self.state.pending_allocations[self.state.thread_key()] = origin
        except Exception as exc:
            self.state.log(
                "trace_error",
                where="context_reuse",
                error=str(exc),
                traceback=traceback.format_exc(),
            )
        return False


class ContextAllocatedBreakpoint(HardwareBreakpoint):
    def stop(self):
        try:
            origin = self.state.pending_allocations.pop(self.state.thread_key(), None)
            if origin is not None:
                self.state.remember_origin(_reg("rdx"), origin)
        except Exception as exc:
            self.state.log(
                "trace_error",
                where="context_allocated",
                error=str(exc),
                traceback=traceback.format_exc(),
            )
        return False


class SetTeamStubBreakpoint(HardwareBreakpoint):
    def stop(self):
        try:
            count = _i32(CONTEXT_STACK_COUNT)
            array = _u64(CONTEXT_STACK_ARRAY)
            team_context = (
                _u64(array + (count - 2) * 8)
                if array and count is not None and count >= 2
                else None
            )
            side_context = (
                _u64(array + (count - 1) * 8)
                if array and count is not None and count >= 1
                else None
            )
            rsp = _reg("rsp")
            self.state.log(
                "set_team_stub_entry",
                context_stack_count=count,
                team_context=team_context,
                team_context_hex=(_read(team_context, 0x40) or b"").hex(),
                team_value=_i32(team_context + 0x10) if team_context else None,
                side_context=side_context,
                side_value=_i32(side_context + 0x10) if side_context else None,
                matched_origin=self.state.find_origin(team_context),
                caller_return_address=_u64(rsp),
                entry_registers={
                    name: _reg(name)
                    for name in ("rcx", "rdx", "r8", "r9")
                },
                backtrace=gdb.execute("bt 32", to_string=True),
                total_constructor_hits=self.state.total_constructor_hits,
                interesting_constructor_hits=self.state.interesting_constructor_hits,
            )
        except Exception as exc:
            self.state.log(
                "trace_error",
                where="set_team_stub",
                error=str(exc),
                traceback=traceback.format_exc(),
            )
        return False


def _on_exit(event):
    if _STATE is not None:
        _STATE.log(
            "inferior_exited",
            detail=str(event),
            total_constructor_hits=_STATE.total_constructor_hits,
            interesting_constructor_hits=_STATE.interesting_constructor_hits,
        )


def start_trace(log_path, _cards_base):
    global _STATE
    open(log_path, "w", encoding="utf-8").close()
    _STATE = State(log_path)
    points = {
        "context_reuse": ContextReuseBreakpoint(_STATE, CONTEXT_REUSE),
        "context_allocated": ContextAllocatedBreakpoint(_STATE, CONTEXT_ALLOCATED),
        "set_team_stub": SetTeamStubBreakpoint(_STATE, SET_TEAM_STUB),
    }
    gdb.events.exited.connect(_on_exit)
    _STATE.log(
        "trace_armed",
        breakpoints={
            name: {"number": point.number, "va": point.address}
            for name, point in points.items()
        },
        hardware_only=True,
        client_memory_writes=False,
        matching="exact_context_pointer",
    )