"""Bounded script execution and task-owned evidence for the device-test plugin."""

from __future__ import annotations

import asyncio
from contextlib import suppress
import hashlib
import ipaddress
import json
import mimetypes
import os
from pathlib import Path
import re
import shutil
import signal
import subprocess
import sys
import time
from typing import Any
from urllib.parse import parse_qsl, unquote, urlsplit
import uuid

from mcp.types import CallToolResult, TextContent


PACKAGE = Path(__file__).resolve().parent
SCRIPTS = PACKAGE / "skill" / "scripts"
MAX_FILE_BYTES = 64 * 1024 * 1024
MAX_OUTPUT_BYTES = 4 * 1024 * 1024
SECRET_KEY = re.compile(
    r"(?:^|_)(?:password|passwd|pwd|secret|token|access_token|api_key|apikey|authorization)$",
    re.IGNORECASE,
)


class InputError(ValueError):
    """Only constant, public messages belong in this exception."""


def local_file(value: str) -> Path:
    if (
        not value
        or "://" in value
        or value.startswith(("\\\\", "//"))
        or re.search(r"[\x00-\x1f]", value)
        or not Path(value).is_absolute()
    ):
        raise InputError("需要本机文件的绝对路径；不接受 URL 或网络共享。")
    path = Path(value).resolve()
    if not path.is_file():
        raise InputError("本机输入文件不存在。")
    return path


def unicast_host(value: str) -> str:
    try:
        if "%" in value:
            raise ValueError
        address = ipaddress.ip_address(value)
        effective = address.ipv4_mapped if isinstance(address, ipaddress.IPv6Address) else None
        effective = effective or address
        if (
            effective.is_unspecified
            or effective.is_multicast
            or str(effective) == "255.255.255.255"
        ):
            raise ValueError
    except ValueError:
        raise InputError("需要一个明确的单播 IP；不接受主机名、网段或广播地址。") from None
    return str(address)


def operation_budget() -> int:
    value = os.environ.get("ARTEMIS_DEVICE_TEST_TIMEOUT", "45")
    if not re.fullmatch(r"[0-9]+", value) or not 10 <= int(value) <= 3300:
        raise InputError("插件操作时限必须是 10..3300 秒的整数。")
    return int(value)


def check_budget(required: float) -> None:
    if required > operation_budget():
        raise InputError("所需时限超过插件操作预算；请缩短操作或同时调整宿主与插件时限。")


def artifact_root() -> Path:
    value = os.environ.get("ARTEMIS_ARTIFACTS_DIR", "")
    if not value:
        raise InputError("缺少 Artemis 注入的任务资料目录；请通过 Artemis 任务调用。")
    path = Path(value).absolute()
    if path.resolve() != path or path.is_symlink():
        raise InputError("任务资料目录不能通过链接指向其他位置。")
    path.mkdir(parents=True, exist_ok=True)
    return path


def new_artifact(prefix: str, suffix: str) -> Path:
    return artifact_root() / f"{prefix}-{uuid.uuid4().hex}{suffix}"


