"""Plugin author helpers; no scheduler, device actions or implicit retries.

Use snapshot_file before approving a write and consume its immutable bytes, not the
original pathname. run_process owns the tool and its descendants on every exit.
"""

from __future__ import annotations

import asyncio
from collections.abc import Iterator, Mapping, Sequence
from dataclasses import dataclass, field
import hashlib
import json
import math
import os
from pathlib import Path
import re
import signal
import stat
import subprocess
import sys
import tempfile

from artemis.plugins.recording_process import ProcessJob

MAX_INPUT_BYTES = 64 * 1024 * 1024
DEFAULT_OUTPUT_BYTES = 4 * 1024 * 1024


class FileVerificationError(ValueError):
    """The input cannot be bound to its explicitly approved digest."""


@dataclass(frozen=True)
class VerifiedFile:
    """Bounded, immutable input. No original path is exposed for reopening."""

    data: bytes = field(repr=False)
    sha256: str = field(init=False)

    def __post_init__(self):
        if not isinstance(self.data, bytes):
            raise TypeError("快照内容必须是不可变 bytes")
        object.__setattr__(self, "sha256", hashlib.sha256(self.data).hexdigest())

    @property
    def size(self) -> int:
        return len(self.data)

    def chunks(self, size: int = 256 * 1024) -> Iterator[memoryview]:
        if type(size) is not int or size <= 0:
            raise ValueError("分块大小必须是正整数")
        content = memoryview(self.data)
        for offset in range(0, len(content), size):
            yield content[offset : offset + size]


def snapshot_file(
    path: Path | str, expected_sha256: str, *, max_bytes: int = MAX_INPUT_BYTES
) -> VerifiedFile:
    """Copy once, then verify the exact bytes that will be consumed by the caller.

    The source must already be authorized by the plugin. This is not a path or
    permission sandbox. Memory use is bounded by max_bytes plus one read byte.
    """
    if type(max_bytes) is not int or max_bytes <= 0:
        raise ValueError("文件大小上限必须是正整数")
    if not isinstance(expected_sha256, str) or not re.fullmatch(
        r"[a-fA-F0-9]{64}", expected_sha256
    ):
        raise FileVerificationError("必须明确提供已批准的 SHA256")
    source = Path(path)
    if source.is_symlink() or not stat.S_ISREG(source.lstat().st_mode):
        raise FileVerificationError("输入必须是普通文件，不能是链接或设备")
    flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) | getattr(os, "O_NONBLOCK", 0)
    descriptor = os.open(source, flags)
    with os.fdopen(descriptor, "rb") as handle:
        info = os.fstat(handle.fileno())
        if not stat.S_ISREG(info.st_mode) or info.st_size > max_bytes:
            raise FileVerificationError("输入不是普通文件或超过大小上限")
        data = handle.read(max_bytes + 1)
    if len(data) > max_bytes:
        raise FileVerificationError("输入超过大小上限")
    snapshot = VerifiedFile(data)
    if snapshot.sha256 != expected_sha256.lower():
        raise FileVerificationError("输入 SHA256 与批准值不一致；不得开始写操作")
    return snapshot


class ManagedProcessError(RuntimeError):
    """Stable failure metadata. Do not infer a business verdict or retry a write."""

    def __init__(self, code: str, *, may_have_started: bool = True):
        super().__init__(code)
        self.error_code = code
        self.may_have_started = may_have_started


class ProcessTimedOut(ManagedProcessError, TimeoutError):
    pass


@dataclass(frozen=True)
class ProcessResult:
    returncode: int
    stdout: bytes = field(repr=False)
    stderr: bytes = field(repr=False)


async def _cleanup(process: asyncio.subprocess.Process, job: ProcessJob | None) -> None:
    error = None
    try:
        if job is not None:
            job.terminate()
            async with asyncio.timeout(5):
                while job.active_processes():
                    await asyncio.sleep(0.02)
        elif os.name != "nt":
            # The bootstrap stays alive until cleanup, pinning the owned PGID.
            try:
                os.killpg(process.pid, signal.SIGKILL)
            except ProcessLookupError:
                pass
    except (OSError, TimeoutError) as exc:
        error = exc
    finally:
        if job is not None:
            job.close()
        if process.returncode is None:
            try:
                process.kill()
            except ProcessLookupError:
                pass
    # Process.wait also waits for pipe EOF. A paused, full output reader must be
    # drained separately; it cannot serve as evidence that the OS process lives.
    try:
        async with asyncio.timeout(5):
            while process.returncode is None:
                await asyncio.sleep(0.02)
    except TimeoutError as exc:
        error = exc
    if error is not None:
        raise ManagedProcessError("process_cleanup_unconfirmed") from error


async def _finish_cleanup(task: asyncio.Task[None]) -> None:
    cancelled = False
    while not task.done():
        try:
            await asyncio.shield(task)
        except asyncio.CancelledError:
            cancelled = True
    task.result()
    if cancelled:
        raise asyncio.CancelledError


async def _read_output(stream: asyncio.StreamReader, consumed: list[int], limit: int) -> bytes:
    chunks = []
    while chunk := await stream.read(65536):
        consumed[0] += len(chunk)
        if consumed[0] > limit:
            raise ManagedProcessError("process_output_limit")
        chunks.append(chunk)
    return b"".join(chunks)


async def _drain(stream: asyncio.StreamReader | None) -> None:
    if stream is not None:
        while await stream.read(65536):
            pass


