"""SWD/DAP operations — Debug Port and Access Port register access."""

import logging

from fastmcp import Context
from fastmcp.contrib.mcp_mixin import MCPMixin, mcp_tool
from openocd.errors import OpenOCDError

from mcjtag.mixins.helpers import parse_hex, require_session
from mcjtag.models import APInfoResult, APListResult, APRegResult, DAPInfoResult, DPRegResult
from mcjtag.state import SessionState

log = logging.getLogger(__name__)


def _get_state(ctx: Context) -> SessionState:
    return ctx.request_context.lifespan_context


class SWDMixin(MCPMixin):
    """Tools for SWD/DAP operations: DAP identification and DP/AP register access.

    The ARM Debug Access Port (DAP) has two layers:
    - Debug Port (DP): Controls the debug interface itself (always present).
    - Access Ports (APs): Provide access to system resources. AP 0 is usually
      a MEM-AP for memory bus access. Other APs may provide JTAG-AP or
      additional MEM-APs for different bus domains.

    Typical workflow: swd_info -> swd_list_aps -> swd_read_ap/swd_write_ap.
    For standard memory and register access, prefer the higher-level
    read_memory/write_memory/read_registers tools instead.
    """

    @mcp_tool()
    async def swd_info(self, ctx: Context, dap: str | None = None) -> DAPInfoResult:
        """Get Debug Access Port identification -- DPIDR, AP count, and raw info.

        Identifies the DAP hardware revision and shows how many Access Ports
        are present. The DPIDR value encodes the DP version (DPv0/DPv1/DPv2),
        designer, and part number. Use this as the first SWD operation to
        verify the debug connection is working.

        On single-DAP boards (most Cortex-M), the dap parameter can be
        omitted. Multi-DAP systems (e.g. multi-core SoCs) require specifying
        which DAP to query.

        Args:
            dap: DAP instance name (e.g. "stm32f1x.dap"). Omit for auto-detection
                 on single-DAP targets.
        """
        state = _get_state(ctx)
        require_session(state)

        info = await state.session.swd.info(dap=dap)
        return DAPInfoResult(
            name=info.name,
            dpidr=f"0x{info.dpidr:08x}",
            ap_count=info.ap_count,
            raw_info=info.raw_info,
        )

    @mcp_tool()
    async def swd_list_aps(self, ctx: Context, dap: str | None = None) -> APListResult:
        """Enumerate all Access Ports on the DAP.

        Probes each AP index (0, 1, 2, ...) by reading its IDR register
        (offset 0xFC) until a zero IDR is found. Returns the AP type
        (MEM-AP, JTAG-AP, etc.), IDR value, and ROM table base address.

        Common AP types:
        - MEM-AP (AHB/APB/AXI): Provides memory bus access. AP 0 is
          usually AHB-AP for Cortex-M system bus access.
        - JTAG-AP: Bridges to a JTAG chain behind the DAP.

        Use the returned AP index numbers with swd_read_ap/swd_write_ap
        to access specific AP registers.

        Args:
            dap: DAP instance name. Omit for auto-detection on single-DAP targets.
        """
        state = _get_state(ctx)
        require_session(state)

        aps = await state.session.swd.list_aps(dap=dap)
        return APListResult(
            aps=[
                APInfoResult(
                    index=ap.index,
                    idr=f"0x{ap.idr:08x}",
                    base=f"0x{ap.base:08x}",
                    ap_type=ap.ap_type,
                )
                for ap in aps
            ],
            count=len(aps),
        )

    @mcp_tool()
    async def swd_read_dp(self, ctx: Context, address: str, dap: str | None = None) -> DPRegResult:
        """Read a Debug Port register.

        DP registers control the debug interface itself. These are the
        lowest-level SWD registers, below the Access Port layer.

        DP register map:
        - 0x0  DPIDR: DP Identification. Read-only. Contains DP version,
          designer code, and part number.
        - 0x4  CTRL/STAT: Control and status bits. Includes power-up
          request/acknowledge, overrun detection, and sticky error flags.
          Reading shows status; writing controls power and error clearing.
        - 0x8  SELECT: AP and register bank select. Bits [31:24] select
          the AP number, bits [7:4] select the register bank within that AP.
          Normally managed automatically by higher-level AP operations.
        - 0xC  RDBUFF: Read buffer. Returns the result of the last AP read.
          Used internally by the SWD protocol; rarely needed directly.
        - 0x24 TARGETID: Target identification (DPv2+ only). Contains
          the JEDEC manufacturer ID and part number of the target device.

        For most debugging tasks, use swd_info instead of reading DP
        registers directly.

        Args:
            address: DP register offset in hex (e.g. "0x0", "0x4", "0x24").
            dap: DAP instance name. Omit for auto-detection on single-DAP targets.
        """
        state = _get_state(ctx)
        require_session(state)

        addr = parse_hex(address, "DP register address")
        value = await state.session.swd.dpreg(addr, dap=dap)
        return DPRegResult(
            address=f"0x{addr:x}",
            value=f"0x{value:08x}",
            written=False,
        )

    @mcp_tool()
    async def swd_write_dp(
        self, ctx: Context, address: str, value: str, dap: str | None = None
    ) -> DPRegResult:
        """Write a Debug Port register.

        CAUTION: DP registers control the debug interface itself. Writing
        incorrect values to CTRL/STAT (0x4) or SELECT (0x8) can lock out
        the debugger, requiring a power cycle to recover.

        Common write operations:
        - CTRL/STAT (0x4): Write "0x50000000" to request system and debug
          power-up. Write "0x00000010" to clear sticky error flags.
        - SELECT (0x8): Write AP number in bits [31:24] and bank in [7:4]
          to select which AP and register bank subsequent AP reads access.

        DPIDR (0x0) and RDBUFF (0xC) are read-only; writes are ignored.

        Args:
            address: DP register offset in hex (e.g. "0x4" for CTRL/STAT).
            value: Value to write in hex (e.g. "0x50000000").
            dap: DAP instance name. Omit for auto-detection on single-DAP targets.
        """
        state = _get_state(ctx)
        require_session(state)

        addr = parse_hex(address, "DP register address")
        val = parse_hex(value, "DP register value")
        await state.session.swd.dpreg(addr, val, dap=dap)
        return DPRegResult(
            address=f"0x{addr:x}",
            value=f"0x{val:08x}",
            written=True,
        )

    @mcp_tool()
    async def swd_read_ap(
        self, ctx: Context, ap: int, address: str, dap: str | None = None
    ) -> APRegResult:
        """Read an Access Port register by AP index and register offset.

        This is a low-level DAP operation. For reading target memory, prefer
        read_memory which handles CSW/TAR/DRW sequencing automatically.

        MEM-AP register map (offsets are multiples of 0x04):
        - 0x00 CSW: Control/Status Word. Bits [2:0] set access size
          (0=byte, 1=halfword, 2=word). Bits [5:4] set address increment
          mode (1=increment, 2=packed).
        - 0x04 TAR: Transfer Address Register. Set this to the target
          memory address before reading/writing DRW.
        - 0x0C DRW: Data Read/Write register. Reading DRW fetches data
          from the address in TAR. Writing DRW stores data at TAR.
        - 0x10 BD0-0x1C BD3: Banked Data registers (direct word access
          to TAR+0x00 through TAR+0x0C).
        - 0xF8 BASE: Debug ROM table base address. Points to CoreSight
          component discovery tables.
        - 0xFC IDR: AP Identification Register. Encodes AP type, class,
          designer, and revision. Read this to identify what kind of AP
          you are talking to.

        To enumerate APs, use swd_list_aps instead of manually reading
        IDR registers.

        Args:
            ap: Access Port index (0-255). Get valid indices from swd_list_aps.
            address: AP register offset in hex. Must be 4-byte aligned.
                     Common values: "0x00" (CSW), "0x04" (TAR), "0x0C" (DRW),
                     "0xF8" (BASE), "0xFC" (IDR).
            dap: DAP instance name. Omit for auto-detection on single-DAP targets.
        """
        state = _get_state(ctx)
        require_session(state)

        if ap < 0 or ap > 255:
            raise OpenOCDError(f"AP number must be 0-255, got {ap}")
        addr = parse_hex(address, "AP register address")
        value = await state.session.swd.apreg(ap, addr, dap=dap)
        return APRegResult(
            ap=ap,
            address=f"0x{addr:x}",
            value=f"0x{value:08x}",
            written=False,
        )

    @mcp_tool()
    async def swd_write_ap(
        self, ctx: Context, ap: int, address: str, value: str, dap: str | None = None
    ) -> APRegResult:
        """Write an Access Port register by AP index and register offset.

        CAUTION: AP registers directly control memory bus access. Writing
        TAR (0x04) sets the target address, and writing DRW (0x0C) performs
        a bus write to that address. Incorrect CSW (0x00) values can cause
        bus faults or lock up the debug interface.

        This bypasses the safe-write-range checks that write_memory enforces.
        For writing target memory with safety validation, prefer write_memory.

        Common write sequences for manual MEM-AP memory access:
          1. Write CSW (0x00) with "0x23000002" (32-bit, auto-increment)
          2. Write TAR (0x04) with the target address
          3. Write DRW (0x0C) with the data value

        Args:
            ap: Access Port index (0-255). Get valid indices from swd_list_aps.
            address: AP register offset in hex. Must be 4-byte aligned.
                     Common values: "0x00" (CSW), "0x04" (TAR), "0x0C" (DRW).
            value: Value to write in hex (e.g. "0x23000002" for CSW word-size
                   auto-increment mode).
            dap: DAP instance name. Omit for auto-detection on single-DAP targets.
        """
        state = _get_state(ctx)
        require_session(state)

        if ap < 0 or ap > 255:
            raise OpenOCDError(f"AP number must be 0-255, got {ap}")
        addr = parse_hex(address, "AP register address")
        val = parse_hex(value, "AP register value")
        await state.session.swd.apreg(ap, addr, val, dap=dap)
        return APRegResult(
            ap=ap,
            address=f"0x{addr:x}",
            value=f"0x{val:08x}",
            written=True,
        )