def task_evidence(name: str, *, allow_missing: bool = False) -> Path:
    if (
        not name
        or len(name) > 180
        or name != Path(name).name
        or any(char in name for char in '/\\:*?"<>|')
        or name in {".", ".."}
        or name.rstrip(" .") != name
        or re.search(r"[\x00-\x1f]", name)
    ):
        raise InputError("证据必须引用当前任务中的普通附件文件名。")
    path = artifact_root() / name
    if path.is_symlink() or path.resolve().parent != artifact_root():
        raise InputError("证据文件不存在或超出当前任务资料目录。")
    if not path.is_file() and not (allow_missing and not path.exists()):
        raise InputError("证据文件不存在或超出当前任务资料目录。")
    if path.is_file() and path.stat().st_size > MAX_FILE_BYTES:
        raise InputError("单份证据不能超过 64 MiB。")
    return path


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def sanitized(value: Any, secrets: tuple[str, ...] = ()) -> Any:
    if isinstance(value, dict):
        return {
            key: "[REDACTED]" if SECRET_KEY.search(str(key)) else sanitized(item, secrets)
            for key, item in value.items()
        }
    if isinstance(value, list):
        return [sanitized(item, secrets) for item in value]
    if not isinstance(value, str):
        return value
    for secret in sorted(filter(None, secrets), key=len, reverse=True):
        if value == secret:
            return "[REDACTED]"
        if len(secret) >= 4:
            value = value.replace(secret, "[REDACTED]")
    value = re.sub(r"(?i)rtsps?://[^\s\"<>]+", "[RTSP_URL_REDACTED]", value)
    return re.sub(
        r"(?i)(password|passwd|pwd|token|secret|api_key|authorization)\s*[:=]\s*[^\s,;]+",
        r"\1=[REDACTED]",
        value,
    )


def rtsp_environment(name: str) -> tuple[dict[str, str], tuple[str, ...]]:
    if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]{0,99}", name):
        raise InputError("RTSP 输入必须是环境变量名称。")
    value = os.environ.get(name, "")
    config = os.environ.get("ARTEMIS_DEVICE_TEST_ENV_FILE", "")
    if not value and config:
        from dotenv import dotenv_values

        value = dotenv_values(local_file(config), interpolate=False).get(name) or ""
    try:
        parsed = urlsplit(value)
        if (
            parsed.scheme not in {"rtsp", "rtsps"}
            or not parsed.hostname
            or parsed.port == 0
            or parsed.fragment
            or re.search(r"[\x00-\x20\x7f]", value)
        ):
            raise ValueError
    except ValueError:
        raise InputError("指定的 RTSP 环境变量未配置或格式无效；请在本机凭据环境中配置。") from None
    secrets = tuple(
        filter(
            None,
            (
                value,
                parsed.username,
                parsed.password,
                unquote(parsed.password or ""),
                *(item for _, item in parse_qsl(parsed.query)),
            ),
        )
    )
    return {name: value}, secrets


def runner_arguments(command: str, arguments: list[str]) -> list[str]:
    budget = operation_budget()
    result = ["--timeout", str(budget)]
    profile = os.environ.get("ARTEMIS_DEVICE_TEST_PROFILE", "")
    tool_root = os.environ.get("ARTEMIS_DEVICE_TEST_TOOL_ROOT", "")
    if profile:
        result += ["--profile", str(local_file(profile))]
    if tool_root:
        root = Path(tool_root)
        if not root.is_absolute() or not root.is_dir() or tool_root.startswith(("\\\\", "//")):
            raise InputError("工具目录必须是本机已解压目录的绝对路径。")
        result += ["--tool-root", str(root.resolve())]
    # Leave time for the runner to terminate its own child before the outer deadline.
    result[1] = str(max(1, budget - 3))
    return [*result, command, *arguments]


def failure(
    message: str,
    code: str = "invalid_input",
    *,
    execution: str = "not_started",
    side_effect: str = "none",
) -> CallToolResult:
    return CallToolResult(
        isError=True,
        content=[TextContent(type="text", text=message)],
        structuredContent={
            "test_verdict": "not_evaluated",
            "artemis": {
                "version": 1,
                "execution": execution,
                "side_effect": side_effect,
                "retryable": False,
                "error_code": code,
                "attachments": [],
            },
        },
    )


def attach_file(path: Path) -> dict:
    if path.is_symlink() or not path.is_file() or path.stat().st_size > MAX_FILE_BYTES:
        raise InputError("附件缺失、使用链接或超过 64 MiB；不能声称已归档。")
    root = artifact_root()
    if path.parent.resolve() != root:
        destination = new_artifact("evidence", path.suffix[:16])
        temporary = destination.with_name(destination.name + ".tmp")
        with path.open("rb") as source, temporary.open("xb") as target:
            shutil.copyfileobj(source, target)
        if sha256(path) != sha256(temporary):
            raise InputError("证据复制后的 SHA-256 不匹配。")
        temporary.rename(destination)
        path = destination
    return {
        "name": path.name,
        "media_type": mimetypes.guess_type(path.name)[0] or "application/octet-stream",
        "sha256": sha256(path),
    }