async def _wait_result(process: asyncio.subprocess.Process, path: Path) -> int:
    while not path.exists():
        if process.returncode is not None:
            raise ManagedProcessError("process_worker_interrupted")
        await asyncio.sleep(0.02)
    if path.stat().st_size > 4096:
        raise ManagedProcessError("process_invalid_status")
    try:
        result = json.loads(path.read_text("utf-8"))
        if not isinstance(result, dict):
            raise ValueError
    except (OSError, ValueError) as exc:
        raise ManagedProcessError("process_invalid_status") from exc
    if result.get("launch_failed") is True:
        raise ManagedProcessError("process_launch_failed", may_have_started=False)
    if type(result.get("returncode")) is not int:
        raise ManagedProcessError("process_invalid_status")
    return result["returncode"]


async def run_process(
    argv: Sequence[str],
    *,
    cwd: Path | str,
    timeout: float,
    env: Mapping[str, str] | None = None,
    input_data: bytes | None = None,
    max_output_bytes: int = DEFAULT_OUTPUT_BYTES,
) -> ProcessResult:
    """Run one fixed argv without a shell; reclaim descendants even on normal exit.

    Windows assigns a waiting bootstrap to a Job Object *before* it can start the
    tool. POSIX uses a dedicated session. Cancellation completes cleanup before
    propagating. Inputs are at most 64 MiB; stdout+stderr share a bounded budget.
    Return codes/output are tool results, never a test verdict. No retries occur.
    """
    if (
        not isinstance(argv, Sequence)
        or isinstance(argv, (str, bytes))
        or not argv
        or any(not isinstance(arg, str) or "\0" in arg for arg in argv)
    ):
        raise ValueError("必须提供固定的 argv 字符串列表，不能使用 shell 命令文本")
    if not argv[0] or isinstance(timeout, bool) or not math.isfinite(timeout) or timeout <= 0:
        raise ValueError("启动程序和有限正数超时是必填项")
    if type(max_output_bytes) is not int or max_output_bytes <= 0:
        raise ValueError("输出上限必须是正整数")
    if input_data is not None and (
        not isinstance(input_data, bytes) or len(input_data) > MAX_INPUT_BYTES
    ):
        raise ValueError("标准输入必须是至多 64 MiB 的 bytes")
    directory = Path(cwd).resolve(strict=True)
    if not directory.is_dir():
        raise ValueError("工作目录必须是已存在的目录")
    control = (
        json.dumps(
            {
                "argv": list(argv),
                "has_input": input_data is not None,
                "input_size": len(input_data or b""),
            }
        ).encode()
        + b"\n"
    )
    if len(control) > 1024 * 1024:
        raise ValueError("启动参数超过大小上限")
    job = None
    process: asyncio.subprocess.Process | None = None
    cleanup: asyncio.Task[None] | None = None
    tasks: list[asyncio.Task] = []

    async def reclaim():
        nonlocal cleanup
        if process is None:
            if job is not None:
                job.close()
            return
        if cleanup is None:
            cleanup = asyncio.create_task(_cleanup(process, job))
        await _finish_cleanup(cleanup)

    with tempfile.TemporaryDirectory(prefix="artemis-plugin-process-") as temporary:
        status = Path(temporary) / "status.json"
        try:
            job = ProcessJob() if os.name == "nt" else None
        except OSError as exc:
            raise ManagedProcessError("process_ownership_failed", may_have_started=False) from exc
        launch = asyncio.create_task(
            asyncio.create_subprocess_exec(
                sys.executable,
                "-I",
                "-B",
                str(Path(__file__).with_name("_process_worker.py")),
                str(status),
                str(os.getpid()),
                cwd=str(directory),
                env=dict(env) if env is not None else None,
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                creationflags=subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0,
                start_new_session=os.name != "nt",
            )
        )
        try:
            cancelled = False
            while process is None:
                try:
                    process = await asyncio.shield(launch)
                except asyncio.CancelledError:
                    cancelled = True
                except OSError as exc:
                    raise ManagedProcessError(
                        "process_launch_failed", may_have_started=False
                    ) from exc
            if cancelled:
                raise asyncio.CancelledError
            if job is not None:
                try:
                    job.assign(process.pid)
                except OSError as exc:
                    raise ManagedProcessError(
                        "process_ownership_failed", may_have_started=False
                    ) from exc
            stdin, stdout, stderr = process.stdin, process.stdout, process.stderr
            assert stdin is not None and stdout is not None and stderr is not None
            consumed = [0]

            async def feed():
                try:
                    stdin.write(control)
                    await stdin.drain()
                    if input_data:
                        stdin.write(input_data)
                        await stdin.drain()
                except (BrokenPipeError, ConnectionResetError):
                    pass
                finally:
                    stdin.close()

            async def wait_tool():
                code = await _wait_result(process, status)
                await reclaim()
                return code

            tasks = [
                asyncio.create_task(feed()),
                asyncio.create_task(_read_output(stdout, consumed, max_output_bytes)),
                asyncio.create_task(_read_output(stderr, consumed, max_output_bytes)),
                asyncio.create_task(wait_tool()),
            ]
            try:
                _, stdout, stderr, code = await asyncio.wait_for(asyncio.gather(*tasks), timeout)
            except TimeoutError as exc:
                raise ProcessTimedOut("process_timeout") from exc
            return ProcessResult(code, stdout, stderr)
        finally:
            for task in tasks:
                task.cancel()

            async def finalize():
                try:
                    await reclaim()
                finally:
                    await asyncio.gather(*tasks, return_exceptions=True)
                    if process is not None:
                        try:
                            async with asyncio.timeout(5):
                                await asyncio.gather(_drain(process.stdout), _drain(process.stderr))
                                await process.wait()
                        except TimeoutError as exc:
                            raise ManagedProcessError("process_cleanup_unconfirmed") from exc

            await _finish_cleanup(asyncio.create_task(finalize()))