"""Target state inspection and execution control."""

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 TargetControlResult, TargetStateResult
from mcjtag.state import SessionState

log = logging.getLogger(__name__)


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


class TargetMixin(MCPMixin):
    """Tools for inspecting and controlling target execution state.

    The target must be halted for most debug operations (memory reads,
    register reads, single-step). Use target_state to check, and
    target_control(action="halt") to stop it.
    """

    @mcp_tool()
    async def target_state(self, ctx: Context) -> TargetStateResult:
        """Get the current target execution state and program counter.

        Returns the target state as one of: "halted", "running", "reset",
        or "debug-running". The program counter (pc) is only available when
        the target is halted.

        Call this before operations that require a halted target (memory
        reads, register access, single-step) to verify the target state.
        If state is "running", call target_control(action="halt") first.
        """
        state = _get_state(ctx)
        require_session(state)

        ts = await state.session.target.state()
        pc_hex = f"0x{ts.current_pc:08x}" if ts.current_pc is not None else None
        return TargetStateResult(
            state=ts.state,
            pc=pc_hex,
            halted=ts.state == "halted",
            name=ts.name,
        )

    @mcp_tool()
    async def target_control(
        self,
        ctx: Context,
        action: Literal["halt", "resume", "step", "reset_halt", "reset_run"],
        address: str | None = None,
    ) -> TargetControlResult:
        """Control target execution -- halt, resume, single-step, or reset.

        Actions:
        - halt: Stop the CPU at its current instruction. Required before
          reading registers, accessing memory, or setting software breakpoints.
        - resume: Continue execution from the current PC, or from a specific
          address if provided. The target enters "running" state.
        - step: Execute exactly one instruction and halt again. Useful for
          stepping through code. Optionally starts from a specific address.
        - reset_halt: Assert hardware reset and halt at the reset vector
          (typically 0x08000000 on STM32). The chip is fully reinitialized
          but stopped before any code executes. Use this to start a fresh
          debug session.
        - reset_run: Assert hardware reset and let the chip run freely.
          The target enters "running" state. Debug state (breakpoints,
          watchpoints) is preserved across reset on most targets.

        The address parameter is only used with resume and step. It must be
        2-byte aligned (ARM Thumb requirement). If omitted, execution
        continues from the current PC.

        MCJTAG_SAFE_WRITE_RANGES does NOT apply to target control -- these
        actions do not write to memory. However, resuming execution at an
        incorrect address can crash the target or cause unpredictable behavior.

        Args:
            action: Control action -- "halt", "resume", "step", "reset_halt",
                    or "reset_run".
            address: Optional hex address for resume/step (e.g. "0x08001000").
                     Must be 2-byte aligned. Ignored for halt/reset actions.
        """
        state = _get_state(ctx)
        require_session(state)

        addr = None
        if address:
            addr = parse_hex(address)
            # M4: ARM Thumb instructions require 2-byte aligned addresses
            if addr % 2 != 0:
                raise OpenOCDError(
                    f"Address 0x{addr:08x} is not 2-byte aligned. "
                    f"ARM Thumb instructions require even addresses."
                )

        target = state.session.target

        if action == "halt":
            ts = await target.halt()
        elif action == "resume":
            await target.resume(addr)
            ts = await target.state()
        elif action == "step":
            ts = await target.step(addr)
        elif action == "reset_halt":
            await target.reset("halt")
            ts = await target.state()
        elif action == "reset_run":
            await target.reset("run")
            ts = await target.state()
        else:
            raise OpenOCDError(f"Unknown action: {action}")

        pc_hex = f"0x{ts.current_pc:08x}" if ts.current_pc is not None else None
        return TargetControlResult(
            action=action,
            state=ts.state,
            pc=pc_hex,
            halted=ts.state == "halted",
        )
