"""VNC connection management tools."""

import asyncio
import json

import asyncvnc2
from fastmcp import Context
from fastmcp.exceptions import ToolError

from mcvnc._state import CONNECTION_TIMEOUT, DEFAULT_PORT, VncSession, get_state, mcp


@mcp.tool()
async def vnc_connect(
    ctx: Context,
    host: str,
    port: int = DEFAULT_PORT,
    password: str | None = None,
    session_id: str | None = None,
) -> str:
    """Connect to a VNC server.

    Args:
        host: VNC server hostname or IP address.
        port: VNC server port (default from MCVNC_DEFAULT_PORT env, usually 5900).
        password: VNC password if the server requires authentication.
        session_id: Optional alias for this session. Defaults to "host:port".
    """
    state = get_state(ctx)
    sid = session_id or f"{host}:{port}"

    if sid in state.sessions:
        raise ToolError(f"Session '{sid}' already exists. Disconnect first or use a different ID.")

    cm = asyncvnc2.connect(host, port=port, password=password)
    try:
        client = await asyncio.wait_for(cm.__aenter__(), timeout=CONNECTION_TIMEOUT)
    except (asyncio.TimeoutError, Exception) as e:
        # Clean up the context manager's socket if __aenter__ opened one
        # before failing. asyncvnc2 opens the TCP connection early in __aenter__,
        # so cancellation or auth failure can leave an orphaned StreamWriter.
        try:
            await cm.__aexit__(type(e), e, e.__traceback__)
        except Exception:
            pass
        if isinstance(e, asyncio.TimeoutError):
            raise ToolError(
                f"Connection to {host}:{port} timed out after {CONNECTION_TIMEOUT}s. "
                f"Verify the VNC server is running and reachable."
            )
        raise ToolError(f"Failed to connect to {host}:{port}: {e}")

    session = VncSession(
        client=client,
        host=host,
        port=port,
        width=client.video.width,
        height=client.video.height,
        _cm=cm,
    )
    state.sessions[sid] = session

    return json.dumps({
        "session_id": sid,
        "host": host,
        "port": port,
        "desktop_name": client.video.name,
        "width": session.width,
        "height": session.height,
    })


@mcp.tool()
async def vnc_disconnect(ctx: Context, session_id: str) -> str:
    """Disconnect from a VNC server.

    Args:
        session_id: Session to disconnect (e.g. "host:port" or your alias).
    """
    state = get_state(ctx)
    return await state.disconnect(session_id)


@mcp.tool()
async def vnc_list_sessions(ctx: Context) -> str:
    """List all active VNC sessions."""
    state = get_state(ctx)
    if not state.sessions:
        return "No active VNC sessions. Use vnc_connect to establish a connection."

    sessions = []
    for sid, s in state.sessions.items():
        sessions.append({
            "session_id": sid,
            "host": s.host,
            "port": s.port,
            "width": s.width,
            "height": s.height,
        })
    return json.dumps(sessions)


@mcp.tool()
async def vnc_session_info(ctx: Context, session_id: str) -> str:
    """Get detailed info about a VNC session.

    Args:
        session_id: Session to inspect.
    """
    state = get_state(ctx)
    s = state.get(session_id)
    return json.dumps({
        "session_id": session_id,
        "host": s.host,
        "port": s.port,
        "desktop_name": s.client.video.name,
        "width": s.width,
        "height": s.height,
        "color_mode": s.client.video.mode,
    })
