"""JTAG chain scanning and raw shift operations."""

import logging
from typing import Literal

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 JTAGChain, JTAGShiftResult, TAPResult
from mcjtag.state import SessionState

log = logging.getLogger(__name__)


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


class JTAGMixin(MCPMixin):
    """Tools for JTAG chain enumeration and low-level scan operations."""

    @mcp_tool()
    async def jtag_scan(self, ctx: Context) -> JTAGChain:
        """Scan the JTAG chain and enumerate all TAPs (Test Access Ports).

        Returns each TAP's name, IDCODE, instruction register length (IR),
        and enabled status. Use this to discover connected devices, verify
        the JTAG chain integrity, and obtain TAP names needed by jtag_shift.

        The TAP name (e.g. "stm32f1x.cpu") is required for all subsequent
        JTAG operations. Run this first after connecting via JTAG transport.
        """
        state = _get_state(ctx)
        require_session(state)

        taps = await state.session.jtag.scan_chain()
        tap_results = [
            TAPResult(
                name=tap.name,
                idcode=f"0x{tap.idcode:08x}",
                ir_length=tap.ir_length,
                enabled=tap.enabled,
            )
            for tap in taps
        ]

        return JTAGChain(taps=tap_results, count=len(tap_results))

    @mcp_tool()
    async def jtag_shift(
        self,
        ctx: Context,
        tap: str,
        operation: Literal["irscan", "drscan"],
        value: str,
        bits: int | None = None,
    ) -> JTAGShiftResult:
        """Perform a raw JTAG IR scan or DR scan operation.

        Low-level JTAG shift register access. Use this for direct TAP
        interaction when higher-level tools (memory, registers, SWD) do
        not cover your protocol. Prefer target/memory/register tools for
        standard Cortex-M debugging.

        Operations:
        - irscan: Loads an instruction into the TAP's Instruction Register.
          Selects which Data Register subsequent drscan operations access.
          The bit width is determined by the TAP's IR length (from jtag_scan).
          Example: irscan with value "0x0E" selects the IDCODE register on
          many ARM TAPs.
        - drscan: Shifts data through the currently selected Data Register.
          The `bits` parameter is REQUIRED and must match the width of the
          selected data register (e.g. 32 for IDCODE, 35 for a boundary
          scan register). Returns the data shifted out of the DR.

        The response contains value_out (hex string) which is the data
        captured from TDO during the shift. For irscan, this is the
        previous IR contents. For drscan, this is the DR contents before
        the new value was shifted in.

        Typical workflow:
          1. jtag_scan() to discover TAPs and get TAP names
          2. irscan to select a data register (e.g. "0x0E" for IDCODE)
          3. drscan to read/write the selected register (e.g. bits=32)

        Args:
            tap: TAP name from jtag_scan (e.g. "stm32f1x.cpu", "stm32f1x.bs").
            operation: "irscan" to load an instruction, "drscan" to shift data.
            value: Hex value to shift in (e.g. "0x0E", "0x00000000"). For drscan,
                   this is the data to shift into DR; use "0x0" to read without
                   modifying.
            bits: Number of bits to shift for drscan. REQUIRED for drscan, ignored
                  for irscan. Must match the target data register width.
        """
        state = _get_state(ctx)
        require_session(state)

        int_value = parse_hex(value, "shift value")
        jtag = state.session.jtag

        if operation == "irscan":
            result = await jtag.irscan(tap, int_value)
            return JTAGShiftResult(
                operation="irscan",
                tap=tap,
                value_in=value,
                value_out=f"0x{result:x}",
            )
        elif operation == "drscan":
            if bits is None:
                raise OpenOCDError("bits parameter is required for drscan")
            result = await jtag.drscan(tap, bits, int_value)
            return JTAGShiftResult(
                operation="drscan",
                tap=tap,
                value_in=value,
                value_out=f"0x{result:x}",
                bits=bits,
            )
        else:
            raise OpenOCDError(f"Unknown operation: {operation}. Use 'irscan' or 'drscan'.")