def receipt(
    operation: str,
    data: dict,
    *,
    target: str = "local-environment",
    kind: str = "terminal",
    execution: str = "completed",
    side_effect: str = "none",
    error_code: str = "",
    files: list[Path] | None = None,
) -> CallToolResult:
    data = sanitized(data)
    attachments = []
    issues = []
    for path in files or []:
        try:
            attachments.append(attach_file(path))
        except (InputError, OSError):
            issues.append("某份输出未能回读归档；可能缺失、使用链接或超过 64 MiB。")
    if issues:
        data["attachment_issues"] = issues
        error_code = error_code or "attachment_unavailable"
    data["test_verdict"] = "not_evaluated"
    log = new_artifact(operation, ".json")
    with log.open("x", encoding="utf-8") as stream:
        json.dump(data, stream, ensure_ascii=False, indent=2)
    attachments.append(attach_file(log))
    target_id = (
        target if len(target) <= 256 else "sha256:" + hashlib.sha256(target.encode()).hexdigest()
    )
    data["artemis"] = {
        "version": 1,
        "execution": execution,
        "side_effect": side_effect,
        "retryable": False,
        "source": {
            "kind": kind,
            "target_id": target_id,
            "label": operation,
            "captured_at": time.time(),
        },
        "attachments": attachments,
    }
    if error_code:
        data["artemis"]["error_code"] = error_code
    text = json.dumps(data, ensure_ascii=False)
    if len(text) > 16000:
        text = text[:15000] + "\n结果较长；完整结构与任务 JSON 附件保留原始字段。"
    return CallToolResult(
        isError=bool(error_code),
        content=[TextContent(type="text", text=text)],
        structuredContent=data,
    )


async def stop_process(process: asyncio.subprocess.Process) -> None:
    """Reap the worker and its fixed tool descendants before returning from cancellation."""
    if os.name == "nt":
        command = str(
            Path(os.environ.get("SystemRoot", r"C:\Windows")) / "System32" / "taskkill.exe"
        )
        killer = await asyncio.create_subprocess_exec(
            command,
            "/PID",
            str(process.pid),
            "/T",
            "/F",
            stdout=asyncio.subprocess.DEVNULL,
            stderr=asyncio.subprocess.DEVNULL,
            creationflags=subprocess.CREATE_NO_WINDOW,
        )
        try:
            await asyncio.wait_for(killer.wait(), 5)
        except TimeoutError:
            killer.kill()
            await killer.wait()
    else:
        with suppress(ProcessLookupError):
            os.killpg(process.pid, signal.SIGKILL)
    if process.returncode is None:
        with suppress(ProcessLookupError):
            process.kill()
    await process.wait()


async def read_output(process: asyncio.subprocess.Process) -> bytes:
    chunks = []
    size = 0
    while chunk := await process.stdout.read(65536):
        size += len(chunk)
        if size > MAX_OUTPUT_BYTES:
            raise InputError("工具输出超过 4 MiB，已停止；部分输出不作为有效观测。")
        chunks.append(chunk)
    await process.wait()
    return b"".join(chunks)


