#!/usr/bin/env python3
"""Run a small AirXDB Computer MCP smoke test with screenshot artifacts."""

from __future__ import annotations

import argparse
import base64
import json
import os
import shutil
import socket
import subprocess
import tarfile
import tempfile
import threading
import time
import urllib.error
import urllib.request
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Tuple

REQUIRED_BASE_ENV = [
    "MIDSCENE_MODEL_NAME",
    "MIDSCENE_MODEL_BASE_URL",
    "MIDSCENE_MODEL_API_KEY",
]
REQUIRED_SEMANTIC_ENV = [
    *REQUIRED_BASE_ENV,
    "MIDSCENE_MODEL_FAMILY",
]
VALID_MODEL_FAMILIES = [
    "doubao-vision",
    "doubao-seed",
    "gemini",
    "qwen2.5-vl",
    "qwen3-vl",
    "qwen3.5",
    "qwen3.6",
    "vlm-ui-tars",
    "vlm-ui-tars-doubao",
    "vlm-ui-tars-doubao-1.5",
    "glm-v",
    "auto-glm",
    "auto-glm-multilingual",
    "gpt-5",
]


def now_stamp() -> str:
    return datetime.now().strftime("%Y%m%d-%H%M%S")


def load_env(project_root: Path) -> Dict[str, str]:
    env = os.environ.copy()
    env_file = project_root / "AirPlan" / "state" / "airxdb" / "midscene.local.env"
    if not env_file.exists():
        return env

    for raw_line in env_file.read_text(encoding="utf-8-sig").splitlines():
        line = raw_line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, value = line.split("=", 1)
        key = key.strip()
        value = value.strip().strip('"').strip("'")
        if key and value and key not in env:
            env[key] = value
    return env


def infer_model_family(model_name: str) -> str:
    normalized = model_name.strip().lower()
    if normalized.startswith("gpt-5"):
        return "gpt-5"
    if "qwen3" in normalized:
        return "qwen3-vl"
    if "qwen2.5" in normalized or "qwen-2.5" in normalized:
        return "qwen2.5-vl"
    if "gemini" in normalized:
        return "gemini"
    if "doubao" in normalized:
        return "doubao-seed"
    if "glm" in normalized:
        return "glm-v"
    return ""


def validate_env(env: Dict[str, str], semantic: bool) -> List[str]:
    if not semantic:
        return []
    required = REQUIRED_SEMANTIC_ENV if semantic else REQUIRED_BASE_ENV
    missing = [name for name in required if not env.get(name)]
    family = env.get("MIDSCENE_MODEL_FAMILY", "")
    if semantic and family and family not in VALID_MODEL_FAMILIES:
        missing.append("MIDSCENE_MODEL_FAMILY(valid value)")
    return missing


def redact(text: str, env: Dict[str, str]) -> str:
    secret = env.get("MIDSCENE_MODEL_API_KEY")
    if secret:
        text = text.replace(secret, "<redacted>")
    return text


def command_name(name: str) -> str:
    if os.name == "nt":
        candidate = shutil.which(f"{name}.cmd")
        if candidate:
            return candidate
    return shutil.which(name) or name


def free_port() -> int:
    sock = socket.socket()
    sock.bind(("127.0.0.1", 0))
    port = sock.getsockname()[1]
    sock.close()
    return int(port)


def npm_cache_root(env: Dict[str, str]) -> Optional[Path]:
    local_app_data = env.get("LOCALAPPDATA")
    if local_app_data:
        return Path(local_app_data) / "npm-cache"
    home = env.get("HOME") or env.get("USERPROFILE")
    return Path(home) / ".npm" if home else None


def midscene_dist_dirs(project_root: Path, env: Dict[str, str]) -> List[Path]:
    candidates: List[Path] = []
    local = project_root / "node_modules" / "@midscene" / "computer-mcp" / "dist"
    if local.exists():
        candidates.append(local)

    cache = npm_cache_root(env)
    if cache:
        candidates.extend(cache.glob("_npx/*/node_modules/@midscene/computer-mcp/dist"))

    unique: List[Path] = []
    seen = set()
    for path in candidates:
        resolved = str(path.resolve())
        if resolved not in seen:
            seen.add(resolved)
            unique.append(path)
    return unique


