import builtins
import gc
import json
import logging
import multiprocessing
import os
import select
import signal
import socket
import struct
from collections import namedtuple
from typing import Any, Callable, Optional, Union

from pyroute2 import config
from pyroute2.common import USE_DEFAULT_TIMEOUT
from pyroute2.netlink import exceptions as pyroute2_exceptions

log = logging.getLogger(__name__)

ChildProcessReturnValue = namedtuple(
    'ChildProcessReturnValue', ('payload', 'fds')
)
ChildFuncReturnType = Union[None, ChildProcessReturnValue, bytearray, bytes]


def wrapper(
    ctrl: socket.socket,
    func: Callable[..., ChildFuncReturnType],
    argv: list[Any],
) -> None:
    '''Child function wrapper.

    This code will be executed in the child process after running fork().
    The internal data structures might be damaged, so the function `func`
    must be as simple as possible, and use only local variables of simple
    types.

    The garbage collector should be disabled as well to minimize possible
    deadlocks.

    If process doesn't response in time, it will get killed.
    '''
    gc.disable()
    payload: bytes = b''
    ret_data: bytes = b''
    fds: list[int] = []
    sockets: list[socket.socket] = []
    try:
        ret = func(*argv)
        if isinstance(ret, bytes):
            ret_data = ret
        if isinstance(ret, ChildProcessReturnValue):
            ret_data, sockets = ret
        if isinstance(ret_data, bytearray):
            ret_data = bytes(ret_data)
        if not isinstance(ret_data, bytes):
            raise TypeError('return values not supported')
        payload = struct.pack('B', 2) + ret_data
        if sockets:
            fds = [x.fileno() for x in sockets]
    except Exception as e:
        payload = struct.pack('B', 1) + json.dumps(
            {'exception': e.__class__.__name__, 'options': e.args}
        ).encode('utf-8')
        fds = []
    finally:
        socket.send_fds(ctrl, [payload], fds, len(fds))


class ChildProcess:
    ctrl_r: socket.socket
    ctrl_w: socket.socket

    def __init__(
        self, target: Callable[..., ChildFuncReturnType], args: list[Any]
    ):
        self._mode: str = config.child_process_mode
        self._target: Callable[..., ChildFuncReturnType] = target
        self._args: list[Any] = args
        self._proc: Optional[multiprocessing.Process] = None
        self._running: bool = False
        self._exitcode: Optional[int] = None
        self._pid: Optional[int] = None

    def close(self):
        self.stop()

    def __enter__(self):
        self.run()
        return self

    def __exit__(self, *_):
        self.close()

    @property
    def mode(self):
        return self._mode

    @property
    def pid(self):
        return self._pid

    def communicate(
        self, timeout: int = USE_DEFAULT_TIMEOUT
    ) -> tuple[bytes, list[int]]:
        '''Communicate with the child process.

        Raises:

        * struct.error -- error unpacking response from the child process
        * OSError -- OS level error communicating with the child process
        * TimeoutError -- the child process is alive, but doesn't response
        * RuntimeError -- the child process is dead or not started
        * TypeError -- error loading propagated exception
        '''
        if not self._running:
            raise RuntimeError('child process not started yet')
        if timeout == USE_DEFAULT_TIMEOUT:
            timeout = config.default_communicate_timeout
        rl, _, _ = select.select([self.ctrl_r], [], [], timeout)
        if not len(rl):
            # no data received within timeout
            # 1. the child process is killed
            # 2. the child process is stuck
            #
            # So first check the process, if it is killed, raise
            # RuntimeError, otherwise raise TimeoutError.
            #
            c_pid, w_status = os.waitpid(self.pid, os.WNOHANG)
            if c_pid != 0:
                # process is dead
                self._exitcode = os.waitstatus_to_exitcode(w_status)
                raise RuntimeError(
                    f'child process is dead, status: {self.exitcode}'
                )
            # process is alive, but stuck, kill it
            self.stop(kill=True, reason='no response from the child')
            raise TimeoutError(f'no response from the child pid {self.pid}')

        ret_data = b''
        # raises OSError
        (raw_data, fds, _, _) = socket.recv_fds(self.ctrl_r, 1024, 1)
        # get the return type
        # raises struct.error
        (ret_type,) = struct.unpack('B', raw_data[:1])
        raw_data = raw_data[1:]
        if ret_type == 1:
            # exception
            payload = json.loads(raw_data.decode('utf-8'))
            if not set(payload.keys()) == set(('exception', 'options')):
                raise TypeError('error loading child exception')
            if payload.get('exception') is not None:
                error_class = getattr(builtins, payload['exception'], None)
                if error_class is None:
                    error_class = getattr(
                        pyroute2_exceptions, payload['exception'], None
                    )
                if error_class is None:
                    error_class = Exception
                if not issubclass(error_class, Exception):
                    raise TypeError('error loading child error')
                raise error_class(*payload['options'])
        elif ret_type == 2:
            # raw_data
            ret_data = raw_data
        return ret_data, fds

    def get_data(self, timeout: int = USE_DEFAULT_TIMEOUT) -> bytes:
        return self.communicate(timeout)[0]

    def get_fds(self, timeout: int = USE_DEFAULT_TIMEOUT) -> list[int]:
        return self.communicate(timeout)[1]

    @property
    def proc(self) -> multiprocessing.Process:
        if self._proc is None:
            raise RuntimeError('not started')
        return self._proc

    @proc.setter
    def proc(self, value: Optional[multiprocessing.Process]) -> None:
        self._proc = value

    def _unsupported(self) -> Exception:
        return TypeError('unsupported mode')

    def run(self) -> None:
        if self._running:
            return
        self._running = True
        self.ctrl_r, self.ctrl_w = socket.socketpair(
            socket.AF_UNIX, socket.SOCK_DGRAM
        )
        if self.mode == 'fork':
            self._pid = os.fork()
            if self._pid == 0:
                wrapper(self.ctrl_w, self._target, self._args)
                os._exit(0)
        elif self.mode == 'mp':
            self.proc = multiprocessing.Process(
                target=wrapper, args=[self.ctrl_w, self._target, self._args]
            )
            self.proc.start()
            self._pid = self.proc.pid
        else:
            raise self._unsupported()

    def stop(self, kill: bool = False, reason: Optional[str] = None) -> None:
        if not self._running:
            return
        self._running = False
        self.ctrl_r.close()
        self.ctrl_w.close()
        if config.force_gc:
            gc.collect()
        if self.mode == 'fork':
            try:
                if kill:
                    os.kill(self.pid, signal.SIGKILL)
                else:
                    os.kill(self.pid, signal.SIGTERM)
                _, status = os.waitpid(self.pid, 0)
                self._exitcode = os.waitstatus_to_exitcode(status)
            except ProcessLookupError:
                # the child process has already exited
                pass
        elif self.mode == 'mp':
            if kill:
                self.proc.kill()
            else:
                self.proc.terminate()
            self.proc.join()
        else:
            raise self._unsupported()
        if reason is not None:
            log.warning(reason)

    @property
    def exitcode(self) -> Optional[int]:
        if self.mode == 'fork':
            return self._exitcode
        elif self.mode == 'mp':
            return self.proc.exitcode
        else:
            raise self._unsupported()