async def run_cli(
    script: str,
    arguments: list[str],
    *,
    operation: str,
    target: str = "local-environment",
    kind: str = "terminal",
    write: bool = False,
    files: list[Path] | None = None,
    env: dict[str, str] | None = None,
    secrets: tuple[str, ...] = (),
    timeout: float | None = None,
) -> CallToolResult:
    artifact_root()
    budget = timeout or operation_budget()
    native_scripts = {"network_inventory.py", "dependency_runner.py"}
    if script not in {"tool_runner.py", "device_adapter.py", "evidence.py", *native_scripts}:
        raise InputError("不支持的脚本入口。")
    script_path = (PACKAGE if script in native_scripts else SCRIPTS) / script
    child_env = os.environ.copy()
    for name in ("FFREPORT", "PYTHONPATH", "PYTHONHOME"):
        child_env.pop(name, None)
    child_env.update(env or {})
    options = (
        {"creationflags": subprocess.CREATE_NO_WINDOW | subprocess.CREATE_NEW_PROCESS_GROUP}
        if os.name == "nt"
        else {"start_new_session": True}
    )
    try:
        process = await asyncio.create_subprocess_exec(
            sys.executable,
            "-I",
            "-B",
            str(script_path),
            *arguments,
            cwd=str(PACKAGE),
            env=child_env,
            stdin=asyncio.subprocess.DEVNULL,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.DEVNULL,
            **options,
        )
    except OSError:
        return failure("固定脚本无法启动；未执行设备动作。", "worker_launch_failed")
    try:
        output = await asyncio.wait_for(read_output(process), budget)
    except (TimeoutError, InputError) as exc:
        await stop_process(process)
        return receipt(
            operation,
            {"execution_status": "timed_out" if isinstance(exc, TimeoutError) else "output_limit"},
            target=target,
            kind=kind,
            execution="partial",
            side_effect="possible" if write else "none",
            error_code="worker_timeout" if isinstance(exc, TimeoutError) else "output_limit",
            files=files,
        )
    except asyncio.CancelledError:
        await asyncio.shield(stop_process(process))
        raise
    except OSError:
        await stop_process(process)
        return failure(
            "工具输出通道中断；执行结果待核对，不能自动重试。",
            "worker_output_interrupted",
            execution="unknown",
            side_effect="possible" if write else "none",
        )
    try:
        decoded = output.decode("utf-8")
        try:
            result = json.loads(decoded)
        except ValueError:
            # V2 discovery emits one JSON event per line, rather than one envelope.
            result = {"events": [json.loads(line) for line in decoded.splitlines() if line.strip()]}
        if not isinstance(result, dict):
            raise ValueError
    except (UnicodeError, ValueError):
        return receipt(
            operation,
            {"returncode": process.returncode, "execution_status": "invalid_output_suppressed"},
            target=target,
            kind=kind,
            execution="unknown" if write else "completed",
            side_effect="possible" if write else "none",
            error_code="invalid_worker_output",
            files=files,
        )
    result = sanitized(result, secrets)
    if script == "evidence.py":
        result.pop("output", None)
        result["output_location"] = "artemis_task_attachments"
    status = result.get("execution_status", "")
    not_started = status in {
        "invalid_arguments",
        "tool_missing",
        "configuration_error",
        "launch_failed",
    }
    attempted = bool(
        result.get("write_executed")
        or result.get("write_may_have_started")
        or result.get("may_have_started")
    )
    verified_mismatch = (
        script == "device_adapter.py"
        and arguments[0] == "verify-version"
        and result.get("status") in {"FAIL", "INCOMPLETE"}
    )
    error = process.returncode != 0 and not verified_mismatch
    execution = (
        "not_started"
        if not_started
        else "unknown"
        if status in {"cleanup_unconfirmed", "unknown"}
        else "partial"
        if status in {"timed_out", "output_limit"} or (error and attempted)
        else "completed"
    )
    side_effect = (
        "confirmed"
        if write and result.get("write_executed")
        else "possible"
        if write and attempted
        else "none"
    )
    if write and script != "device_adapter.py" and not not_started:
        side_effect = "possible" if error else "confirmed"
    return receipt(
        operation,
        {"returncode": process.returncode, "tool_result": result, "stderr_suppressed": True},
        target=target,
        kind=kind,
        execution=execution,
        side_effect=side_effect,
        error_code=str(result.get("error_code") or "tool_execution_failed") if error else "",
        files=files,
    )