"""Flash memory info and programming tools."""

import logging
from pathlib import Path

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

from mcjtag.mixins.helpers import require_session
from mcjtag.models import FlashBankInfo, FlashBanksResult, FlashProgramResult
from mcjtag.state import SessionState

log = logging.getLogger(__name__)

# C3: Allowed firmware file extensions
_ALLOWED_EXTENSIONS = {".bin", ".hex", ".elf", ".ihex", ".srec", ".s19"}

# C3: Maximum flash image size (16 MB covers virtually all embedded targets)
_MAX_FLASH_IMAGE_SIZE = 16 * 1024 * 1024


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


class FlashMixin(MCPMixin):
    """Tools for flash memory inspection and programming.

    Flash programming is a destructive operation that overwrites the
    target's firmware. Use flash_info first to understand the flash layout,
    then flash_program to write a new image.
    """

    @mcp_tool()
    async def flash_info(self, ctx: Context) -> FlashBanksResult:
        """List all flash banks on the target with size, base address, and sector count.

        Returns flash bank topology: driver name, base address, total size,
        number of sectors, and target association for each bank. Use this
        before flash_program to verify the flash layout and confirm the
        expected base address.

        A sectors value of -1 means the detailed sector map could not be
        read (the bank exists but sector info is unavailable).

        Typical results for STM32F103: 1 bank at 0x08000000, 128KB, 128 sectors.
        """
        state = _get_state(ctx)
        require_session(state)

        banks = await state.session.flash.banks()
        bank_infos = []

        for bank in banks:
            # M2: Log sector-count failures (use -1 sentinel, not silent 0)
            try:
                detailed = await state.session.flash.info(bank.index)
                sector_count = len(detailed.sectors)
            except Exception as exc:
                log.warning(
                    "Could not get sector info for bank %d: %s",
                    bank.index,
                    exc,
                )
                sector_count = -1

            bank_infos.append(
                FlashBankInfo(
                    index=bank.index,
                    name=bank.name,
                    base=f"0x{bank.base:08x}",
                    size=bank.size,
                    sectors=sector_count,
                    target=bank.target,
                )
            )

        return FlashBanksResult(banks=bank_infos, count=len(bank_infos))

    @mcp_tool()
    async def flash_program(
        self,
        ctx: Context,
        path: str,
        erase: bool = True,
        verify: bool = True,
        address: str | None = None,
    ) -> FlashProgramResult:
        """Program a firmware image into flash memory.

        DESTRUCTIVE: This overwrites the target's flash contents. The
        operation performs up to three phases: erase affected sectors, write
        the image data, and read-back verify. Progress is reported for each
        phase.

        Supported file formats:
        - .elf: Recommended. Contains code, address info, and debug symbols.
          Address is extracted from ELF program headers automatically.
        - .hex / .ihex: Intel HEX format. Contains address info per record.
        - .srec / .s19: Motorola S-record. Contains address info per record.
        - .bin: Raw binary. Contains NO address metadata -- you MUST provide
          the `address` parameter (e.g. "0x08000000" for STM32 flash base)
          or the data will be written to address 0x00000000.

        Maximum file size: 16 MB. Empty files (0 bytes) are rejected.

        Erase behavior: When erase=True (default), only the flash sectors
        overlapping the image are erased -- not the entire flash. This is
        sector-granular, so a 1KB image may erase a 4KB sector.

        Verify behavior: When verify=True (default), the flash contents are
        read back after writing and compared byte-by-byte against the image.
        If verification fails, the result reports success=False.

        After programming, the target is typically left halted. Use
        target_control(action="reset_run") to boot the new firmware.

        Args:
            path: Absolute or relative path to the firmware image file.
                  Must have one of the allowed extensions (.bin, .hex, .elf,
                  .ihex, .srec, .s19).
            erase: Erase affected flash sectors before writing (default True).
                   Set False only if sectors were pre-erased or for partial
                   updates within an already-erased sector.
            verify: Read back and verify flash contents after writing
                    (default True). Disable for faster programming when
                    verification is not needed.
            address: Base address for .bin files, in hex (e.g. "0x08000000").
                     REQUIRED for .bin files -- they have no embedded address.
                     Must be 4-byte aligned. Ignored for .hex/.elf/.ihex/.srec
                     which carry their own address information.
        """
        state = _get_state(ctx)
        require_session(state)

        image_path = Path(path).resolve()

        # Validate address if provided
        address_int: int | None = None
        if address is not None:
            try:
                # Strip 0x/0X prefix if present, then parse as hex
                clean = address.strip()
                if clean.lower().startswith("0x"):
                    clean = clean[2:]
                address_int = int(clean, 16)
            except (ValueError, TypeError):
                return FlashProgramResult(
                    path=str(image_path),
                    erased=False,
                    verified=False,
                    success=False,
                    message=f"Invalid address '{address}': must be a hex value (e.g. '0x08000000')",
                )
            if address_int % 4 != 0:
                return FlashProgramResult(
                    path=str(image_path),
                    erased=False,
                    verified=False,
                    success=False,
                    message=(
                        f"Address 0x{address_int:08X} is not word-aligned (must be 4-byte aligned)"
                    ),
                )

        # C3: File existence
        if not image_path.exists():
            return FlashProgramResult(
                path=str(image_path),
                erased=False,
                verified=False,
                success=False,
                message=f"Firmware image not found: {image_path}",
            )

        # C3: File extension validation
        if image_path.suffix.lower() not in _ALLOWED_EXTENSIONS:
            return FlashProgramResult(
                path=str(image_path),
                erased=False,
                verified=False,
                success=False,
                message=(
                    f"Unsupported file type '{image_path.suffix}'. "
                    f"Allowed: {', '.join(sorted(_ALLOWED_EXTENSIONS))}"
                ),
            )

        # C3: File size validation
        file_size = image_path.stat().st_size
        if file_size == 0:
            return FlashProgramResult(
                path=str(image_path),
                erased=False,
                verified=False,
                success=False,
                message="Firmware image is empty (0 bytes)",
            )
        if file_size > _MAX_FLASH_IMAGE_SIZE:
            return FlashProgramResult(
                path=str(image_path),
                erased=False,
                verified=False,
                success=False,
                message=(
                    f"Firmware image is {file_size} bytes, "
                    f"exceeds {_MAX_FLASH_IMAGE_SIZE} byte limit"
                ),
            )

        # Progress: 3 phases — erase, write, verify
        total_phases = 1 + int(erase) + int(verify)
        phase = 0

        if erase:
            await ctx.report_progress(phase, total_phases, "Erasing flash...")
            phase += 1

        await ctx.report_progress(phase, total_phases, f"Writing {image_path.name}...")

        # M1: Only catch OpenOCD/Flash/OS errors, not programming bugs
        try:
            await state.session.flash.write_image(
                image_path,
                erase=erase,
                verify=verify,
                address=address_int,
            )
            await ctx.report_progress(
                total_phases,
                total_phases,
                f"Programmed {image_path.name} successfully",
            )
            return FlashProgramResult(
                path=str(image_path),
                erased=erase,
                verified=verify,
                success=True,
                message=f"Programmed {image_path.name} successfully",
            )
        except (OpenOCDError, FlashError, OSError) as exc:
            log.error("Flash programming failed: %s", exc)
            return FlashProgramResult(
                path=str(image_path),
                erased=erase,
                verified=False,
                success=False,
                message=f"Flash error: {exc}",
            )
