# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Model Context Protocol (MCP) server command (artemis mcp)."""

import json
import os
from pathlib import Path
import re
import sys
import threading
import tomllib
from typing import Annotated

from mcp_server.base import mcp as agent_mcp
import mcp_server.tools  # noqa: F401
from mcp_server.utils import env_utils
from artemis.mcp.adb_server import mcp as adb_mcp
from artemis.runtime import shutdown_awake_service, start_awake_service
from artemis.utils.logger import get_logger
from rich.console import Console
from rich.syntax import Syntax
import typer

logger = get_logger(__name__)
console = Console()

ARTEMIS_MCP_TOOLS = (
    "mobile_run_task",
    "mobile_manage_task",
    "mobile_get_device_state",
    "mobile_inspect_trace",
    "mobile_diagnose",
    "mobile_workspace",
    "mobile_config",
    "mobile_task_file",
    "mobile_batch",
    "mobile_search_history",
    "mobile_export_report",
)


def _get_config_snippet(client: str, python_exe: str, project_root: str) -> dict:
    """Generate the supported Codex MCP configuration."""
    if client != "codex":
        raise ValueError("Only Codex MCP configuration is supported")
    return {
        "mcp_servers": {
            "artemis": {
                "command": python_exe,
                "args": ["-m", "mcp_server"],
                "cwd": project_root,
                "env": {
                    "PYTHONUNBUFFERED": "1",
                    "PYTHONPATH": project_root,
                    "ARTEMIS_DESKTOP_NOTIFY": "true",
                },
                "enabled": True,
                "required": True,
                "startup_timeout_sec": 120,
                "enabled_tools": list(ARTEMIS_MCP_TOOLS),
            }
        }
    }


def _codex_toml_block(server_config: dict) -> str:
    """Renders a managed Codex MCP server block using TOML-compatible JSON strings."""
    command = json.dumps(str(server_config["command"]), ensure_ascii=False)
    args = json.dumps(server_config.get("args", []), ensure_ascii=False)
    cwd = json.dumps(str(server_config["cwd"]), ensure_ascii=False)
    env = server_config.get("env", {})
    lines = [
        "# BEGIN ARTEMIS MCP CONFIG",
        "[mcp_servers.artemis]",
        f"command = {command}",
        f"args = {args}",
        f"cwd = {cwd}",
    ]
    for key in ("enabled", "required"):
        if key in server_config:
            lines.append(f"{key} = {'true' if server_config[key] else 'false'}")
    if "startup_timeout_sec" in server_config:
        lines.append(f"startup_timeout_sec = {int(server_config['startup_timeout_sec'])}")
    if "enabled_tools" in server_config:
        lines.append(
            f"enabled_tools = {json.dumps(server_config['enabled_tools'], ensure_ascii=False)}"
        )
    if env:
        lines.extend(["", "[mcp_servers.artemis.env]"])
        for key, value in env.items():
            lines.append(f"{key} = {json.dumps(str(value), ensure_ascii=False)}")
    lines.append("# END ARTEMIS MCP CONFIG")
    return "\n".join(lines) + "\n"


