"""mcjtag — JTAG/SWD MCP Server for hardware debugging via OpenOCD."""

from __future__ import annotations

import logging
import os
from contextlib import asynccontextmanager
from pathlib import Path

from fastmcp import FastMCP

from mcjtag.mixins import (
    BreakpointMixin,
    ConnectionMixin,
    DiagnosticsMixin,
    FlashMixin,
    JTAGMixin,
    MemoryMixin,
    RawMixin,
    RegistersMixin,
    RTTMixin,
    SVDMixin,
    SWDMixin,
    TargetMixin,
    TransportMixin,
)
from mcjtag.state import SessionState

log = logging.getLogger(__name__)


# ======================================================================
# Lifespan — session management
# ======================================================================


@asynccontextmanager
async def lifespan(server: FastMCP):
    """Manage the OpenOCD session lifecycle.

    Supports three startup modes via environment variables:
    - OPENOCD_CONFIG: Spawn OpenOCD with this config, then connect
    - OPENOCD_HOST + OPENOCD_PORT: Connect to existing OpenOCD
    - Neither: Start with no connection (LLM calls connect manually)

    OPENOCD_SVD: Auto-load an SVD file on startup (requires a connection)
    """
    from openocd.session import Session

    state = SessionState()

    config = os.environ.get("OPENOCD_CONFIG")
    host = os.environ.get("OPENOCD_HOST")

    # M5: Graceful port parsing with clear error message
    port_str = os.environ.get("OPENOCD_PORT", "6666")
    try:
        port = int(port_str)
    except ValueError:
        log.error(
            "Invalid OPENOCD_PORT=%r (must be an integer), using default 6666",
            port_str,
        )
        port = 6666

    if config:
        log.info("Auto-starting OpenOCD with config: %s", config)
        state.session = await Session.start(config, tcl_port=port)
        state.managed_process = True
    elif host:
        log.info("Auto-connecting to OpenOCD at %s:%d", host, port)
        state.session = await Session.connect(host, port)

    # Auto-load SVD if configured
    svd_path = os.environ.get("OPENOCD_SVD")
    if svd_path and state.session:
        try:
            await state.session.svd.load(Path(svd_path))
            state.svd_path = svd_path
            log.info("Auto-loaded SVD: %s", svd_path)
        except Exception:
            log.warning("Failed to load SVD file: %s", svd_path, exc_info=True)

    # Make state available to resource functions
    from mcjtag.resources import set_state

    set_state(state)

    try:
        yield state
    finally:
        # H5: Close session BEFORE clearing state, and never let
        # close() failure prevent state cleanup.
        if state.session:
            try:
                await state.session.close()
                log.info("OpenOCD session closed")
            except Exception:
                log.warning("Error closing OpenOCD session", exc_info=True)
        set_state(SessionState())


# ======================================================================
# Server class — combines all mixins
# ======================================================================


class MCJTAGServer(
    ConnectionMixin,
    DiagnosticsMixin,
    TargetMixin,
    MemoryMixin,
    RegistersMixin,
    FlashMixin,
    JTAGMixin,
    SWDMixin,
    BreakpointMixin,
    RTTMixin,
    TransportMixin,
    SVDMixin,
    RawMixin,
):
    """JTAG/SWD MCP Server for hardware debugging via OpenOCD."""


# ======================================================================
# FastMCP app
# ======================================================================

mcp = FastMCP(
    "mcjtag",
    instructions="""JTAG/SWD MCP Server for hardware debugging via OpenOCD.

Getting started:
- If not auto-connected, call connect(host, port) or start_openocd(config)
- Call probe_diagnostics() to see what's working and what capabilities are available
- Use disconnect() to cleanly close the session

Typical workflows:
1. Run diagnostics: probe_diagnostics() to check probe, target, and capabilities
2. Identify chip: jtag_scan() to see IDCODE, or swd_info() for DAP identification
3. Inspect state: target_state() to check if halted/running
4. Read memory: read_memory("0x08000000", count=16) for flash contents
5. Read registers: read_registers() for all CPU registers
6. Flash firmware: flash_program("/path/to/firmware.bin")
7. Decode peripherals: svd_inspect(svd_path="/path/to/chip.svd", peripheral="GPIOA")
8. Set breakpoints: breakpoint_add("0x08001234") or watchpoint_add("0x20001000", 4)
9. SWD/DAP access: swd_list_aps() to enumerate Access Ports, swd_read_dp/swd_read_ap for registers
10. RTT logging: rtt_setup() + rtt_start() + rtt_read() for SEGGER Real-Time Transfer
11. Transport/adapter: transport_info() for current transport, adapter_speed() to adjust clock
12. Raw escape hatch: raw_command("version") for any OpenOCD command

The target must be halted for register access and most memory operations.
Use target_control(action="halt") to stop the CPU first.
""",
    lifespan=lifespan,
)

