from __future__ import annotations

import io
import json
import os
import queue
import shutil
import subprocess
import threading
from multiprocessing.context import AuthenticationError
from multiprocessing.connection import Client

from ._errors import (
    DlpConfigurationError,
    DlpHelperStartError,
    DlpProtocolError,
    DlpSessionBusyError,
    DlpTransportError,
)
from ._helper import (
    CONTROL_LIMIT,
    HANDSHAKE_PREFIX,
    PROTOCOL_VERSION,
    render_helper_source,
)
from ._raw import _PipeRawIO

_SHUTDOWN_TIMEOUT = 5


def _locate_python() -> list[str] | None:
    for variable in ("DLP_IO_PYTHON", "DATAPACKER_DLP_PYTHON"):
        value = os.environ.get(variable)
        if value:
            return [value]
    executable = shutil.which("python")
    if executable:
        return [executable]
    if shutil.which("py"):
        return ["py", "-3"]
    return None


def _helper_process_options(os_name=None) -> dict:
    if (os_name or os.name) == "nt":
        return {"creationflags": subprocess.CREATE_NO_WINDOW}
    return {}


def _stop_process(process) -> None:
    if process.poll() is not None:
        return
    process.terminate()
    try:
        process.wait(timeout=5)
    except subprocess.TimeoutExpired:
        process.kill()
        try:
            process.wait(timeout=5)
        except subprocess.TimeoutExpired:
            pass


def _close_process_pipes(process) -> None:
    for name in ("stdin", "stdout", "stderr"):
        stream = getattr(process, name, None)
        if stream is None:
            continue
        try:
            stream.close()
        except OSError:
            pass


