"""Tool definitions — schemas (Anthropic API format) paired with execution functions."""

from __future__ import annotations

import json
import os
import time
import uuid
from dataclasses import dataclass
from typing import Any, Callable

import requests
from parallel import Parallel

from .utils import ToolResult, resize_and_encode_png


@dataclass
class Tool:
    """A tool definition paired with its optional execution function.

    *execute* is ``None`` for dispatch-only tools (handled by the orchestrator).
    """

    schema: dict[str, Any]
    execute: Callable[..., ToolResult] | None = None


def run_python_on_vm(env: Any, code: str) -> ToolResult:
    """Execute Python code on the VM via the /execute endpoint."""
    command = ["python3", "-c", code]
    payload = json.dumps({"command": command, "shell": False})
    headers = {"Content-Type": "application/json"}
    http_server = f"http://{env.vm_ip}:{env.server_port}"

    try:
        response = requests.post(
            http_server + "/execute",
            headers=headers,
            data=payload,
            timeout=120,
        )
        if response.status_code != 200:
            return ToolResult(error=f"HTTP {response.status_code}: {response.text}")
        result = response.json()
    except requests.exceptions.ReadTimeout:
        return ToolResult(error="Python script timed out")
    except Exception as e:
        return ToolResult(error=f"Failed to execute Python script: {e}")

    rc = result.get("returncode", -1)
    output = result.get("output", "")
    error = result.get("error", "")
    if rc != 0:
        return ToolResult(error=error or f"Exit code {rc}", output=output)
    return ToolResult(output=output or "(no output)")


def run_bash_on_vm(env: Any, script: str, timeout: int = 30) -> ToolResult:
    """Execute a bash script on the VM via the /execute endpoint (shell mode).

    Uses the /execute endpoint available on all AMIs rather than
    /run_bash_script which requires a custom server build.
    """
    payload = json.dumps({"command": script, "shell": True})
    headers = {"Content-Type": "application/json"}
    http_server = f"http://{env.vm_ip}:{env.server_port}"

    try:
        response = requests.post(
            http_server + "/execute",
            headers=headers,
            data=payload,
            timeout=timeout + 30,
        )
        if response.status_code != 200:
            return ToolResult(error=f"HTTP {response.status_code}: {response.text}")
        result = response.json()
    except requests.exceptions.ReadTimeout:
        return ToolResult(error=f"Script timed out after {timeout + 30}s")
    except Exception as e:
        return ToolResult(error=f"Failed to execute bash script: {e}")

    rc = result.get("returncode", -1)
    output = result.get("output", "")
    error = result.get("error", "")
    if rc != 0:
        return ToolResult(error=error or f"Exit code {rc}", output=output)
    return ToolResult(output=output or "(no output)")


def take_screenshot_on_vm(env: Any, timeout: int = 15) -> tuple[ToolResult, bytes | None]:
    """Capture a screenshot on the VM and return (ToolResult, raw_png_bytes).

    Two-call flow: /execute runs gnome-screenshot, /file retrieves the PNG bytes.
    DISPLAY=:0 because only X0 is bound on the OSWorld AMI. Uses a per-call UUID
    filename to avoid collisions between concurrent captures.
    """
    http_server = f"http://{env.vm_ip}:{env.server_port}"
    remote_path = f"/tmp/pointer_screenshot_{uuid.uuid4().hex}.png"

    capture = run_bash_on_vm(
        env,
        f"DISPLAY=:0 gnome-screenshot -f {remote_path}",
        timeout=timeout,
    )
    if capture.error:
        return ToolResult(error=f"screenshot capture failed: {capture.error}"), None

    try:
        resp = requests.post(
            http_server + "/file",
            files={"file_path": (None, remote_path)},
            timeout=timeout,
        )
        if resp.status_code != 200:
            return (
                ToolResult(error=f"HTTP {resp.status_code} from /file: {resp.text[:200]}"),
                None,
            )
        png_bytes = resp.content
    except Exception as e:
        return ToolResult(error=f"Failed to download screenshot: {e}"), None

    b64 = resize_and_encode_png(png_bytes)
    if not b64:
        return ToolResult(error="screenshot downloaded but was empty"), None
    return ToolResult(
        output=f"The screenshot was saved at {remote_path} on the virtual machine.",
        base64_data=b64,
        media_type="image/png",
    ), png_bytes


_BINARY_EXTENSIONS = frozenset(
    {
        ".png",
        ".jpg",
        ".jpeg",
        ".gif",
        ".bmp",
        ".tiff",
        ".tif",
        ".ico",
        ".webp",
        ".svg",
        ".pdf",
        ".doc",
        ".docx",
        ".xls",
        ".xlsx",
        ".ppt",
        ".pptx",
        ".zip",
        ".tar",
        ".gz",
        ".bz2",
        ".xz",
        ".7z",
        ".rar",
        ".mp3",
        ".mp4",
        ".avi",
        ".mov",
        ".mkv",
        ".wav",
        ".flac",
        ".ogg",
        ".exe",
        ".dll",
        ".so",
        ".dylib",
        ".o",
        ".pyc",
        ".class",
        ".xcf",
        ".psd",
        ".ai",
        ".eps",
        ".ods",
        ".odt",
        ".odp",
        ".odb",
        ".sqlite",
        ".db",
    }
)
_MAX_FILE_READ_CHARS = 50_000