def _merge_codex_toml(file_path: Path, server_config: dict) -> bool:
    """Adds or updates Artemis in Codex config.toml while preserving unrelated settings."""
    try:
        file_path.parent.mkdir(parents=True, exist_ok=True)
        existing = file_path.read_text(encoding="utf-8") if file_path.exists() else ""
        managed_block = _codex_toml_block(server_config)
        begin_marker = "# BEGIN ARTEMIS MCP CONFIG"
        end_marker = "# END ARTEMIS MCP CONFIG"

        if begin_marker in existing and end_marker in existing:
            managed_pattern = re.compile(
                rf"(?ms)^{re.escape(begin_marker)}\n.*?^{re.escape(end_marker)}\n?"
            )
            # Windows paths contain backslashes; a replacement string would
            # interpret them as regex escapes on a second installation.
            updated = managed_pattern.sub(lambda _: managed_block, existing, count=1)
        else:
            # Remove an older, unmanaged Artemis table before adding the managed block.
            kept_lines: list[str] = []
            skipping_artemis = False
            header_pattern = re.compile(r"^\s*\[{1,2}\s*(.+?)\s*\]{1,2}\s*(?:#.*)?$")
            for line in existing.splitlines(keepends=True):
                header_match = header_pattern.match(line.rstrip("\r\n"))
                if header_match:
                    normalized = re.sub(r"[\s\"']", "", header_match.group(1)).lower()
                    skipping_artemis = normalized == "mcp_servers.artemis" or normalized.startswith(
                        "mcp_servers.artemis."
                    )
                if not skipping_artemis:
                    kept_lines.append(line)
            preserved = "".join(kept_lines).rstrip()
            updated = f"{preserved}\n\n{managed_block}" if preserved else managed_block

        # Never replace a user's config with invalid TOML.
        tomllib.loads(updated)
        if updated != existing:
            file_path.write_text(updated, encoding="utf-8")
        return True
    except Exception as e:
        if file_path.exists():
            try:
                backup_path = file_path.with_suffix(file_path.suffix + ".bak")
                backup_path.write_text(file_path.read_text(encoding="utf-8"), encoding="utf-8")
                logger.warning(
                    f"Could not update {file_path}; backup created at {backup_path}: {e}"
                )
            except (OSError, UnicodeDecodeError) as backup_exc:
                logger.warning(
                    f"Could not update {file_path} ({e}); backup also failed: {backup_exc}"
                )
        else:
            logger.warning(f"Could not create Codex MCP config file {file_path}: {e}")
        return False


def _inject_rules_block(file_path: Path, content: str) -> bool:
    """Injects or updates a managed ARTEMIS block in a global rules file without overwriting user content."""
    try:
        file_path.parent.mkdir(parents=True, exist_ok=True)
        begin_marker = "<!-- BEGIN ARTEMIS MOBILE TESTING RULES -->"
        end_marker = "<!-- END ARTEMIS MOBILE TESTING RULES -->"
        new_block = f"{begin_marker}\n{content.strip()}\n{end_marker}\n"

        existing = ""
        if file_path.exists():
            try:
                existing = file_path.read_text(encoding="utf-8")
            except Exception:
                existing = ""

        if begin_marker in existing and end_marker in existing:
            prefix = existing.split(begin_marker)[0]
            suffix = existing.split(end_marker)[1]
            updated = f"{prefix.rstrip()}\n\n{new_block}\n{suffix.lstrip()}".strip() + "\n"
        else:
            updated = f"{existing.rstrip()}\n\n{new_block}".strip() + "\n"

        if existing == updated:
            return True

        file_path.write_text(updated, encoding="utf-8")
        return True
    except Exception as e:
        logger.warning(f"Could not inject rules block into {file_path}: {e}")
        return False


def install_rules(client: str, project_root: str) -> list[str]:
    """Install testing rules only into the Codex instructions file."""
    if client != "codex":
        raise ValueError("Only Codex MCP configuration is supported")
    rules_src = Path(project_root) / "mcp_server" / "rules.md"
    if not rules_src.exists():
        rules_src = Path(__file__).resolve().parents[4] / "mcp_server" / "rules.md"
    try:
        raw_rules = rules_src.read_text(encoding="utf-8").strip()
    except OSError:
        logger.warning(f"Could not read rules file {rules_src}")
        return []
    codex_home = Path(os.environ.get("CODEX_HOME") or (Path.home() / ".codex"))
    override_file = codex_home / "AGENTS.override.md"
    agents_file = (
        override_file
        if override_file.exists() and override_file.read_text(encoding="utf-8").strip()
        else codex_home / "AGENTS.md"
    )
    return [str(agents_file)] if _inject_rules_block(agents_file, raw_rules) else []


def install_mcp_config(client: str, python_exe: str, project_root: str) -> list[str]:
    """Install ARTEMIS configuration and rules only for Codex."""
    server_cfg = _get_config_snippet(client, python_exe, project_root)["mcp_servers"]["artemis"]
    codex_home = Path(os.environ.get("CODEX_HOME") or (Path.home() / ".codex"))
    codex_path = codex_home / "config.toml"
    if not _merge_codex_toml(codex_path, server_cfg):
        raise typer.Exit(1)
    return [str(codex_path), *install_rules(client, project_root)]


