"""MCP Server for driving Textual TUI apps headlessly."""

from __future__ import annotations

import io
import json

from mcp.server.fastmcp import FastMCP
from mcp.types import ToolAnnotations
from rich.console import Console

from textual_mcp.app_loader import load_app
from textual_mcp.serializers.widget_state import WidgetStateExtractor
from textual_mcp.session_manager import SessionManager

mcp = FastMCP("textual_mcp")
_session_manager = SessionManager()

# Register tools on the module-level server


def _register_tools(server: FastMCP, manager: SessionManager) -> None:
    """Register lifecycle tools on the given FastMCP server."""

    @server.tool(
        name="textual_launch",
        description="Launch a Textual app headlessly for QA interaction.",
        annotations=ToolAnnotations(destructiveHint=False, openWorldHint=False),
    )
    async def textual_launch(
        app_path: str,
        size: list[int] | None = None,
        session_id: str | None = None,
    ) -> str:
        if size is None:
            size = [120, 40]
        try:
            app = load_app(app_path)
            session = manager.create_session(session_id)
            await session.launch(app, size=tuple(size))
            return json.dumps({
                "session_id": session.session_id,
                "app_title": session.app.title,
                "screen_name": session.app.screen.name if session.app.screen else None,
            })
        except Exception as e:
            return json.dumps({"error": str(e)})

    @server.tool(
        name="textual_stop",
        description="Stop a running Textual app session.",
        annotations=ToolAnnotations(destructiveHint=True),
    )
    async def textual_stop(session_id: str) -> str:
        try:
            session = manager.get_session(session_id)
            errors = session.error_collector.drain()
            await manager.remove_session(session_id)
            return json.dumps({"stopped": True, "errors": errors})
        except Exception as e:
            return json.dumps({"error": str(e)})


# Register on the module-level server
_register_tools(mcp, _session_manager)


# --- Interaction tools (module-level for testability) ---


async def textual_press(session_id: str, keys: list[str]) -> str:
    """Simulate key presses in a running app session."""
    try:
        session = _session_manager.get_session(session_id)
        await session.pilot.press(*keys)
        await session.pilot.pause()
        focused = session.app.screen.focused
        return json.dumps({"focused_after": focused.id if focused else None})
    except Exception as e:
        return json.dumps({"error": str(e)})


async def textual_click(
    session_id: str,
    selector: str,
    offset: list[int] | None = None,
    times: int = 1,
) -> str:
    """Click a widget identified by CSS selector."""
    if offset is None:
        offset = [0, 0]
    try:
        session = _session_manager.get_session(session_id)
        with session.app_context():
            hit = await session.pilot.click(
                selector, offset=tuple(offset), times=times
            )
            await session.pilot.pause()
        return json.dumps({
            "widget_clicked": selector,
            "hit_target": hit,
        })
    except Exception as e:
        return json.dumps({"error": str(e)})


async def textual_type_text(
    session_id: str,
    text: str,
    submit: bool = False,
) -> str:
    """Type text into the currently focused input widget."""
    try:
        session = _session_manager.get_session(session_id)
        for char in text:
            await session.pilot.press(char)
        if submit:
            await session.pilot.press("enter")
        await session.pilot.pause()
        focused = session.app.screen.focused
        text_after = None
        if focused is not None and hasattr(focused, "value"):
            text_after = focused.value
        return json.dumps({"text_after": text_after})
    except Exception as e:
        return json.dumps({"error": str(e)})


async def textual_hover(
    session_id: str,
    selector: str,
    offset: list[int] | None = None,
) -> str:
    """Hover the mouse over a widget identified by CSS selector."""
    if offset is None:
        offset = [0, 0]
    try:
        session = _session_manager.get_session(session_id)
        with session.app_context():
            await session.pilot.hover(selector, offset=tuple(offset))
            await session.pilot.pause()
        return json.dumps({"widget_under": selector})
    except Exception as e:
        return json.dumps({"error": str(e)})


def _register_interaction_tools(server: FastMCP) -> None:
    """Register interaction tools on the FastMCP server."""
    server.tool(
        name="textual_press",
        description="Simulate key presses in a running Textual app.",
        annotations=ToolAnnotations(destructiveHint=False, openWorldHint=False),
    )(textual_press)
    server.tool(
        name="textual_click",
        description="Click a widget by CSS selector in a running Textual app.",
        annotations=ToolAnnotations(destructiveHint=False, openWorldHint=False),
    )(textual_click)
    server.tool(
        name="textual_type_text",
        description="Type text into the focused input widget.",
        annotations=ToolAnnotations(destructiveHint=False, openWorldHint=False),
    )(textual_type_text)
    server.tool(
        name="textual_hover",
        description="Hover the mouse over a widget by CSS selector.",
        annotations=ToolAnnotations(destructiveHint=False, openWorldHint=False),
    )(textual_hover)


_register_interaction_tools(mcp)


# --- Observation tools (module-level for testability) ---

_extractor = WidgetStateExtractor()


