"""MCP tool: vm_viewport_action — zero-image-token screen observation for vm-pilot.

Provider hierarchy (first available wins):

1. Claude CLI — if ``uv`` and ``claude`` are both on PATH.
2. OpenRouter — if ``OPENROUTER_API_KEY`` env var is set; delegates to
   ``scripts/vision_query.py`` subprocess.
3. Inline fallback — returns a text-only response; zero image tokens.

On JSON parse failure from providers 1 or 2, ``provider`` is set to
``'inline_after_parse_failure'`` to surface the failure mode transparently.
FastMCP ``FileSystemProvider`` auto-discovers this tool — no manual registration
in ``server.py`` is needed beyond importing this module.
"""

from __future__ import annotations

import io
import json
import logging
import os
import shutil
import subprocess
import tempfile
from pathlib import Path

import anyio
from fastmcp.tools import tool
from PIL import Image as PilImage

from mcp_vm_blackbox.models import ScaleInfo, ScreenState, SoMMarksLegend, ViewportResponse, VisionProvider
from mcp_vm_blackbox.vm_inspection import _screenshot_via_api, _to_jpeg
from mcp_vm_blackbox.vm_interaction import update_marks_cache

_log = logging.getLogger(__name__)

_SCRIPTS_DIR = Path(__file__).resolve().parent.parent.parent / "scripts"


# ---------------------------------------------------------------------------
# Provider availability checks
# ---------------------------------------------------------------------------


def _can_use_claude_cli() -> bool:
    """Return True when both ``uv`` and ``claude`` are on PATH.

    Returns:
        True if both executables are found via ``shutil.which``.
    """
    return shutil.which("uv") is not None and shutil.which("claude") is not None


def _can_use_openrouter() -> bool:
    """Return True when ``OPENROUTER_API_KEY`` is present in the environment.

    Returns:
        True if the environment variable is set to a non-empty string.
    """
    return bool(os.environ.get("OPENROUTER_API_KEY"))


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_mcp_config() -> dict[str, object]:
    """Return an MCP server config dict for a vm-blackbox subprocess.

    Returns:
        Dict suitable for writing as the ``--mcp-config`` JSON file passed
        to the ``claude`` CLI.
    """
    return {"mcpServers": {"vm-blackbox": {"command": "uv", "args": ["run", "mcp-vm-blackbox"]}}}


async def _take_screenshot_to_tempfile(
    vm_name: str, apply_som: bool = False
) -> tuple[Path, ScaleInfo, SoMMarksLegend | None] | None:
    """Capture a screenshot to a temporary JPEG file, optionally with SoM overlay.

    Captures a PNG via the vboxapi, converts it to JPEG using :func:`_to_jpeg`,
    and when *apply_som* is ``True`` attempts to detect UI elements and stamp
    Set-of-Marks labels onto the image before saving.

    Args:
        vm_name: VirtualBox VM name to screenshot.
        apply_som: When ``True``, detect UI elements and stamp SoM labels.
            Falls back to the unmarked JPEG when detection yields no elements.

    Returns:
        A ``(path, scale_info, marks_legend)`` tuple on success, where
        ``marks_legend`` is ``None`` when *apply_som* is ``False`` or no
        elements were detected.  Returns ``None`` when the screenshot capture
        itself failed.
    """
    tmp_dir = Path(tempfile.mkdtemp(prefix="viewport-"))
    out_path = tmp_dir / f"{vm_name}-viewport.png"
    result = _screenshot_via_api(vm_name, out_path)
    if result is None or result.error is not None:
        # Clean up empty temp dir on failure
        await anyio.Path(out_path).unlink(missing_ok=True)
        await anyio.Path(tmp_dir).rmdir()
        return None

    # Read PNG bytes and downscale to JPEG for context delivery
    png_bytes = out_path.read_bytes()
    jpeg_bytes, scale_info = _to_jpeg(png_bytes)

    marks_legend: SoMMarksLegend | None = None

    if apply_som:
        from mcp_vm_blackbox.som_overlay import detect_elements, stamp_marks  # noqa: PLC0415

        elements = await detect_elements(vm_name=vm_name, target="local")
        if elements:
            with PilImage.open(io.BytesIO(jpeg_bytes)) as jpeg_img:
                stamped_img, legend_draft = stamp_marks(jpeg_img.copy(), elements, scale_info)
                buf = io.BytesIO()
                stamped_img.convert("RGB").save(buf, format="JPEG", quality=85)
                jpeg_bytes = buf.getvalue()

            marks_legend = legend_draft.model_copy(update={"vm_name": vm_name})
            update_marks_cache(vm_name, marks_legend)
            _log.debug("viewport: SoM stamped %d marks for %s", len(marks_legend.marks), vm_name)
        else:
            _log.debug("viewport: SoM requested but no elements detected for %s", vm_name)

    # Overwrite temp file path with JPEG (callers get a single file path)
    jpeg_path = tmp_dir / f"{vm_name}-viewport.jpg"
    jpeg_path.write_bytes(jpeg_bytes)
    out_path.unlink(missing_ok=True)

    return jpeg_path, scale_info, marks_legend


