"""CPU register read and write operations."""

import logging

from fastmcp import Context
from fastmcp.contrib.mcp_mixin import MCPMixin, mcp_tool

from mcjtag.mixins.helpers import parse_hex, require_session
from mcjtag.models import RegisterSet, RegisterWriteResult
from mcjtag.state import SessionState

log = logging.getLogger(__name__)


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


class RegistersMixin(MCPMixin):
    """Tools for reading and writing CPU registers."""

    @mcp_tool()
    async def read_registers(
        self,
        ctx: Context,
        names: list[str] | None = None,
    ) -> RegisterSet:
        """Read CPU registers from the halted target.

        If no names are given, reads all registers. Otherwise reads only
        the specified registers. The target must be halted.

        Common ARM Cortex-M registers: r0-r12, sp, lr, pc, xPSR, msp, psp

        Args:
            names: List of register names to read, or None for all registers.
        """
        state = _get_state(ctx)
        require_session(state)

        regs = state.session.registers

        if names is None:
            all_regs = await regs.read_all()
            reg_dict = {name: f"0x{r.value:0{r.size // 4}x}" for name, r in all_regs.items()}
        else:
            values = await regs.read_many(names)
            reg_dict = {name: f"0x{v:08x}" for name, v in values.items()}

        return RegisterSet(registers=reg_dict, count=len(reg_dict))

    @mcp_tool()
    async def write_register(
        self,
        ctx: Context,
        name: str,
        value: str,
    ) -> RegisterWriteResult:
        """Write a value to a CPU register on the halted target.

        The target must be halted for register writes.

        Args:
            name: Register name (e.g. "pc", "r0", "sp").
            value: Value to write in hex (e.g. "0x08001000").
        """
        state = _get_state(ctx)
        require_session(state)

        int_value = parse_hex(value, "register value")
        await state.session.registers.write(name, int_value)

        return RegisterWriteResult(
            name=name,
            value=f"0x{int_value:08x}",
            written=True,
        )