# Register all mixin tools
server = MCJTAGServer()
server.register_all(mcp)

# Import resources to register them
import mcjtag.resources  # noqa: E402, F401, I001


# ======================================================================
# Prompts
# ======================================================================


@mcp.prompt()
def identify_chip() -> str:
    """Identify the connected chip by scanning the JTAG chain."""
    return """Identify the connected chip:
1. Call jtag_scan() to enumerate the JTAG chain
2. Note the IDCODE of each TAP
3. Look up the IDCODE to identify the chip manufacturer and part number
4. Call target_state() to confirm the debug connection is working
5. Report the chip identification and connection status"""


@mcp.prompt()
def debug_crash() -> str:
    """Diagnose a crash by inspecting CPU state and stack."""
    return """Diagnose a crash on the target:
1. Call target_control(action="halt") to stop the CPU
2. Call read_registers(names=["pc", "lr", "sp", "xPSR"]) to read key registers
3. The PC shows where execution stopped
4. The LR shows the return address (caller)
5. Read the stack: read_memory(address=<sp value>, count=16, width=32)
6. Check if xPSR indicates a fault (look for exception number in bits 0-8)
7. If in a fault handler, read the stacked frame at the SP value
8. Report the crash location, call chain, and likely cause"""


@mcp.prompt()
def reverse_engineer_peripheral(peripheral: str) -> str:
    """Decode all registers of a peripheral using SVD metadata."""
    return f"""Reverse-engineer the {peripheral} peripheral:
1. If no SVD is loaded, ask for the SVD file path and call svd_inspect(svd_path=...)
2. Call svd_inspect(peripheral="{peripheral}") to read and decode all registers
3. For each register, analyze the bitfield values
4. Determine the current configuration of the peripheral
5. Report what the peripheral is doing based on the register values
6. Note any unusual or potentially problematic settings"""


@mcp.prompt()
def flash_and_verify(image_path: str) -> str:
    """Full flash workflow: erase, program, verify, and reset."""
    return f"""Flash firmware and verify:
1. Call target_control(action="halt") to stop the CPU
2. Call flash_info() to check flash bank topology
3. Call flash_program(path="{image_path}", erase=True, verify=True)
4. Check the result for success/failure
5. If successful, call target_control(action="reset_run") to start the new firmware
6. Wait a moment, then call target_state() to confirm it's running
7. Report the flash result and post-reset status"""


@mcp.prompt()
def dap_discovery() -> str:
    """Enumerate and identify the Debug Access Port and its Access Ports."""
    return """Discover the DAP architecture:
1. Call swd_info() to get the DPIDR and AP count
2. Call swd_list_aps() to enumerate all Access Ports
3. For each MEM-AP, note the ROM table base address (from the 'base' field)
4. Read the DPIDR via swd_read_dp(address="0x0") — decode the version, designer, and part
5. If DPv2+, read TARGETID via swd_read_dp(address="0x24") for chip identification
6. Report the DAP topology: which APs exist, their types, and the debug architecture version
7. If a MEM-AP is present, read CPUID at 0xE000ED00 via read_memory to identify the core"""


@mcp.prompt()
def memory_map() -> str:
    """Build a memory map by reading the vector table and probing regions."""
    return """Build a memory map of the target:
1. Call target_control(action="halt") to stop the CPU
2. Read the vector table: read_memory(address="0x08000000", count=8, width=32)
   - Word 0: Initial stack pointer (top of RAM)
   - Word 1: Reset vector (start of code)
   - Words 2-7: NMI, HardFault, and other exception vectors
3. Call flash_info() to get flash bank sizes and base addresses
4. From the initial SP, deduce the RAM region (SP points to end of RAM)
5. Read a few bytes at common peripheral base addresses to verify they're accessible:
   - 0x40000000 (APB1 peripherals)
   - 0x40010000 (APB2 peripherals)
   - 0xE000E000 (System control block)
6. Build and report the memory map:
   - Flash: base, size, sector layout
   - RAM: base, estimated size
   - Peripheral regions
   - Vector table summary"""


# ======================================================================
# Entry point
# ======================================================================


def main():
    """STDIO MCP server entry point."""
    try:
        from importlib.metadata import version

        package_version = version("mcjtag")
    except Exception:
        package_version = "2026.2.12"

    print(f"\U0001f50c mcjtag v{package_version}", flush=True)
    mcp.run()


if __name__ == "__main__":
    main()