def mcp_command(
    server_type: Annotated[
        str,
        typer.Option(
            "--type",
            "-t",
            help="Type of MCP server to start: 'agent' (default, universal IDE mobile agent), 'adb' (raw adb), 'xml' (xml fuzzy search).",
        ),
    ] = "agent",
    transport: Annotated[
        str,
        typer.Option(
            "--transport",
            help="MCP transport protocol ('stdio' or 'sse').",
        ),
    ] = "stdio",
    host: Annotated[
        str,
        typer.Option(
            "--host",
            help="Host address when running with SSE transport.",
        ),
    ] = "127.0.0.1",
    port: Annotated[
        int,
        typer.Option(
            "--port",
            "-p",
            help="Port number when running with SSE transport.",
        ),
    ] = 8001,
    install_config: Annotated[
        str | None,
        typer.Option(
            "--install",
            "-i",
            help="Install ARTEMIS MCP configuration and testing rules for Codex: --install codex.",
        ),
    ] = None,
    generate_config: Annotated[
        str | None,
        typer.Option(
            "--generate-config",
            "-g",
            help="Output ready-to-use Codex TOML configuration: --generate-config codex.",
        ),
    ] = None,
) -> None:
    """Launch or configure the ARTEMIS Model Context Protocol (MCP) server."""
    project_root = env_utils.get_project_root()
    python_exe = env_utils.resolve_python_executable(project_root)

    if install_config:
        client = install_config.lower()
        if client != "codex":
            console.print(
                "[bold red]目前仅支持 Codex MCP 配置，请使用 --install codex。[/bold red]"
            )
            raise typer.Exit(1)
        installed_paths = install_mcp_config(client, python_exe, project_root)
        console.print(
            "[bold green]✔ Successfully installed ARTEMIS MCP server configuration & rules to:[/bold green]"
        )
        for path in installed_paths:
            console.print(f"  • [cyan]{path}[/cyan]")
        console.print(
            "\n[dim]Please restart or reload your IDE window to activate the Artemis MCP tools.[/dim]"
        )
        raise typer.Exit(0)

    if generate_config:
        client = generate_config.lower()
        if client != "codex":
            console.print(
                "[bold red]目前仅支持 Codex MCP 配置，请使用 --generate-config codex。[/bold red]"
            )
            raise typer.Exit(1)
        server_cfg = _get_config_snippet(client, python_exe, project_root)["mcp_servers"]["artemis"]
        json_str = _codex_toml_block(server_cfg)
        syntax_language = "toml"

        syntax = Syntax(json_str, syntax_language, theme="monokai", line_numbers=False)
        console.print(f"[bold cyan]MCP Configuration for {client.upper()}:[/bold cyan]")
        console.print(syntax)
        raise typer.Exit(0)

    threading.Thread(target=start_awake_service, daemon=True, name="artemis-awake-init").start()
    try:
        st = server_type.lower()
        if st in ("agent", "mobile", "artemis", "default"):
            logger.info(f"Starting Artemis Mobile Agent MCP Server over {transport}...")
            if transport.lower() == "sse":
                agent_mcp.run(transport="sse", host=host, port=port)
            else:
                agent_mcp.run(transport="stdio")
        elif st == "adb":
            if transport.lower() == "sse":
                logger.info(f"Starting Artemis ADB MCP Server over {transport}...")
                adb_mcp.run(transport="sse", host=host, port=port)
            else:
                # stdio carries the MCP JSON-RPC stream: redirect logging and detach
                # from the parent DataEngine before any output can corrupt it.
                from artemis.mcp.adb_server import configure_stdio_mode

                configure_stdio_mode()
                logger.info("Starting Artemis ADB MCP Server over stdio...")
                adb_mcp.run(transport="stdio")
        elif st == "xml":
            from artemis.mcp.xml_search_server import mcp as xml_mcp

            logger.info(f"Starting Artemis XML Fuzzy Search MCP Server over {transport}...")
            if transport.lower() == "sse":
                xml_mcp.run(transport="sse", host=host, port=port)
            else:
                xml_mcp.run(transport="stdio")
        else:
            logger.error(
                f"Unsupported MCP server type: {server_type}. Use 'agent', 'adb', or 'xml'."
            )
            raise typer.Exit(1)
    finally:
        shutdown_awake_service()