def ensure_windows_screenshot_assets(project_root: Path, env: Dict[str, str]) -> Tuple[str, List[str]]:
    if os.name != "nt":
        return "skipped-non-windows", []

    dist_dirs = midscene_dist_dirs(project_root, env)
    missing_dirs = [
        path for path in dist_dirs
        if not (path / "screenCapture_1.3.2.bat").exists() or not (path / "app.manifest").exists()
    ]
    if not missing_dirs:
        return "ok", [str(path) for path in dist_dirs]
    if not dist_dirs:
        return "missing-midscene-package", []

    npm = command_name("npm")
    with tempfile.TemporaryDirectory(prefix="airxdb-screenshot-desktop-") as tmp:
        tmp_path = Path(tmp)
        subprocess.run(
            [npm, "pack", "screenshot-desktop@1.15.3", "--pack-destination", str(tmp_path)],
            cwd=str(project_root),
            env=env,
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            encoding="utf-8",
            errors="replace",
        )
        tgz = next(tmp_path.glob("screenshot-desktop-1.15.3.tgz"))
        with tarfile.open(tgz, "r:gz") as archive:
            archive.extract("package/lib/win32/screenCapture_1.3.2.bat", path=tmp_path)
            archive.extract("package/lib/win32/app.manifest", path=tmp_path)
        bat = tmp_path / "package" / "lib" / "win32" / "screenCapture_1.3.2.bat"
        manifest = tmp_path / "package" / "lib" / "win32" / "app.manifest"
        for dist in missing_dirs:
            shutil.copy2(bat, dist / "screenCapture_1.3.2.bat")
            shutil.copy2(manifest, dist / "app.manifest")
    return "repaired", [str(path) for path in missing_dirs]


def decode_response(raw: str, content_type: str) -> Optional[Dict[str, Any]]:
    if "text/event-stream" in content_type or raw.startswith("event:") or raw.startswith("data:"):
        data_lines = [line[5:].strip() for line in raw.splitlines() if line.startswith("data:")]
        raw = "\n".join(data_lines).strip()
    return json.loads(raw) if raw.strip() else None


class MCPClient:
    def __init__(self, url: str, env: Dict[str, str], timeout: int) -> None:
        self.url = url
        self.env = env
        self.timeout = timeout
        self.session_id: Optional[str] = None
        self.seq = 1

    def post(self, method: str, params: Optional[Dict[str, Any]] = None, expect_id: bool = True) -> Optional[Dict[str, Any]]:
        payload: Dict[str, Any] = {"jsonrpc": "2.0", "method": method, "params": params or {}}
        if expect_id:
            payload["id"] = self.seq
            self.seq += 1

        headers = {
            "Content-Type": "application/json",
            "Accept": "application/json, text/event-stream",
        }
        if self.session_id:
            headers["mcp-session-id"] = self.session_id

        request = urllib.request.Request(
            self.url,
            data=json.dumps(payload).encode("utf-8"),
            headers=headers,
            method="POST",
        )
        try:
            with urllib.request.urlopen(request, timeout=self.timeout) as response:
                raw = response.read().decode("utf-8", "replace")
                raw = redact(raw, self.env)
                self.session_id = response.headers.get("mcp-session-id") or self.session_id
                return decode_response(raw, response.headers.get("content-type", ""))
        except urllib.error.HTTPError as error:
            raw = error.read().decode("utf-8", "replace")
            raise RuntimeError(f"HTTP {error.code} for {method}: {redact(raw, self.env)}") from error

    def call_tool(self, name: str, arguments: Optional[Dict[str, Any]] = None) -> Optional[Dict[str, Any]]:
        return self.post("tools/call", {"name": name, "arguments": arguments or {}})


def content_texts(result: Optional[Dict[str, Any]]) -> List[str]:
    content = (result or {}).get("result", {}).get("content", [])
    return [item.get("text", "") for item in content if item.get("type") == "text"]


def save_images(result: Optional[Dict[str, Any]], output_dir: Path, stem: str) -> List[str]:
    saved: List[str] = []
    content = (result or {}).get("result", {}).get("content", [])
    for index, item in enumerate(content):
        if item.get("type") != "image" or not item.get("data"):
            continue
        mime = item.get("mimeType", "image/png")
        ext = ".jpg" if "jpeg" in mime or "jpg" in mime else ".png"
        path = output_dir / f"{stem}-{index}{ext}"
        path.write_bytes(base64.b64decode(item["data"]))
        saved.append(str(path))
    return saved


def wait_for_port(port: int, proc: subprocess.Popen[str], logs: List[str], timeout: int) -> None:
    deadline = time.time() + timeout
    while time.time() < deadline:
        if proc.poll() is not None:
            raise RuntimeError(f"computer-mcp exited {proc.returncode}\n" + "\n".join(logs[-80:]))
        try:
            with socket.create_connection(("127.0.0.1", port), timeout=0.2):
                return
        except OSError:
            time.sleep(0.25)
    raise TimeoutError("computer-mcp HTTP server did not open a port\n" + "\n".join(logs[-80:]))