_IMAGE_EXTENSIONS: dict[str, str] = {
    ".png": "image/png",
    ".jpg": "image/jpeg",
    ".jpeg": "image/jpeg",
    ".gif": "image/gif",
    ".webp": "image/webp",
}

_UNSUPPORTED_BINARY_MESSAGE = (
    "This is an unsupported file type. If you want to read the contents "
    "of the file, use code tools to view parts of the file"
)


def _download_file_from_vm(env: Any, path: str, timeout: int = 15) -> bytes | None:
    """Download raw bytes of a file from the VM via the /file endpoint."""
    http_server = f"http://{env.vm_ip}:{env.server_port}"
    try:
        resp = requests.post(
            http_server + "/file",
            files={"file_path": (None, path)},
            timeout=timeout,
        )
        if resp.status_code != 200:
            return None
        return resp.content
    except Exception:
        return None


def read_file_on_vm(env: Any, path: str) -> ToolResult:
    import base64
    import os

    ext = os.path.splitext(path)[1].lower()

    # PDF -> document content block
    if ext == ".pdf":
        raw_bytes = _download_file_from_vm(env, path)
        if raw_bytes is None:
            return ToolResult(error=f"Failed to download file: {path}")
        b64 = base64.b64encode(raw_bytes).decode("utf-8")
        return ToolResult(
            output=f"PDF file: {path}",
            base64_data=b64,
            media_type="application/pdf",
        )

    # Image -> image content block with correct media type
    if ext in _IMAGE_EXTENSIONS:
        raw_bytes = _download_file_from_vm(env, path)
        if raw_bytes is None:
            return ToolResult(error=f"Failed to download file: {path}")
        b64 = base64.b64encode(raw_bytes).decode("utf-8")
        return ToolResult(
            output=f"Image file: {path}",
            base64_data=b64,
            media_type=_IMAGE_EXTENSIONS[ext],
        )

    # Other binary files -> unsupported
    if ext in _BINARY_EXTENSIONS:
        return ToolResult(output=_UNSUPPORTED_BINARY_MESSAGE)

    # Text files -> cat and return
    result = run_bash_on_vm(env, f"cat {path!r}", timeout=15)
    if result.error:
        return result
    text = result.output or ""
    if len(text) > _MAX_FILE_READ_CHARS:
        text = (
            text[:_MAX_FILE_READ_CHARS]
            + f"\n... [truncated — file is {len(text)} chars, showing first {_MAX_FILE_READ_CHARS}]"
        )
        return ToolResult(output=text)
    return result


def write_file_on_vm(env: Any, path: str, content: str) -> ToolResult:
    code = (
        f"import os\n"
        f"os.makedirs(os.path.dirname({path!r}) or '.', exist_ok=True)\n"
        f"with open({path!r}, 'w') as f:\n"
        f"    f.write({content!r})\n"
        f"print('File written successfully')"
    )
    return run_python_on_vm(env, code)


def run_in_background_on_vm(
    env: Any,
    command: str,
    poll_interval: int = 60,
    max_polls: int = 10,
) -> ToolResult:
    """Launch a command in the background on the VM and poll until completion.

    Avoids the /execute endpoint's 120s subprocess timeout by using nohup
    with output redirection, then polling for process completion from the
    agent side with time.sleep between HTTP calls.
    """
    job_id = uuid.uuid4().hex[:12]
    log_file = f"/tmp/bg_{job_id}.log"

    # Wrap in `sh -c` so env-var prefixes (FOO=bar cmd), pipelines, &&, and
    # redirection inside `command` are parsed by a shell rather than passed
    # as literal argv to nohup.
    escaped = command.replace("'", "'\\''")
    launch_script = f"nohup sh -c '{escaped}' > {log_file} 2>&1 & echo $!"
    launch_result = run_bash_on_vm(env, launch_script, timeout=10)
    if launch_result.error and not launch_result.output:
        return launch_result

    pid = (launch_result.output or "").strip().split("\n")[-1].strip()
    if not pid.isdigit():
        return ToolResult(error=f"Failed to parse PID from launch output: {launch_result.output}")

    timed_out = True
    for i in range(max_polls):
        time.sleep(poll_interval)
        poll_result = run_bash_on_vm(
            env,
            f"kill -0 {pid} 2>/dev/null && echo RUNNING || echo FINISHED",
            timeout=10,
        )
        status = (poll_result.output or "").strip()
        if "FINISHED" in status:
            timed_out = False
            break

    read_result = run_bash_on_vm(env, f"cat {log_file}", timeout=30)
    output = read_result.output or read_result.error or "(no output)"

    run_bash_on_vm(env, f"rm -f {log_file}", timeout=5)

    if timed_out:
        output = f"WARNING: Process may still be running (pid={pid}). Output so far:\n{output}"

    return ToolResult(output=output)