# ---------------------------------------------------------------------------
# Provider implementations
# ---------------------------------------------------------------------------


def _provider_claude_cli(
    vm_name: str, instruction: str, max_turns: int, marks_legend: SoMMarksLegend | None = None
) -> ViewportResponse:
    """Spawn a Claude CLI process with vm-blackbox MCP access.

    Writes a temporary MCP config pointing at the vm-blackbox server, then
    runs ``claude --bare -p --output-format json``.  On JSON parse failure the
    response carries ``provider='inline_after_parse_failure'``.

    When *marks_legend* is provided, a SoM reference is prepended to the
    system prompt so the micro-agent knows to reference elements by mark number.

    Args:
        vm_name: VirtualBox VM name to observe.
        instruction: Natural-language task for the viewport micro-agent.
        max_turns: Maximum interaction turns allowed for the claude process.
        marks_legend: SoMMarksLegend from the SoM overlay step, or ``None``
            when SoM was not applied.

    Returns:
        ViewportResponse with ``provider='claude'`` on success, or
        ``provider='inline_after_parse_failure'`` on any failure.
    """
    with tempfile.NamedTemporaryFile(
        encoding="utf-8", mode="w", suffix=".json", prefix="vm-blackbox-mcp-", delete=False
    ) as tmp:
        json.dump(_make_mcp_config(), tmp)
        mcp_config_path = tmp.name

    try:
        som_prefix = ""
        if marks_legend is not None:
            som_prefix = "Elements labeled with numbered marks [1], [2], etc. Reference by mark number. "
        system_prompt = (
            f"{som_prefix}"
            f"You are a viewport micro-agent for VM '{vm_name}'. "
            "Use vm_screenshot to observe the current screen, then return ONLY a single "
            "valid JSON object — no markdown, no commentary — with this structure: "
            '{"action_taken": "<what you observed or did>", '
            '"screen_state": {"description": "<screen summary>", '
            '"visible_elements": [], "warnings": [], "errors": []}, '
            '"provider": "claude"}'
        )
        cmd = [
            "claude",
            "--bare",
            "-p",
            "--mcp-config",
            mcp_config_path,
            "--allowedTools",
            "mcp__vm-blackbox__vm_screenshot,mcp__vm-blackbox__vm_screenshot_api",
            "--max-turns",
            str(max_turns),
            "--model",
            "claude-haiku-4-5",
            "--output-format",
            "json",
            "--system-prompt",
            system_prompt,
            f"Observe VM '{vm_name}' and: {instruction}",
        ]
        result = subprocess.run(cmd, check=False, capture_output=True, text=True, timeout=60)

        if result.returncode != 0:
            return ViewportResponse(
                provider=VisionProvider.INLINE_AFTER_PARSE_FAILURE,
                action_taken=f"claude CLI exited with code {result.returncode}",
                screen_state=ScreenState(
                    description=(result.stderr or "no stderr output")[:500],
                    errors=[f"claude CLI exit code {result.returncode}"],
                ),
            )

        # claude --output-format json may wrap payload in {"result": ...}
        try:
            raw: object = json.loads(result.stdout)
            if isinstance(raw, dict) and "result" in raw:
                raw = raw["result"]
            if isinstance(raw, str):
                raw = json.loads(raw)
            return ViewportResponse.model_validate(raw)
        except Exception:  # noqa: BLE001
            return ViewportResponse(
                provider=VisionProvider.INLINE_AFTER_PARSE_FAILURE,
                action_taken="claude CLI returned non-JSON output",
                screen_state=ScreenState(
                    description=(result.stdout or "no output")[:500],
                    warnings=["JSON parse failed; raw CLI output is in description"],
                ),
            )
    finally:
        Path(mcp_config_path).unlink(missing_ok=True)