def start_log_pump(stream: Any, prefix: str, logs: List[str], env: Dict[str, str]) -> None:
    def pump() -> None:
        for line in iter(stream.readline, ""):
            logs.append(prefix + redact(line.rstrip(), env))

    threading.Thread(target=pump, daemon=True).start()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run AirXDB Computer MCP smoke test.")
    parser.add_argument("--project", default=".")
    parser.add_argument("--prompt", default="Windows taskbar Start button")
    parser.add_argument("--action", choices=["connect", "screenshot", "mousemove", "act"], default="mousemove")
    parser.add_argument("--output-dir", default="")
    parser.add_argument("--timeout", type=int, default=240)
    parser.add_argument("--no-repair-screenshot-assets", action="store_true")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    project_root = Path(args.project).expanduser().resolve()
    output_dir = Path(args.output_dir).expanduser().resolve() if args.output_dir else project_root / "AirPlan" / "docs" / "debug" / "airxdb-artifacts"
    output_dir.mkdir(parents=True, exist_ok=True)
    env = load_env(project_root)
    semantic = args.action in {"mousemove", "act"}
    missing = validate_env(env, semantic)
    if missing:
        print("airxdb_smoke=blocked")
        print("missing_midscene_env=" + ",".join(missing))
        suggestion = infer_model_family(env.get("MIDSCENE_MODEL_NAME", ""))
        if suggestion:
            print(f"suggested_midscene_model_family={suggestion}")
        print("valid_midscene_model_families=" + ",".join(VALID_MODEL_FAMILIES))
        raise SystemExit(2)

    port = free_port()
    logs: List[str] = []
    npx = command_name("npx")
    proc = subprocess.Popen(
        [npx, "-y", "@midscene/computer-mcp", "--mode", "http", "--host", "127.0.0.1", "--port", str(port)],
        cwd=str(project_root),
        env=env,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        encoding="utf-8",
        errors="replace",
    )
    assert proc.stdout is not None
    assert proc.stderr is not None
    start_log_pump(proc.stdout, "OUT ", logs, env)
    start_log_pump(proc.stderr, "ERR ", logs, env)

    steps: List[Dict[str, Any]] = []
    asset_status = "skipped"
    asset_paths: List[str] = []
    try:
        wait_for_port(port, proc, logs, min(args.timeout, 60))
        if not args.no_repair_screenshot_assets:
            asset_status, asset_paths = ensure_windows_screenshot_assets(project_root, env)

        client = MCPClient(f"http://127.0.0.1:{port}/mcp", env, args.timeout)
        client.post(
            "initialize",
            {
                "protocolVersion": "2024-11-05",
                "capabilities": {},
                "clientInfo": {"name": "airxdb-computer-mcp-smoke", "version": "0.2.1"},
            },
        )
        client.post("notifications/initialized", {}, expect_id=False)

        actions: List[Tuple[str, Dict[str, Any], str]] = [
            ("computer_connect", {}, "connect"),
        ]
        if args.action == "mousemove":
            actions.append(("MouseMove", {"locate": {"prompt": args.prompt}}, "mousemove"))
        elif args.action == "act":
            actions.append(("act", {"prompt": args.prompt}, "act"))
        if args.action in {"screenshot", "mousemove", "act"}:
            actions.append(("take_screenshot", {}, "final"))

        for tool, arguments, stem in actions:
            result = client.call_tool(tool, arguments)
            steps.append(
                {
                    "tool": tool,
                    "arguments": arguments,
                    "texts": content_texts(result),
                    "images": save_images(result, output_dir, f"{now_stamp()}-{stem}"),
                }
            )
            time.sleep(0.3)

        try:
            result = client.call_tool("computer_disconnect", {})
            steps.append({"tool": "computer_disconnect", "arguments": {}, "texts": content_texts(result), "images": []})
        except Exception as error:
            steps.append({"tool": "computer_disconnect", "arguments": {}, "texts": [str(error)], "images": []})

        failed = [
            step for step in steps
            if any(text.lower().startswith(("warning:", "failed", "error")) for text in step["texts"])
        ]
        status = "failed" if failed else "ok"
        report = {
            "status": status,
            "action": args.action,
            "prompt": args.prompt,
            "modelName": env.get("MIDSCENE_MODEL_NAME", ""),
            "modelFamily": env.get("MIDSCENE_MODEL_FAMILY", ""),
            "assetRepair": {"status": asset_status, "paths": asset_paths},
            "steps": steps,
            "logsTail": logs[-80:],
        }
        report_path = output_dir / f"{now_stamp()}-airxdb-smoke.json"
        report_path.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")

        print(f"airxdb_smoke={status}")
        print(f"capture_mode={'semantic' if semantic else 'screenshot'}")
        print(f"report={report_path}")
        print(f"asset_repair={asset_status}")
        for step in steps:
            print(f"step={step['tool']}")
            for text in step["texts"]:
                print("text=" + text.replace("\n", " | ")[:1600])
            for image in step["images"]:
                print(f"image={image}")
        raise SystemExit(1 if failed else 0)
    finally:
        proc.terminate()
        try:
            proc.wait(timeout=8)
        except subprocess.TimeoutExpired:
            proc.kill()


if __name__ == "__main__":
    main()