async def textual_snapshot(
    session_id: str,
    max_depth: int | None = None,
    include_styles: bool = False,
) -> str:
    """Snapshot current app state — widget tree with refs, focus, bindings, errors."""
    try:
        session = _session_manager.get_session(session_id)
        tree_text = session.serializer.serialize(
            session.app, max_depth=max_depth, include_styles=include_styles
        )
        focused = session.app.screen.focused
        focused_ref = None
        if focused is not None:
            for ref, widget in session.serializer._ref_map.items():
                if widget is focused:
                    focused_ref = ref
                    break
        bindings = []
        for key, ab in session.app.active_bindings.items():
            bindings.append({
                "key": key,
                "description": ab.binding.description,
            })
        errors = session.error_collector.drain()
        return json.dumps({
            "tree": tree_text,
            "focused_ref": focused_ref,
            "bindings": bindings,
            "errors": errors,
        })
    except Exception as e:
        return json.dumps({"error": str(e)})


async def textual_screenshot(
    session_id: str,
    format: str = "text",
) -> str:
    """Capture a visual screenshot of the running app."""
    try:
        session = _session_manager.get_session(session_id)
        if format == "svg":
            content = session.app.export_screenshot()
        else:
            # Text format: render screen to plain text via Rich Console
            width, height = session.app.size
            console = Console(
                width=width,
                height=height,
                file=io.StringIO(),
                force_terminal=True,
                color_system="truecolor",
                record=True,
                legacy_windows=False,
            )
            screen_render = session.app.screen._compositor.render_update(
                full=True,
                screen_stack=session.app._background_screens,
            )
            console.print(screen_render)
            content = console.export_text()
        return json.dumps({"content": content, "format": format})
    except Exception as e:
        return json.dumps({"error": str(e)})


async def textual_query(
    session_id: str,
    selector: str,
) -> str:
    """Query widgets matching a CSS selector."""
    try:
        session = _session_manager.get_session(session_id)
        results = []
        for widget in session.app.screen.query(selector):
            results.append({
                "type": type(widget).__name__,
                "id": widget.id,
                "classes": sorted(widget.classes),
                "properties": _extractor.extract(widget),
            })
        return json.dumps({"widgets": results})
    except Exception as e:
        return json.dumps({"error": str(e)})


async def textual_get_screen_stack(session_id: str) -> str:
    """Get the current screen stack."""
    try:
        session = _session_manager.get_session(session_id)
        screens = []
        for screen in session.app.screen_stack:
            screens.append({
                "name": screen.name or type(screen).__name__,
                "type": type(screen).__name__,
                "is_modal": screen.is_modal,
            })
        return json.dumps({"screens": screens})
    except Exception as e:
        return json.dumps({"error": str(e)})


def _register_observation_tools(server: FastMCP) -> None:
    """Register observation tools on the FastMCP server."""
    ro = ToolAnnotations(readOnlyHint=True, destructiveHint=False, openWorldHint=False)
    server.tool(
        name="textual_snapshot",
        description="Snapshot current app state — widget tree with refs, focus, bindings, errors.",
        annotations=ro,
    )(textual_snapshot)
    server.tool(
        name="textual_screenshot",
        description="Capture a visual screenshot of the running app (text or SVG).",
        annotations=ro,
    )(textual_screenshot)
    server.tool(
        name="textual_query",
        description="Query widgets matching a CSS selector with their properties.",
        annotations=ro,
    )(textual_query)
    server.tool(
        name="textual_get_screen_stack",
        description="Get the current screen stack with modal indicators.",
        annotations=ro,
    )(textual_get_screen_stack)


_register_observation_tools(mcp)


# --- Assertion / wait tools ---


def _register_assertion_tools(server: FastMCP, manager: SessionManager) -> None:
    """Register assertion/wait tools on the given FastMCP server."""

    @server.tool(
        name="textual_wait_for",
        description="Wait for a condition: idle, animation, workers, or widget.",
        annotations=ToolAnnotations(
            readOnlyHint=True, destructiveHint=False, openWorldHint=False
        ),
    )
    async def textual_wait_for(
        session_id: str,
        condition: str,
        selector: str | None = None,
        text: str | None = None,
        timeout: float = 5.0,
    ) -> str:
        import asyncio
        import time

        try:
            session = manager.get_session(session_id)
            start = time.time()
            met = False

            if condition == "idle":
                await session.pilot.pause()
                met = True
            elif condition == "animation":
                await session.pilot.wait_for_animation()
                met = True
            elif condition == "workers":
                deadline = start + timeout
                while time.time() < deadline:
                    active = [w for w in session.app.workers if not w.is_finished]
                    if not active:
                        met = True
                        break
                    await asyncio.sleep(0.1)
            elif condition == "widget" and selector:
                deadline = start + timeout
                while time.time() < deadline:
                    results = session.app.screen.query(selector)
                    if len(results) > 0:
                        met = True
                        break
                    await asyncio.sleep(0.1)

            elapsed = time.time() - start
            return json.dumps({"met": met, "elapsed": round(elapsed, 3)})
        except Exception as e:
            return json.dumps({"error": str(e)})

    @server.tool(
        name="textual_check_errors",
        description="Check for collected errors and app exceptions.",
        annotations=ToolAnnotations(
            readOnlyHint=True, destructiveHint=False, openWorldHint=False
        ),
    )
    async def textual_check_errors(session_id: str, clear: bool = True) -> str:
        try:
            session = manager.get_session(session_id)
            app_exception = session.error_collector.check_app_exception(session.app)
            if clear:
                errors = session.error_collector.drain()
            else:
                errors = session.error_collector.errors.copy()
            return json.dumps({"errors": errors, "app_exception": app_exception})
        except Exception as e:
            return json.dumps({"error": str(e)})


_register_assertion_tools(mcp, _session_manager)


def main():
    mcp.run()


if __name__ == "__main__":
    main()