async def _provider_openrouter(vm_name: str, instruction: str) -> ViewportResponse:
    """Call ``scripts/vision_query.py`` as a subprocess for OpenRouter analysis.

    Takes a screenshot first, then passes the JPEG path and instruction to
    the PEP 723 vision script.  Cleans up the temporary screenshot on exit.

    Args:
        vm_name: VirtualBox VM name.
        instruction: Natural-language observation task.

    Returns:
        ViewportResponse with ``provider='openrouter'`` on success, or
        ``provider='inline_after_parse_failure'`` on any failure.
    """
    script_path = _SCRIPTS_DIR / "vision_query.py"
    if not script_path.exists():
        return ViewportResponse(
            provider=VisionProvider.INLINE_AFTER_PARSE_FAILURE,
            action_taken="OpenRouter provider unavailable: vision_query.py not found",
            screen_state=ScreenState(
                description=f"Script not found at {script_path}",
                errors=["vision_query.py missing; cannot use OpenRouter provider"],
            ),
        )

    capture = await _take_screenshot_to_tempfile(vm_name)
    if capture is None:
        return ViewportResponse(
            provider=VisionProvider.INLINE_AFTER_PARSE_FAILURE,
            action_taken="OpenRouter provider: screenshot capture failed",
            screen_state=ScreenState(
                description="Could not capture a screenshot for OpenRouter analysis",
                errors=["screenshot capture failed"],
            ),
        )

    screenshot_path, _scale_info, _marks_legend = capture
    try:
        cmd = ["uv", "run", str(script_path), str(screenshot_path), instruction]
        result = await anyio.run_process(cmd, check=False)
        stdout = result.stdout.decode(errors="replace")
        stderr = result.stderr.decode(errors="replace")

        if result.returncode != 0:
            return ViewportResponse(
                provider=VisionProvider.INLINE_AFTER_PARSE_FAILURE,
                action_taken="OpenRouter provider: vision_query.py failed",
                screen_state=ScreenState(
                    description=(stderr or "no stderr")[:500], errors=[f"vision_query.py exit code {result.returncode}"]
                ),
            )

        try:
            data = json.loads(stdout)
            return ViewportResponse.model_validate(data)
        except Exception:  # noqa: BLE001
            return ViewportResponse(
                provider=VisionProvider.INLINE_AFTER_PARSE_FAILURE,
                action_taken="OpenRouter provider: JSON parse failed",
                screen_state=ScreenState(
                    description=(stdout or "no output")[:500],
                    warnings=["JSON parse failed; raw output is in description"],
                ),
            )
    finally:
        await anyio.Path(screenshot_path).unlink(missing_ok=True)
        await anyio.Path(screenshot_path.parent).rmdir()


def _provider_inline_fallback(vm_name: str) -> ViewportResponse:
    """Return a text-only fallback when no vision provider is available.

    No image bytes enter the pilot's context — zero image tokens.

    Args:
        vm_name: VirtualBox VM name (used for descriptive output only).

    Returns:
        ViewportResponse with ``provider='inline_fallback'``.
    """
    return ViewportResponse(
        provider=VisionProvider.INLINE_FALLBACK,
        action_taken="screenshot captured inline",
        screen_state=ScreenState(description="Inline screenshot (no vision provider available)"),
    )


# ---------------------------------------------------------------------------
# MCP tool
# ---------------------------------------------------------------------------


@tool
async def vm_viewport_action(vm_name: str, instruction: str, max_turns: int = 5) -> ViewportResponse:
    """Observe the current VM screen and optionally perform one interaction.

    Returns a structured text description with element coordinates — zero image
    tokens in the pilot's context for routine screen observations.

    Provider hierarchy (first available wins):

    1. **Claude CLI** — if ``uv`` and ``claude`` are both on PATH.
    2. **OpenRouter** — if ``OPENROUTER_API_KEY`` env var is set.
    3. **Inline fallback** — text-only response, no image tokens.

    On JSON parse failure from providers 1 or 2, ``provider`` is set to
    ``'inline_after_parse_failure'`` to signal the failure mode clearly.

    Args:
        vm_name: VirtualBox VM name to observe.
        instruction: Natural-language instruction for the viewport micro-agent.
            Examples: ``"describe the current state"``,
            ``"find the Next button and return its coordinates"``.
        max_turns: Maximum turns for the Claude CLI provider (default 5).

    Returns:
        ViewportResponse with structured screen state and provider metadata.
        Includes ``scale`` and ``marks_legend`` when a SoM overlay was applied.
    """
    capture = await _take_screenshot_to_tempfile(vm_name, apply_som=True)

    scale: float | None = None
    marks_legend: SoMMarksLegend | None = None

    if capture is not None:
        _screenshot_path, scale_info, marks_legend = capture
        # Compute uniform scale factor for ViewportResponse (use the minimum of x/y)
        scale = min(scale_info.scale_x, scale_info.scale_y)

    if _can_use_claude_cli():
        response = _provider_claude_cli(vm_name, instruction, max_turns, marks_legend=marks_legend)
        return response.model_copy(update={"scale": scale, "marks_legend": marks_legend})
    if _can_use_openrouter():
        response = await _provider_openrouter(vm_name, instruction)
        return response.model_copy(update={"scale": scale, "marks_legend": marks_legend})
    response = _provider_inline_fallback(vm_name)
    return response.model_copy(update={"scale": scale, "marks_legend": marks_legend})