# ---------------------------------------------------------------------------
# Chrome DevTools Protocol (CDP) helpers
# ---------------------------------------------------------------------------
#
# Architectural note: Chrome binds --remote-debugging-port to 127.0.0.1 even
# with --remote-debugging-address=0.0.0.0 (silently ignored on modern Chrome
# for security). So CDP cannot be reached from the agent host — we ship Python
# snippets into the VM via /execute, where they talk to 127.0.0.1:9222 locally.
#
# We use pychrome (pure-Python CDP client, ~90KB) rather than Playwright to
# avoid shipping a browser bundle onto every VM.

# Runs inside the VM's python3. Sets a `browser` handle over CDP, execs the
# user's snippet, prints the json-serialized `result` so the agent can read it.
# We keep the wrapper minimal — the pychrome background thread produces a
# harmless traceback on interpreter exit which we squelch by dropping stderr
# on success so clean output dominates the ToolResult.
_CDP_SNIPPET_WRAPPER = """\
import json, sys, traceback
try:
    import pychrome
    browser = pychrome.Browser(url="http://127.0.0.1:{cdp_port}")
    result = None
{user_code}
    print("__CDP_RESULT__" + json.dumps(result, default=str))
except Exception:
    traceback.print_exc()
    sys.exit(1)
"""


def _cdp_version_probe(env: Any) -> ToolResult:
    """Run a tiny VM-side python snippet that hits http://127.0.0.1:9222/json/version.

    Returns output=<version-json> on success, error=<reason> on failure.
    Avoids shelling out to curl (not present on the OSWorld AMI).
    """
    probe = (
        "import json, urllib.request, urllib.error\n"
        "try:\n"
        "    d = urllib.request.urlopen("
        f"'http://127.0.0.1:{env.chromium_port}/json/version', timeout=3).read().decode()\n"
        "    print(d)\n"
        "except Exception as e:\n"
        "    import sys; sys.stderr.write(str(e)); sys.exit(1)\n"
    )
    return run_python_on_vm(env, probe)


def start_chrome_with_cdp(env: Any) -> ToolResult:
    """Ensure Chrome is running on the VM with CDP reachable on env.chromium_port.

    Flow:
    1. If /json/version responds -> already connected, return ok.
    2. If Chrome is running but CDP unreachable -> refuse (the user launched
       Chrome without --remote-debugging-port, and relaunching would lose state).
    3. Otherwise: ensure pychrome is installed, launch Chrome with CDP, poll
       /json/version for up to ~10s, then return ok.
    """
    # Idempotent — pip no-ops if already installed. Capped under /execute's 120s.
    install = run_in_background_on_vm(env=env, command="pip3 install pychrome", poll_interval=60, max_polls=10)
    if install.error and not install.output:
        return ToolResult(error=f"pychrome install failed: {install.error}")

    version = _cdp_version_probe(env)
    if not version.error:
        return ToolResult(output=f"CDP already reachable on :{env.chromium_port}\n{version.output or ''}")

    running = run_bash_on_vm(env, "pgrep -f '/opt/google/chrome|/usr/bin/chromium' | head -1", timeout=10)
    if (running.output or "").strip():
        return ToolResult(
            error=(
                "Chrome is running but not listening for CDP. "
                "Close Chrome first (e.g. via close_application), then retry so "
                "we can relaunch it with --remote-debugging-port enabled."
            )
        )

    http_server = f"http://{env.vm_ip}:{env.server_port}"
    launch_payload = json.dumps(
        {
            "command": [
                "google-chrome",
                f"--remote-debugging-port={env.chromium_port}",
            ],
            "shell": False,
        }
    )
    try:
        resp = requests.post(
            http_server + "/setup/launch",
            headers={"Content-Type": "application/json"},
            data=launch_payload,
            timeout=15,
        )
        if resp.status_code != 200:
            return ToolResult(error=f"/setup/launch HTTP {resp.status_code}: {resp.text[:200]}")
    except Exception as e:
        return ToolResult(error=f"Failed to launch Chrome: {e}")

    # Poll for CDP readiness — Chrome takes 1-3s to bind the port.
    for _ in range(10):
        time.sleep(1)
        check = _cdp_version_probe(env)
        if not check.error:
            return ToolResult(output=f"Chrome launched with CDP on :{env.chromium_port}")
    return ToolResult(error=f"Chrome launched but CDP did not come up on :{env.chromium_port} within 10s")


_CDP_RESULT_SENTINEL = "__CDP_RESULT__"


