"""Connection management — connect, start_openocd, disconnect."""

import contextlib
import logging

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

from mcjtag.models import ConnectionStatus
from mcjtag.state import SessionState

log = logging.getLogger(__name__)


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


class ConnectionMixin(MCPMixin):
    """Tools for connecting to and disconnecting from OpenOCD."""

    @mcp_tool()
    async def connect(
        self,
        ctx: Context,
        host: str = "localhost",
        port: int = 6666,
    ) -> ConnectionStatus:
        """Connect to a running OpenOCD instance over TCL RPC.

        Use this when OpenOCD is already running (started manually or by
        another tool). For auto-starting OpenOCD, use start_openocd instead.

        Args:
            host: OpenOCD TCL RPC hostname.
            port: OpenOCD TCL RPC port (default 6666).
        """
        state = _get_state(ctx)
        async with state.lock:
            if state.session:
                return ConnectionStatus(
                    connected=True,
                    host=host,
                    port=port,
                    message="Already connected. Disconnect first to reconnect.",
                )

            try:
                state.session = await Session.connect(host, port)
            except OpenOCDError as exc:
                return ConnectionStatus(
                    connected=False,
                    host=host,
                    port=port,
                    message=f"Connection failed: {exc}",
                )

            log.info("Connected to OpenOCD at %s:%d", host, port)
            return ConnectionStatus(
                connected=True,
                host=host,
                port=port,
                message=f"Connected to OpenOCD at {host}:{port}",
            )

    @mcp_tool()
    async def start_openocd(
        self,
        ctx: Context,
        config: str,
        tcl_port: int = 6666,
        extra_args: str = "",
    ) -> ConnectionStatus:
        """Spawn an OpenOCD process with the given config, then connect.

        This starts a new OpenOCD subprocess and connects to its TCL RPC
        interface. The process will be stopped when you call disconnect.

        Args:
            config: OpenOCD config string (e.g. "-f interface/cmsis-dap.cfg").
            tcl_port: TCL RPC port to use (default 6666).
            extra_args: Additional CLI arguments for OpenOCD (space-separated).
        """
        state = _get_state(ctx)
        async with state.lock:
            if state.session:
                return ConnectionStatus(
                    connected=True,
                    message="Already connected. Disconnect first to start a new instance.",
                )

            extras = extra_args.split() if extra_args else None
            try:
                state.session = await Session.start(
                    config,
                    tcl_port=tcl_port,
                    extra_args=extras,
                )
            except OpenOCDError as exc:
                return ConnectionStatus(
                    connected=False,
                    host="localhost",
                    port=tcl_port,
                    message=f"Failed to start OpenOCD: {exc}",
                )

            state.managed_process = True

            pid = None
            with contextlib.suppress(AttributeError, TypeError):
                pid = state.session._process._proc.pid

            log.info("Started OpenOCD (pid=%s) on port %d", pid, tcl_port)
            return ConnectionStatus(
                connected=True,
                host="localhost",
                port=tcl_port,
                process_pid=pid,
                message=f"OpenOCD started and connected on port {tcl_port}",
            )

    @mcp_tool()
    async def disconnect(self, ctx: Context) -> ConnectionStatus:
        """Close the OpenOCD connection and stop the process if we started it.

        Safe to call even if not connected.
        """
        state = _get_state(ctx)
        async with state.lock:
            if not state.session:
                return ConnectionStatus(
                    connected=False,
                    message="Not connected.",
                )

            was_managed = state.managed_process
            try:
                await state.session.close()
            except Exception:
                log.warning("Error closing OpenOCD session", exc_info=True)
            finally:
                state.session = None
                state.managed_process = False

            msg = "Disconnected and stopped OpenOCD process" if was_managed else "Disconnected"
            log.info(msg)
            return ConnectionStatus(connected=False, message=msg)
