import asyncio
import inspect
import json
import logging
from multiprocessing import Process, Pipe, get_context
from multiprocessing.connection import Connection
from typing import Optional
from types import FunctionType


logger = logging.getLogger()


class ProcessConnector:
    """
    A process connector classes.
    """
    def __init__(self, function):
        self.parent_con, self.child_con = None, None
        self.pid = 0
        self.proc = None
        self.raw_func = function
        self.is_async = inspect.iscoroutinefunction(function)
        self.result = None

    def kill(self):
        self.proc.kill()

    def execute_in_process(self, conn: Connection):
        try:
            input_params = conn.recv()
            args = input_params.get("args", [])
            kwargs = input_params.get("kwargs", {})
            result = self.raw_func(*args, **kwargs)
            try:
                json.dumps(result)
            except Exception as e:
                logger.debug(e)
                result = {}
            conn.send(result)
        except Exception as error:
            logger.debug(error)
            raise error
        finally:
            conn.close()

    async def execute_in_process_event_loop(self, conn: Connection):
        try:
            input_params = conn.recv()
            args = input_params.get("args", [])
            kwargs = input_params.get("kwargs", {})
            result = await self.raw_func(*args, **kwargs)
            try:
                json.dumps(result)
            except Exception as e:
                logger.debug(e)
                result = {}
            conn.send(result)
        except Exception as error:
            logger.debug(error)
            raise error
        finally:
            conn.close()

    def create_and_run(self, *args, **kwargs):
        ctx = get_context("fork")
        self.parent_con, self.child_con = ctx.Pipe()
        if self.is_async:
            self.proc = ctx.Process(
                target=lambda: (
                    asyncio.run(
                        self.execute_in_process_event_loop(
                            self.child_con
                        )
                    )
                )
            )
        else:
            self.proc = ctx.Process(
                target=lambda: (
                    self.execute_in_process(
                        self.child_con
                    )
                )
            )
        self.proc.start()
        self.parent_con.send({"args": args, "kwargs": kwargs})
        return self.parent_con

    def __execute(self):
        self.execute_in_process(self.child_con)

    @property
    def ret(self):
        if self.parent_con and not self.parent_con.closed and self.child_con.closed:
            self.result = self.parent_con.recv()
            self.parent_con.close()
        return self.result

    def wait(self, timeout: Optional[float] = None):
        self.proc.join(timeout)

    def receive(self):
        return self.parent_con.recv()

    @staticmethod
    def create_process(function: FunctionType):
        assert isinstance(function, FunctionType), "only accept function for process"
        if inspect.iscoroutinefunction(function):
            async def run_in_different_proc(*args, **kwargs):
                pc = ProcessConnector(function)
                pc.create_and_run(*args, **kwargs)
                return pc
            return run_in_different_proc
        else:
            def run_in_different_proc(*args, **kwargs):
                pc = ProcessConnector(function)
                pc.create_and_run(*args, **kwargs)
                return pc
            return run_in_different_proc