def run_cdp_on_vm(env: Any, code: str) -> ToolResult:
    """Execute a pychrome snippet on the VM against Chrome's CDP endpoint.

    The snippet runs with `browser = pychrome.Browser(...)` pre-bound; the
    caller assigns their answer to `result` (json-serialized in the return).
    Chrome state (tabs, DOM, JS runtime, cookies) persists across invocations.

    Per the user's rule: a non-empty stdout is authoritative — stderr is only
    reported as an error when stdout is empty, to tolerate pychrome's noisy
    background-thread teardown.
    """
    ensured = start_chrome_with_cdp(env)
    if ensured.error:
        return ensured

    indented = "\n".join("    " + line for line in (code or "").splitlines()) or "    pass"
    wrapped = _CDP_SNIPPET_WRAPPER.format(cdp_port=env.chromium_port, user_code=indented)
    raw = run_python_on_vm(env, wrapped)

    output = (raw.output or "").strip()
    if output:
        # Split off the sentinel line — anything before it was user `print()` noise.
        result_payload = ""
        user_stdout_lines: list[str] = []
        for line in output.splitlines():
            if line.startswith(_CDP_RESULT_SENTINEL):
                result_payload = line[len(_CDP_RESULT_SENTINEL) :]
            else:
                user_stdout_lines.append(line)
        pieces = []
        if user_stdout_lines:
            pieces.append("\n".join(user_stdout_lines))
        if result_payload:
            pieces.append(f"result: {result_payload}")
        return ToolResult(output="\n".join(pieces) or "(no output)")

    # Stdout empty -> treat stderr as real failure.
    return ToolResult(error=raw.error or "CDP snippet produced no output")


def log_to_bash_history(env: Any, entry: str) -> None:
    """Append a command entry to ~/.bash_history on the VM (best-effort)."""
    if not entry or not entry.strip():
        return

    lines = [ln.strip() for ln in entry.strip().splitlines() if ln.strip()]
    flat = " ; ".join(lines)

    escaped = flat.replace("'", "'\\''")
    append_cmd = f"printf '%s\\n' '{escaped}' >> ~/.bash_history"

    payload = json.dumps({"command": append_cmd, "shell": True})
    headers = {"Content-Type": "application/json"}
    http_server = f"http://{env.vm_ip}:{env.server_port}"

    try:
        requests.post(http_server + "/execute", headers=headers, data=payload, timeout=5)
    except Exception:
        pass


def _get_parallel_client() -> Parallel:
    return Parallel(api_key=os.environ["PARALLEL_API_KEY"])


def web_search(query: str, max_results: int = 5) -> ToolResult:
    """Search the web via Parallel Search API."""
    max_results = min(max_results, 20)
    try:
        client = _get_parallel_client()
        response = client.search(
            search_queries=[query],
            objective=query,
            mode="basic",
            advanced_settings={"max_results": max_results},
        )
        results = [
            {
                "title": r.title or "",
                "href": r.url,
                "body": "\n".join(r.excerpts) if r.excerpts else "",
            }
            for r in response.results
        ]
        return ToolResult(output=json.dumps(results, indent=2, ensure_ascii=False))
    except Exception as e:
        return ToolResult(error=f"web_search failed: {e}")


def web_extract(urls: list[str], max_chars: int = 8000) -> ToolResult:
    """Extract content from URLs via Parallel Extract API."""
    urls = urls[:20]
    try:
        client = _get_parallel_client()
        response = client.extract(
            urls=urls,
            advanced_settings={
                "excerpt_settings": {"max_chars_per_result": max_chars},
            },
        )
        results = []
        for r in response.results:
            content = "\n".join(r.excerpts) if r.excerpts else ""
            results.append(
                {
                    "url": r.url,
                    "title": r.title or "",
                    "content": content[:max_chars],
                    "truncated": len(content) > max_chars,
                }
            )
        return ToolResult(output=json.dumps(results, indent=2, ensure_ascii=False))
    except Exception as e:
        return ToolResult(error=f"web_extract failed: {e}")


WEB_SEARCH = Tool(
    schema={
        "name": "web_search",
        "description": (
            "Search the web and return a JSON array of results. "
            "Each result has 'title', 'href', and 'body' (excerpt) fields. "
            "Use this to find information, verify capabilities, or locate "
            "documentation URLs before fetching them."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "The search query.",
                },
                "max_results": {
                    "type": "integer",
                    "description": "Maximum number of results to return (default 5, max 20).",
                },
            },
            "required": ["query"],
        },
    },
)

WEB_FETCH = Tool(
    schema={
        "name": "web_fetch",
        "description": (
            "Extract readable content from one or more web pages as clean "
            "markdown. Returns a JSON array with 'url', 'title', 'content', "
            "and 'truncated' fields per page. "
            "Use web_search first to find relevant URLs, then web_fetch to "
            "read their content in detail."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "urls": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "URLs to fetch (fetched in parallel).",
                },
                "max_chars": {
                    "type": "integer",
                    "description": "Max characters of extracted text per URL (default 8000).",
                },
            },
            "required": ["urls"],
        },
    },
)

# ---------------------------------------------------------------------------
# MARK: Feasibility gate tools
# ---------------------------------------------------------------------------