class DlpSession:
    """One authenticated, short-lived approved-Python helper session."""

    def __init__(self, process, pipe_name: str, authkey: bytes) -> None:
        self._proc = process
        self.pipe_name = pipe_name
        self._authkey = authkey
        self._request_lock = threading.Lock()
        self._stream_lock = threading.Lock()
        self._state_lock = threading.RLock()
        self._active_raw = None
        self._closed = False

    @classmethod
    def start(cls, *, python_cmd=None, timeout=60):
        if os.name != "nt":
            raise DlpConfigurationError("dlp-io bridge reads require Windows AF_PIPE")
        if timeout <= 0:
            raise DlpConfigurationError("startup timeout must be positive")
        command = list(python_cmd) if python_cmd else _locate_python()
        if not command:
            raise DlpHelperStartError(
                "no approved Python found; configure DLP_IO_PYTHON"
            )
        authkey = os.urandom(32)
        source = render_helper_source(authkey)
        try:
            process = subprocess.Popen(
                command + ["-u", "-"],
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                **_helper_process_options(),
            )
        except OSError as exc:
            raise DlpHelperStartError(
                "failed to start approved Python %r: %s" % (command, exc)
            ) from exc
        try:
            process.stdin.write(source.encode("utf-8"))
            process.stdin.close()
        except (BrokenPipeError, OSError) as exc:
            _stop_process(process)
            raise DlpHelperStartError("failed to inject helper source: %s" % exc) from exc

        result = queue.Queue()

        def read_handshake() -> None:
            try:
                result.put(process.stdout.readline())
            except Exception as exc:
                result.put(exc)

        thread = threading.Thread(target=read_handshake, daemon=True)
        thread.start()
        try:
            line = result.get(timeout=timeout)
        except queue.Empty as exc:
            _stop_process(process)
            raise DlpHelperStartError(
                "timed out waiting for helper handshake after %s seconds" % timeout
            ) from exc
        if isinstance(line, Exception):
            _stop_process(process)
            raise DlpHelperStartError("failed to read helper handshake: %s" % line)
        decoded = line.decode("utf-8", "replace").strip()
        if not decoded.startswith(HANDSHAKE_PREFIX):
            _stop_process(process)
            stderr = process.stderr.read().decode("utf-8", "replace").strip()
            raise DlpHelperStartError(
                "invalid helper handshake %r; stderr: %s" % (decoded, stderr)
            )
        return cls(process, decoded[len(HANDSHAKE_PREFIX):], authkey)

    def _ensure_open(self) -> None:
        if self._closed:
            raise RuntimeError("dlp-io session is closed")

    def _validate_header(self, header) -> dict:
        if not isinstance(header, dict):
            raise DlpProtocolError("helper response must be a JSON object")
        if header.get("version") != PROTOCOL_VERSION:
            raise DlpProtocolError("helper response has an unsupported version")
        if not isinstance(header.get("ok"), bool):
            raise DlpProtocolError("helper response is missing boolean ok")
        return header

    def _request(self, request):
        with self._request_lock:
            try:
                encoded = json.dumps(request).encode("utf-8")
            except (TypeError, ValueError) as exc:
                raise DlpProtocolError("helper request is not JSON serializable") from exc
            try:
                conn = Client(self.pipe_name, family="AF_PIPE", authkey=self._authkey)
            except AuthenticationError as exc:
                raise DlpTransportError("failed to authenticate helper pipe") from exc
            except (EOFError, OSError) as exc:
                raise DlpTransportError("failed to connect to helper pipe") from exc
            try:
                try:
                    conn.send_bytes(encoded)
                except (EOFError, OSError) as exc:
                    raise DlpTransportError("failed to send helper request") from exc
                try:
                    payload = conn.recv_bytes(maxlength=CONTROL_LIMIT)
                except (EOFError, OSError) as exc:
                    raise DlpTransportError("failed to receive helper response") from exc
                try:
                    header = json.loads(payload.decode("utf-8"))
                except (UnicodeDecodeError, json.JSONDecodeError) as exc:
                    raise DlpProtocolError("helper response is not valid UTF-8 JSON") from exc
                return conn, self._validate_header(header)
            except Exception:
                conn.close()
                raise

    def ping(self) -> dict:
        with self._state_lock:
            self._ensure_open()
        conn, header = self._request({"version": PROTOCOL_VERSION, "op": "ping"})
        conn.close()
        if not header["ok"]:
            raise DlpProtocolError(header.get("error", "ping failed"))
        return header

    def _open_conn(self, path: str, offset: int = 0):
        conn, header = self._request(
            {
                "version": PROTOCOL_VERSION,
                "op": "read",
                "path": path,
                "offset": offset,
            }
        )
        if not header["ok"]:
            conn.close()
            raise OSError(header.get("error", "helper read failed"))
        size = header.get("size")
        if isinstance(size, bool) or not isinstance(size, int) or size < 0:
            conn.close()
            raise DlpProtocolError("helper returned an invalid file size")
        return conn, size

    def _release_stream(self, raw) -> None:
        with self._state_lock:
            if self._active_raw is raw:
                self._active_raw = None
                self._stream_lock.release()

    def _reopen_stream(self, path: str, offset: int):
        with self._state_lock:
            self._ensure_open()
        return self._open_conn(path, offset)

    def open_raw(self, path):
        with self._state_lock:
            self._ensure_open()
        if not self._stream_lock.acquire(blocking=False):
            raise DlpSessionBusyError("dlp-io session already has an active stream")
        normalized = os.fsdecode(os.fspath(path))
        try:
            conn, size = self._open_conn(normalized)
            raw = _PipeRawIO(
                conn=conn,
                size=size,
                path=normalized,
                offset=0,
                reopen=self._reopen_stream,
                release=self._release_stream,
            )
            with self._state_lock:
                if self._closed:
                    raw.close()
                    raise RuntimeError("dlp-io session is closed")
                self._active_raw = raw
            return raw
        except Exception:
            self._stream_lock.release()
            raise

    def open(
        self,
        file,
        mode="r",
        buffering=-1,
        encoding=None,
        errors=None,
        newline=None,
    ):
        from ._api import _validate_path_mode, _wrap_reader

        buffering = _validate_path_mode(mode, buffering, encoding, errors, newline)
        if "r" not in mode or "+" in mode:
            raise io.UnsupportedOperation("DlpSession.open supports read-only modes")

        return _wrap_reader(
            self.open_raw(file), mode, buffering, encoding, errors, newline
        )

    def open_read(self, path):
        """Compatibility alias returning a buffered binary reader."""
        return self.open(path, "rb")

    def close(self) -> None:
        with self._state_lock:
            if self._closed:
                return
            self._closed = True
            active = self._active_raw
        if active is not None:
            active.close()

        def request_quit() -> None:
            try:
                conn, header = self._request(
                    {"version": PROTOCOL_VERSION, "op": "quit"}
                )
                conn.close()
                if not header["ok"]:
                    raise DlpProtocolError(header.get("error", "helper quit failed"))
            except Exception:
                pass

        quit_thread = threading.Thread(target=request_quit, daemon=True)
        quit_thread.start()
        quit_thread.join(timeout=_SHUTDOWN_TIMEOUT)
        try:
            if quit_thread.is_alive():
                _stop_process(self._proc)
                quit_thread.join(timeout=_SHUTDOWN_TIMEOUT)
            else:
                try:
                    self._proc.wait(timeout=_SHUTDOWN_TIMEOUT)
                except subprocess.TimeoutExpired:
                    _stop_process(self._proc)
        finally:
            _close_process_pipes(self._proc)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.close()
        return False