"""Shared Windows process ownership for the plugin's fixed external tools."""

from __future__ import annotations

import os
import subprocess
import threading
import time


class ManagedProcessError(OSError):
    """Stable error with execution facts; never includes tool arguments/output."""

    def __init__(self, code, *, may_have_started=False):
        super().__init__(code)
        self.error_code = code
        self.may_have_started = may_have_started


class WindowsJob:
    def __init__(self):
        import ctypes
        from ctypes import wintypes

        class IO_COUNTERS(ctypes.Structure):
            _fields_ = [
                (n, ctypes.c_ulonglong)
                for n in (
                    "ReadOperationCount",
                    "WriteOperationCount",
                    "OtherOperationCount",
                    "ReadTransferCount",
                    "WriteTransferCount",
                    "OtherTransferCount",
                )
            ]

        class BASIC_LIMIT(ctypes.Structure):
            _fields_ = [
                ("PerProcessUserTimeLimit", ctypes.c_longlong),
                ("PerJobUserTimeLimit", ctypes.c_longlong),
                ("LimitFlags", wintypes.DWORD),
                ("MinimumWorkingSetSize", ctypes.c_size_t),
                ("MaximumWorkingSetSize", ctypes.c_size_t),
                ("ActiveProcessLimit", wintypes.DWORD),
                ("Affinity", ctypes.c_size_t),
                ("PriorityClass", wintypes.DWORD),
                ("SchedulingClass", wintypes.DWORD),
            ]

        class EXTENDED_LIMIT(ctypes.Structure):
            _fields_ = [
                ("BasicLimitInformation", BASIC_LIMIT),
                ("IoInfo", IO_COUNTERS),
                ("ProcessMemoryLimit", ctypes.c_size_t),
                ("JobMemoryLimit", ctypes.c_size_t),
                ("PeakProcessMemoryUsed", ctypes.c_size_t),
                ("PeakJobMemoryUsed", ctypes.c_size_t),
            ]

        class ACCOUNTING(ctypes.Structure):
            _fields_ = [
                ("TotalUserTime", ctypes.c_longlong),
                ("TotalKernelTime", ctypes.c_longlong),
                ("ThisPeriodTotalUserTime", ctypes.c_longlong),
                ("ThisPeriodTotalKernelTime", ctypes.c_longlong),
                ("TotalPageFaultCount", wintypes.DWORD),
                ("TotalProcesses", wintypes.DWORD),
                ("ActiveProcesses", wintypes.DWORD),
                ("TotalTerminatedProcesses", wintypes.DWORD),
            ]

        self.ctypes = ctypes
        self.accounting = ACCOUNTING
        self.kernel = k = ctypes.WinDLL("kernel32", use_last_error=True)
        self.ntdll = ctypes.WinDLL("ntdll")
        k.CreateJobObjectW.argtypes = [wintypes.LPVOID, wintypes.LPCWSTR]
        k.CreateJobObjectW.restype = wintypes.HANDLE
        k.SetInformationJobObject.argtypes = [
            wintypes.HANDLE,
            ctypes.c_int,
            wintypes.LPVOID,
            wintypes.DWORD,
        ]
        k.SetInformationJobObject.restype = wintypes.BOOL
        k.AssignProcessToJobObject.argtypes = [wintypes.HANDLE, wintypes.HANDLE]
        k.AssignProcessToJobObject.restype = wintypes.BOOL
        k.TerminateJobObject.argtypes = [wintypes.HANDLE, wintypes.UINT]
        k.TerminateJobObject.restype = wintypes.BOOL
        k.QueryInformationJobObject.argtypes = [
            wintypes.HANDLE,
            ctypes.c_int,
            wintypes.LPVOID,
            wintypes.DWORD,
            wintypes.LPVOID,
        ]
        k.QueryInformationJobObject.restype = wintypes.BOOL
        k.CloseHandle.argtypes = [wintypes.HANDLE]
        k.CloseHandle.restype = wintypes.BOOL
        self.ntdll.NtResumeProcess.argtypes = [wintypes.HANDLE]
        self.ntdll.NtResumeProcess.restype = ctypes.c_long
        self.handle = k.CreateJobObjectW(None, None)
        if not self.handle:
            raise ManagedProcessError("process_ownership_failed")
        limits = EXTENDED_LIMIT()
        limits.BasicLimitInformation.LimitFlags = (
            0x2000  # KILL_ON_JOB_CLOSE; handle is not inherited.
        )
        if not k.SetInformationJobObject(
            self.handle, 9, ctypes.byref(limits), ctypes.sizeof(limits)
        ):
            self.close()
            raise ManagedProcessError("process_ownership_failed")

    def assign(self, process):
        if not self.kernel.AssignProcessToJobObject(self.handle, int(process._handle)):
            raise ManagedProcessError("process_ownership_failed")

    def resume(self, process):
        if self.ntdll.NtResumeProcess(int(process._handle)) != 0:
            raise ManagedProcessError(
                "process_resume_unconfirmed", may_have_started=True
            )

    def stop(self):
        c = self.ctypes
        if not self.kernel.TerminateJobObject(self.handle, 1):
            raise OSError("job_termination_unconfirmed")
        deadline = time.monotonic() + 10
        while True:
            account = self.accounting()
            if not self.kernel.QueryInformationJobObject(
                self.handle, 1, c.byref(account), c.sizeof(account), None
            ):
                raise OSError("job_cleanup_unconfirmed")
            if account.ActiveProcesses == 0:
                return
            if time.monotonic() >= deadline:
                raise OSError("job_cleanup_timeout")
            time.sleep(0.02)

    def close(self):
        if self.handle:
            handle, self.handle = self.handle, None
            if not self.kernel.CloseHandle(handle):
                raise OSError("job_handle_close_unconfirmed")