PROBE_BASH = Tool(
    schema={
        "name": "probe_bash",
        "description": (
            "Run a READ-ONLY bash command to inspect system state (ls, which, "
            "cat, grep, dpkg -l, apt list, file, stat, etc.). "
            "Do NOT use this to modify files, install packages, or change settings."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "script": {
                    "type": "string",
                    "description": "Read-only bash command.",
                },
                "timeout": {
                    "type": "integer",
                    "description": "Timeout in seconds (default 30).",
                },
            },
            "required": ["script"],
        },
    },
)

PROBE_PYTHON = Tool(
    schema={
        "name": "probe_python",
        "description": (
            "Run a READ-ONLY Python script to inspect system state. Do NOT use this to modify files or system state."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "code": {
                    "type": "string",
                    "description": "Read-only Python code.",
                },
            },
            "required": ["code"],
        },
    },
)

PROBE_CHROME_CDP = Tool(
    schema={
        "name": "probe_chrome_cdp",
        "description": (
            "Query Chrome state via CDP (READ-ONLY). Check open tabs, URLs, "
            "page content, extensions. Do NOT navigate to new pages or modify "
            "browser state. The snippet runs in Python on the VM with "
            "`browser = pychrome.Browser(...)` pre-bound."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "code": {
                    "type": "string",
                    "description": "Read-only Python CDP snippet.",
                },
            },
            "required": ["code"],
        },
    },
)

REPORT_FEASIBLE = Tool(
    schema={
        "name": "report_feasible",
        "description": ("Report that the task IS feasible on this VM. Call this exactly once with your reasoning."),
        "input_schema": {
            "type": "object",
            "properties": {
                "reason": {
                    "type": "string",
                    "description": "Why this task is feasible, citing evidence from probing.",
                },
            },
            "required": ["reason"],
        },
    },
)

REPORT_INFEASIBLE = Tool(
    schema={
        "name": "report_infeasible",
        "description": ("Report that the task is NOT feasible on this VM. Call this exactly once with your reasoning."),
        "input_schema": {
            "type": "object",
            "properties": {
                "reason": {
                    "type": "string",
                    "description": (
                        "Why this task is infeasible, citing the specific missing "
                        "capability or resource confirmed by probing."
                    ),
                },
            },
            "required": ["reason"],
        },
    },
)

# ---------------------------------------------------------------------------
# Planner tools
# ---------------------------------------------------------------------------

SUBMIT_PLAN = Tool(
    schema={
        "name": "submit_plan",
        "description": (
            "Submit the execution plan as an ordered list of milestones. "
            "Each milestone should describe an outcome / state the world should "
            "reach, not step-by-step instructions. Let the task shape the count. "
            "Call this once you have gathered enough information."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "milestones": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "Ordered list of milestone descriptions.",
                },
                "infeasible_reason": {
                    "type": "string",
                    "description": (
                        "If you strongly suspect this task is infeasible "
                        "(requires hardware / software / a feature that doesn't "
                        "exist on this VM or in the world), set this to a brief "
                        "reason. Leave unset/empty if the task looks doable."
                    ),
                },
            },
            "required": ["milestones"],
        },
    },
)

# ---------------------------------------------------------------------------
# Executor tools
# ---------------------------------------------------------------------------

TASK_DONE = Tool(
    schema={
        "name": "done",
        "description": "Declare the overall task as successfully completed.",
        "input_schema": {
            "type": "object",
            "properties": {
                "summary": {
                    "type": "string",
                    "description": "Brief summary of what was accomplished.",
                },
            },
            "required": ["summary"],
        },
    },
)

TASK_FAIL = Tool(
    schema={
        "name": "fail",
        "description": (
            "Declare the task as infeasible. Use when the task "
            "cannot be completed due to missing features, impossible requirements, "
            "or unrecoverable errors."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "reason": {
                    "type": "string",
                    "description": "Why the task cannot be completed.",
                },
            },
            "required": ["reason"],
        },
    },
)

# ---------------------------------------------------------------------------
# MARK: Code action
# ---------------------------------------------------------------------------

RUN_PYTHON = Tool(
    schema={
        "name": "run_python",
        "description": "Execute a Python script on the VM. Returns stdout/stderr.",
        "input_schema": {
            "type": "object",
            "properties": {
                "code": {"type": "string", "description": "Python code to execute."},
            },
            "required": ["code"],
        },
    },
)

RUN_BASH = Tool(
    schema={
        "name": "run_bash",
        "description": "Execute a bash script on the VM. Returns stdout/stderr and exit code.",
        "input_schema": {
            "type": "object",
            "properties": {
                "script": {"type": "string", "description": "Bash script to execute."},
                "timeout": {
                    "type": "integer",
                    "description": "Timeout in seconds (default 30).",
                },
            },
            "required": ["script"],
        },
    },
)

READ_FILE = Tool(
    schema={
        "name": "read_file",
        "description": "Read the contents of a file on the VM.",
        "input_schema": {
            "type": "object",
            "properties": {
                "path": {
                    "type": "string",
                    "description": "Absolute file path on the VM.",
                },
            },
            "required": ["path"],
        },
    },
)

