"""Memory read, write, and search operations."""

import logging
import os
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 (
    check_write_address,
    format_hexdump,
    parse_hex,
    require_session,
    values_to_bytes,
)
from mcjtag.models import MemoryDump, MemorySearchResult, MemoryWriteResult
from mcjtag.state import SessionState

log = logging.getLogger(__name__)

# H4: Upper bound on read count to prevent accidental multi-MB reads
_MAX_READ_COUNT = 4096

# H3: Upper bound on search range (1 MB default, configurable)
_MAX_SEARCH_RANGE = 1 * 1024 * 1024


def _get_max_search_range() -> int:
    raw = os.environ.get("MCJTAG_MAX_SEARCH_RANGE")
    if raw:
        try:
            return int(raw)
        except ValueError:
            pass
    return _MAX_SEARCH_RANGE


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


class MemoryMixin(MCPMixin):
    """Tools for reading, writing, and searching target memory.

    All addresses are hex strings (e.g. "0x20000000"). Common Cortex-M
    memory regions: 0x00000000 (code/flash alias), 0x08000000 (flash),
    0x20000000 (SRAM), 0x40000000 (peripherals), 0xE0000000 (system/debug).
    """

    @mcp_tool()
    async def read_memory(
        self,
        ctx: Context,
        address: str,
        count: int = 1,
        width: Literal[8, 16, 32, 64] = 32,
    ) -> MemoryDump:
        """Read memory from the target at the given address.

        Returns both raw hex values (in the `values` array) and a formatted
        hexdump view with ASCII representation (like xxd). The target must
        typically be halted for SRAM/flash reads, but peripheral registers
        can often be read while running.

        Alignment: The address must be naturally aligned to the access width
        (e.g. 4-byte aligned for 32-bit reads). Unaligned 16/32/64-bit reads
        will log a warning and may cause a bus fault on Cortex-M0 targets.
        Use width=8 for unaligned access.

        Important: Hardware registers with read-clear semantics (e.g. status
        registers) are only read once -- the hexdump is built from the same
        read, not a second access.

        Args:
            address: Start address in hex (e.g. "0x08000000" for flash,
                     "0x20000000" for SRAM, "0xE000ED00" for CPUID).
            count: Number of elements to read, 1-4096 (default 1). Total bytes
                   read = count * (width / 8). For example, count=256 width=8
                   reads 256 bytes; count=64 width=32 reads 256 bytes.
            width: Element width in bits -- 8, 16, 32, or 64 (default 32).
                   Use 8 for byte-granularity or unaligned addresses. Use 32
                   for standard register/memory reads.
        """
        state = _get_state(ctx)
        require_session(state)

        addr = parse_hex(address)

        # I-5: Warn on unaligned reads (faults on Cortex-M0)
        alignment = width // 8
        if alignment > 1 and addr % alignment != 0:
            log.warning(
                "Unaligned %d-bit read at 0x%08x (not %d-byte aligned). "
                "This may fault on Cortex-M0 targets.",
                width,
                addr,
                alignment,
            )

        # H4: Bound the read count
        if count <= 0:
            raise OpenOCDError(f"count must be positive, got {count}")
        if count > _MAX_READ_COUNT:
            raise OpenOCDError(
                f"count {count} exceeds maximum {_MAX_READ_COUNT}. "
                f"Use multiple smaller reads for large regions."
            )

        mem = state.session.memory

        if width == 8:
            values = await mem.read_u8(addr, count)
        elif width == 16:
            values = await mem.read_u16(addr, count)
        elif width == 32:
            values = await mem.read_u32(addr, count)
        elif width == 64:
            values = await mem.read_u64(addr, count)
        else:
            raise OpenOCDError(f"Invalid width: {width}. Use 8, 16, 32, or 64.")

        hex_values = [f"0x{v:0{width // 4}x}" for v in values]

        # H1: Build hexdump from already-read values (no re-read).
        # Many hardware registers have read-clear semantics — reading
        # twice would give different values or advance internal state.
        byte_data = values_to_bytes(values, width)
        hexdump = format_hexdump(addr, byte_data)

        return MemoryDump(
            address=f"0x{addr:08x}",
            width=width,
            count=count,
            values=hex_values,
            hexdump=hexdump,
        )

    @mcp_tool()
    async def write_memory(
        self,
        ctx: Context,
        address: str,
        values: list[str],
        width: Literal[8, 16, 32] = 32,
    ) -> MemoryWriteResult:
        """Write values to target memory at the given address.

        Safety: By default, writes are restricted to SRAM regions
        (0x20000000-0x20100000). Writes outside this range are rejected.
        Configure allowed ranges with the MCJTAG_SAFE_WRITE_RANGES env var:
        - Custom ranges: "0x20000000-0x20100000,0x10000000-0x10010000"
        - Unrestricted: "none" (disables all address validation)

        Alignment: The address must be naturally aligned to the access width.
        Unaligned 16/32-bit writes are rejected (use width=8 for unaligned).

        Each value must fit within the specified width (e.g. max 0xFF for
        8-bit, 0xFFFF for 16-bit, 0xFFFFFFFF for 32-bit).

        Note: This writes to RAM, not flash. For flash programming, use
        flash_program instead.

        Args:
            address: Start address in hex (e.g. "0x20000000"). Must be within
                     the safe write range and aligned to width.
            values: List of hex values to write sequentially starting at address.
                    Example: ["0xDEADBEEF", "0x12345678"] writes two 32-bit words.
            width: Element width in bits -- 8, 16, or 32 (default 32). Note: 64-bit
                   writes are not supported (use two 32-bit writes instead).
        """
        state = _get_state(ctx)
        require_session(state)

        addr = parse_hex(address)

        # I-5: Reject unaligned writes (can corrupt hardware state)
        alignment = width // 8
        if alignment > 1 and addr % alignment != 0:
            raise OpenOCDError(
                f"Unaligned {width}-bit write at 0x{addr:08x} "
                f"(not {alignment}-byte aligned). "
                f"Use 8-bit writes for unaligned addresses."
            )

        # M3: Validate each value fits in the specified width
        max_value = (1 << width) - 1
        int_values: list[int] = []
        for v_str in values:
            v = parse_hex(v_str, "value")
            if v > max_value:
                raise OpenOCDError(f"Value 0x{v:x} exceeds {width}-bit maximum (0x{max_value:x})")
            int_values.append(v)

        # C2: Check write address falls within safe ranges
        byte_size = len(int_values) * (width // 8)
        check_write_address(addr, byte_size)

        mem = state.session.memory

        if width == 8:
            await mem.write_u8(addr, int_values)
        elif width == 16:
            await mem.write_u16(addr, int_values)
        elif width == 32:
            await mem.write_u32(addr, int_values)
        else:
            raise OpenOCDError(f"Invalid width: {width}. Use 8, 16, or 32.")

        return MemoryWriteResult(
            address=f"0x{addr:08x}",
            width=width,
            count=len(int_values),
            written=True,
        )

    @mcp_tool()
    async def search_memory(
        self,
        ctx: Context,
        pattern: str,
        start: str,
        end: str,
    ) -> MemorySearchResult:
        """Search for a byte pattern within a memory range.

        Reads memory in 4KB chunks and searches for the pattern at every
        byte offset (no alignment requirement on the pattern itself). The
        search is byte-oriented, so multi-byte patterns match in little-endian
        memory order as written. Returns a list of addresses where the
        pattern was found.

        Max search range is 1 MB by default. Set the MCJTAG_MAX_SEARCH_RANGE
        env var (in bytes, decimal) to increase. Larger ranges take longer
        since each 4KB chunk requires a debug probe transfer.

        Pattern format: Hex bytes as a continuous string WITHOUT "0x" prefix.
        Each byte is two hex characters. Examples:
        - "DEADBEEF" matches bytes [0xDE, 0xAD, 0xBE, 0xEF] in memory
        - "48656C6C6F" matches the ASCII string "Hello"
        - "FF" matches any byte with value 0xFF

        Note: On a 32-bit little-endian target, the 32-bit value 0xDEADBEEF
        is stored as bytes [0xEF, 0xBE, 0xAD, 0xDE]. To find the word value
        0xDEADBEEF, search for pattern "EFBEADDE".

        Args:
            pattern: Hex byte string to search for (e.g. "DEADBEEF", "48656C6C6F").
                     No "0x" prefix. Must be an even number of hex characters.
            start: Start address in hex, inclusive (e.g. "0x20000000"). Should be
                   aligned to the region you are searching for best performance.
            end: End address in hex, exclusive (e.g. "0x20010000"). Must be after
                 start. Range (end - start) must not exceed 1 MB (default) or the
                 value of MCJTAG_MAX_SEARCH_RANGE.
        """
        state = _get_state(ctx)
        require_session(state)

        # Validate the hex pattern
        try:
            pattern_bytes = bytes.fromhex(pattern)
        except ValueError as exc:
            raise OpenOCDError(
                f"Invalid hex pattern: {pattern!r}. "
                f"Expected hex bytes like 'DEADBEEF' or '48656C6C6F'."
            ) from exc

        start_addr = parse_hex(start, "start address")
        end_addr = parse_hex(end, "end address")

        # H3: Validate range bounds
        if end_addr <= start_addr:
            raise OpenOCDError(
                f"Invalid range: end (0x{end_addr:08x}) must be after start (0x{start_addr:08x})"
            )

        max_range = _get_max_search_range()
        range_size = end_addr - start_addr
        if range_size > max_range:
            raise OpenOCDError(
                f"Search range {range_size} bytes exceeds maximum "
                f"{max_range} bytes. Narrow the range or set "
                f"MCJTAG_MAX_SEARCH_RANGE env var."
            )

        await ctx.report_progress(
            0,
            range_size,
            f"Searching {range_size} bytes for pattern {pattern}...",
        )

        matches = await state.session.memory.search(pattern_bytes, start_addr, end_addr)

        await ctx.report_progress(range_size, range_size, "Search complete")
        hex_matches = [f"0x{m:08x}" for m in matches]

        return MemorySearchResult(
            pattern=pattern,
            start=f"0x{start_addr:08x}",
            end=f"0x{end_addr:08x}",
            matches=hex_matches,
            count=len(hex_matches),
        )