def run(argv, *, timeout=None, check=False, max_output_bytes=4 * 1024 * 1024, **kwargs):
    """Contain before resuming; bound pipes; reap the tree on every exit path."""
    if os.name != "nt":
        raise ManagedProcessError("process_ownership_requires_windows")
    text = bool(
        kwargs.get("text")
        or kwargs.get("universal_newlines")
        or kwargs.get("encoding")
        or kwargs.get("errors")
    )
    process = None
    assigned = started = False
    job = WindowsJob()
    buffers = {}
    readers = []
    output_limit = threading.Event()
    read_error = threading.Event()
    lock = threading.Lock()
    size = 0

    def read_pipe(name, pipe):
        nonlocal size
        buffers[name] = bytearray()
        try:
            while data := os.read(pipe.fileno(), 65536):
                with lock:
                    size += len(data)
                    if size > max_output_bytes:
                        output_limit.set()
                    elif not output_limit.is_set():
                        buffers[name].extend(data)
        except OSError:
            read_error.set()

    try:
        kwargs["creationflags"] = (
            kwargs.get("creationflags", 0) | subprocess.CREATE_NO_WINDOW | 0x4
        )
        try:
            process = subprocess.Popen(argv, **kwargs)
        except OSError:
            raise ManagedProcessError("process_launch_failed") from None
        job.assign(process)
        assigned = True
        started = True  # Set before resume: failure cannot prove the tool did not run.
        job.resume(process)
        for name in ("stdout", "stderr"):
            if pipe := getattr(process, name):
                thread = threading.Thread(
                    target=read_pipe, args=(name, pipe), daemon=True
                )
                thread.start()
                readers.append(thread)
        deadline = time.monotonic() + timeout if timeout is not None else None
        while process.poll() is None:
            if output_limit.is_set() or read_error.is_set():
                break
            remaining = deadline - time.monotonic() if deadline is not None else 0.05
            if remaining <= 0:
                raise subprocess.TimeoutExpired(argv, timeout)
            try:
                process.wait(timeout=min(0.05, remaining))
            except subprocess.TimeoutExpired:
                continue
    finally:
        cleanup_failed = False
        try:
            if process is not None:
                if assigned:
                    job.stop()
                elif process.poll() is None:
                    process.kill()  # Still suspended; no unowned descendants can exist.
                process.wait(timeout=10)
        except (OSError, subprocess.TimeoutExpired):
            cleanup_failed = True
        finally:
            try:
                job.close()
            except OSError:
                cleanup_failed = True
            for thread in readers:
                thread.join(timeout=2)
                cleanup_failed |= thread.is_alive()
            for name in ("stdin", "stdout", "stderr"):
                if (pipe := getattr(process, name, None)) and not cleanup_failed:
                    try:
                        pipe.close()
                    except OSError:
                        cleanup_failed = True
        if cleanup_failed:
            raise ManagedProcessError(
                "process_cleanup_unconfirmed", may_have_started=started
            ) from None
    if output_limit.is_set():
        raise ManagedProcessError("process_output_limit", may_have_started=True)
    if read_error.is_set():
        raise ManagedProcessError("process_output_interrupted", may_have_started=True)
    output = {}
    for name in ("stdout", "stderr"):
        data = bytes(buffers[name]) if name in buffers else None
        if text and data is not None:
            pipe = getattr(process, name)
            data = (
                data.decode(pipe.encoding, pipe.errors)
                .replace("\r\n", "\n")
                .replace("\r", "\n")
            )
        output[name] = data
    result = subprocess.CompletedProcess(
        argv, process.returncode, output["stdout"], output["stderr"]
    )
    if check:
        result.check_returncode()
    return result