WRITE_FILE = Tool(
    schema={
        "name": "write_file",
        "description": "Write content to a file on the VM (creates or overwrites).",
        "input_schema": {
            "type": "object",
            "properties": {
                "path": {
                    "type": "string",
                    "description": "Absolute file path on the VM.",
                },
                "content": {"type": "string", "description": "File content to write."},
            },
            "required": ["path", "content"],
        },
    },
)

INSTALL_PACKAGE = Tool(
    schema={
        "name": "install_package",
        "description": (
            "Install a package via a long-running command (e.g. apt-get, pip). "
            "The command runs in the background to avoid timeouts. "
            "Optionally provide a verification command to confirm success."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "command": {
                    "type": "string",
                    "description": "The install command (e.g. 'sudo apt-get install -y firefox').",
                },
                "verification": {
                    "type": "string",
                    "description": "Optional command to verify installation (e.g. 'which firefox').",
                },
            },
            "required": ["command"],
        },
    },
)

RUN_BACKGROUND = Tool(
    schema={
        "name": "run_background",
        "description": (
            "Run a long-running bash command in the background. "
            "Use for any command that might exceed the normal execution timeout "
            "(e.g. large downloads, compilations, data processing). "
            "The command runs asynchronously and is polled for completion."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "command": {
                    "type": "string",
                    "description": "The bash command to run in the background.",
                },
                "timeout": {
                    "type": "integer",
                    "description": "Maximum wait time in seconds (default 600, max 600).",
                },
            },
            "required": ["command"],
        },
    },
)

# ---------------------------------------------------------------------------
# MARK: GUI action
# ---------------------------------------------------------------------------

CLICK = Tool(
    schema={
        "name": "click",
        "description": (
            "Click at a screen coordinate. Coordinates are in a 1280x720 space; "
            "aim for the center of the element you want to click."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "coordinate": {
                    "type": "array",
                    "items": {"type": "integer"},
                    "minItems": 2,
                    "maxItems": 2,
                    "description": "[x, y] target coordinate in 1280x720 space.",
                },
                "num_clicks": {
                    "type": "integer",
                    "description": "Number of clicks (1=single, 2=double, 3=triple). Default 1.",
                },
                "button_type": {
                    "type": "string",
                    "enum": ["left", "middle", "right"],
                    "description": "Which mouse button to press. Default 'left'.",
                },
                "hold_keys": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "Keys to hold while clicking (e.g. ['ctrl'], ['shift']). Default [].",
                },
            },
            "required": ["coordinate"],
        },
    },
)

TYPE = Tool(
    schema={
        "name": "type",
        "description": (
            "Type text into a GUI field. Optionally click a coordinate first to focus "
            "the field. Uses clipboard paste for text containing non-ASCII characters."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "text": {
                    "type": "string",
                    "description": "The text to type.",
                },
                "coordinate": {
                    "type": "array",
                    "items": {"type": "integer"},
                    "minItems": 2,
                    "maxItems": 2,
                    "description": (
                        "Optional [x, y] in 1280x720 space to click before typing "
                        "(to focus the field). Omit if the field is already focused."
                    ),
                },
                "overwrite": {
                    "type": "boolean",
                    "description": "If true, select-all + delete existing text before typing. Default false.",
                },
                "enter": {
                    "type": "boolean",
                    "description": "If true, press Enter after typing. Default false.",
                },
            },
            "required": ["text"],
        },
    },
)

SCROLL = Tool(
    schema={
        "name": "scroll",
        "description": (
            "Scroll at a coordinate. Positive `clicks` scrolls up (or right if shift=true); "
            "negative scrolls down (or left)."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "coordinate": {
                    "type": "array",
                    "items": {"type": "integer"},
                    "minItems": 2,
                    "maxItems": 2,
                    "description": "[x, y] in 1280x720 space to scroll at.",
                },
                "clicks": {
                    "type": "integer",
                    "description": "Scroll amount; positive = up/right, negative = down/left.",
                },
                "shift": {
                    "type": "boolean",
                    "description": "If true, scroll horizontally instead of vertically. Default false.",
                },
            },
            "required": ["coordinate", "clicks"],
        },
    },
)

DRAG_AND_DROP = Tool(
    schema={
        "name": "drag_and_drop",
        "description": "Drag from one coordinate to another, optionally holding modifier keys.",
        "input_schema": {
            "type": "object",
            "properties": {
                "start_coordinate": {
                    "type": "array",
                    "items": {"type": "integer"},
                    "minItems": 2,
                    "maxItems": 2,
                    "description": "[x, y] where the drag begins.",
                },
                "end_coordinate": {
                    "type": "array",
                    "items": {"type": "integer"},
                    "minItems": 2,
                    "maxItems": 2,
                    "description": "[x, y] where the drag ends.",
                },
                "hold_keys": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "Keys to hold while dragging (e.g. ['shift']). Default [].",
                },
            },
            "required": ["start_coordinate", "end_coordinate"],
        },
    },
)

HOTKEY = Tool(
    schema={
        "name": "hotkey",
        "description": "Press a keyboard shortcut (e.g. ['ctrl', 'c'] or ['ctrl', 'shift', 't']).",
        "input_schema": {
            "type": "object",
            "properties": {
                "keys": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "List of keys to press together, in order (e.g. ['ctrl', 'c']).",
                },
            },
            "required": ["keys"],
        },
    },
)

HOLD_AND_PRESS = Tool(
    schema={
        "name": "hold_and_press",
        "description": (
            "Hold one set of keys while pressing another sequence. Useful for shortcuts that "
            "require modifiers held across multiple presses (e.g. hold shift, press several arrow keys)."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "hold_keys": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "Keys to hold down throughout (e.g. ['shift']).",
                },
                "press_keys": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "Keys to press one after another while hold_keys are held.",
                },
            },
            "required": ["hold_keys", "press_keys"],
        },
    },
)

WAIT = Tool(
    schema={
        "name": "wait",
        "description": (
            "Sleep for a specified duration to let the UI settle (e.g. after opening a file, "
            "waiting for a network response, or letting an animation finish)."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "seconds": {
                    "type": "number",
                    "description": "Number of seconds to sleep.",
                },
            },
            "required": ["seconds"],
        },
    },
)

OPEN = Tool(
    schema={
        "name": "open",
        "description": (
            "Open an application or file by name. Triggers the Activities overview "
            "(via the 'win' key), types the name, and presses Enter. Use this instead "
            "of clicking through the launcher or file manager manually."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "app_or_filename": {
                    "type": "string",
                    "description": "Name of the application or filename to open (e.g. 'Firefox', 'Files', 'report.txt').",
                },
            },
            "required": ["app_or_filename"],
        },
    },
)

SWITCH_APPLICATIONS = Tool(
    schema={
        "name": "switch_applications",
        "description": (
            "Switch focus to an already-open application by fuzzy-matching its window title. "
            "Uses wmctrl to raise and maximize the matched window. Prefer this over alt-tabbing "
            "when you know the target application name."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "app_code": {
                    "type": "string",
                    "description": "Name (or partial name) of the application to switch to — matched against open window titles via difflib.",
                },
            },
            "required": ["app_code"],
        },
    },
)

CLOSE_APPLICATION = Tool(
    schema={
        "name": "close_application",
        "description": (
            "Close an open application by fuzzy-matching its window title and sending a "
            "graceful close request via wmctrl. If the app shows a save-changes dialog, "
            "handle it with a follow-up `click` or `hotkey` in the next step."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "app_code": {
                    "type": "string",
                    "description": "Name (or partial name) of the application to close — matched against open window titles via difflib.",
                },
            },
            "required": ["app_code"],
        },
    },
)

LIST_WINDOWS = Tool(
    schema={
        "name": "list_windows",
        "description": (
            "List all open windows on the VM desktop via wmctrl. "
            "Returns window ID, desktop, WM_CLASS, and title for each window. "
            "Use to discover what applications are currently open before switching or closing."
        ),
        "input_schema": {"type": "object", "properties": {}, "required": []},
    },
)

SCREENSHOT = Tool(
    schema={
        "name": "screenshot",
        "description": (
            "Capture a fresh screenshot of the VM desktop and return it as an image. "
            "Use this to re-check desktop state after an action whose visual effect "
            "you need to confirm before proceeding."
        ),
        "input_schema": {"type": "object", "properties": {}, "required": []},
    },
)

CHROME_CDP = Tool(
    schema={
        "name": "chrome_cdp",
        "description": (
            "Drive Chrome on the VM via Chrome DevTools Protocol. The snippet you "
            "provide runs in Python on the VM with `browser = pychrome.Browser(...)` "
            "and `result = None` pre-bound — assign `result` to whatever you want "
            "returned (JSON-serialized). If Chrome is not running, the tool starts "
            "it with --remote-debugging-port enabled; if Chrome is running without "
            "CDP the tool will refuse, and you must close Chrome first.\n\n"
            "Chrome state (tabs, DOM, cookies, in-page JS globals) persists across "
            "calls, so you can set a variable in one call and read it in the next. "
            "pychrome usage reminders: `browser.list_tab()` returns all CDP "
            "targets — not all are user-visible tabs (background pages, iframes, "
            "service workers may appear). Filter to the tab you want by its URL "
            "before acting. Call `tab.start()` before any Domain methods, and "
            "`tab.Runtime.evaluate(expression=...)` for JS. `tab.Page.enable()` "
            "then `tab.Page.navigate(url=...)` to load a URL — wait for "
            "document.readyState == 'complete' before reading results. Prefer "
            "this over GUI clicks for Chrome tasks."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "code": {
                    "type": "string",
                    "description": (
                        "Python snippet to run on the VM. `browser` and `result` "
                        "are pre-bound; assign `result` to return a value."
                    ),
                },
            },
            "required": ["code"],
        },
    },
)

MARK_MILESTONE_DONE = Tool(
    schema={
        "name": "mark_milestone_done",
        "description": (
            "Mark the current milestone as completed and advance to the next one. "
            "Call this when you have finished the work described by the current "
            "milestone (marked with [>] in the plan). Returns the updated plan status."
        ),
        "input_schema": {"type": "object", "properties": {}, "required": []},
    },
)

UPDATE_MILESTONE = Tool(
    schema={
        "name": "update_milestone",
        "description": (
            "Update the description of a PENDING milestone. Cannot modify completed milestones. "
            "Use when you realize a milestone's target state needs to be refined or corrected "
            "based on what you've learned during execution."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "milestone_idx": {
                    "type": "integer",
                    "description": "1-based index of the milestone to update (matching the plan display).",
                },
                "new_description": {
                    "type": "string",
                    "description": "The new description for the milestone.",
                },
            },
            "required": ["milestone_idx", "new_description"],
        },
    },
)

ADD_MILESTONE = Tool(
    schema={
        "name": "add_milestone",
        "description": (
            "Insert a new milestone into the plan. If idx is provided, inserts at that "
            "1-based position (shifting later milestones down). If idx is omitted, appends "
            "at the end. Use when you discover additional work not anticipated by the original plan."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "description": {
                    "type": "string",
                    "description": "Description of the new milestone (target state, not instructions).",
                },
                "idx": {
                    "type": "integer",
                    "description": (
                        "Optional 1-based position to insert at. Cannot be before completed milestones. "
                        "If omitted, appends at the end."
                    ),
                },
            },
            "required": ["description"],
        },
    },
)

DELETE_MILESTONE = Tool(
    schema={
        "name": "delete_milestone",
        "description": (
            "Remove a PENDING milestone that is no longer needed. Cannot delete completed "
            "milestones or the current milestone. Use when a milestone turns out to be "
            "unnecessary or redundant."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "milestone_idx": {
                    "type": "integer",
                    "description": "1-based index of the milestone to delete.",
                },
            },
            "required": ["milestone_idx"],
        },
    },
)

SET_SUB_MILESTONES = Tool(
    schema={
        "name": "set_sub_milestones",
        "description": (
            "Set a list of action items for the current milestone. Replaces any previously "
            "set sub-milestones. Use to break down the current milestone into smaller, "
            "trackable steps."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "sub_milestones": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "Ordered list of sub-milestone descriptions for the current milestone.",
                },
            },
            "required": ["sub_milestones"],
        },
    },
)

MARK_SUB_MILESTONE_DONE = Tool(
    schema={
        "name": "mark_sub_milestone_done",
        "description": (
            "Mark a sub-milestone of the current milestone as completed. "
            "Use the 1-based index as displayed in the plan."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "sub_milestone_idx": {
                    "type": "integer",
                    "description": "1-based index of the sub-milestone to mark as done.",
                },
            },
            "required": ["sub_milestone_idx"],
        },
    },
)

# ---------------------------------------------------------------------------
# Verifier tools
# ---------------------------------------------------------------------------

SUBMIT_VERDICT = Tool(
    schema={
        "name": "submit_verdict",
        "description": (
            "Submit the final verification decision. Call exactly once after probing "
            "the VM state. Use 'done' only when at least one probe confirms the "
            "instruction is satisfied; 'failed' when the capability exists but the "
            "state does not match (must include diagnosis and suggestion); "
            "'infeasible' when the required capability genuinely does not exist on "
            "this VM, confirmed by probing."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "decision": {
                    "type": "string",
                    "enum": ["done", "failed", "infeasible"],
                    "description": "Final verdict for the task.",
                },
                "reasoning": {
                    "type": "string",
                    "description": (
                        "Brief (<=3 sentences) explanation citing the specific probe "
                        "outputs that drove the decision. Example: 'openpyxl read of "
                        "/home/user/Invoices.xlsx shows pivot table in Sheet2 with "
                        'col_fields=["Invoice No"] — matches instruction.\''
                    ),
                },
                "diagnosis": {
                    "type": "string",
                    "description": (
                        "REQUIRED when decision='failed'. What the VM state actually "
                        "is vs. what the instruction requires, grounded in probe "
                        'output. Example: \'settings.json contains "editor.fontSize":'
                        " 14 but task asked for 18. The file is saved but with wrong "
                        "value.' Leave empty for 'done' or 'infeasible'."
                    ),
                },
                "suggestion": {
                    "type": "string",
                    "description": (
                        "REQUIRED when decision='failed'. A specific, different "
                        "approach for the next attempt — not 'try again'. State what "
                        "to do differently and why. Example: 'Write settings.json "
                        "directly with python instead of using the Settings UI — the "
                        "UI change did not persist because VS Code was not focused "
                        "when Ctrl+, was pressed.' Leave empty for 'done' or "
                        "'infeasible'."
                    ),
                },
            },
            "required": ["decision", "reasoning"],
        },
    },
)
