import ctypes
import os
import subprocess
import typing as t
from dataclasses import dataclass
from functools import wraps

import cv2
import numpy as np
from matplotlib import pyplot as plt

from mower.utils import config
from mower.utils.csleep import MowerExit
from mower.utils.device.emulator import restart_emulator
from mower.utils.device.exception import EmulatorError, GameError
from mower.utils.log import logger


class LDOpenGLIncompatible(Exception):
    pass


class ImageCaptureError(Exception):
    pass


def retry_ldopengl(func):
    @wraps(func)
    def retry_wrapper(self, *args, **kwargs):
        for _ in range(3):
            try:
                return func(self, *args, **kwargs)
            except MowerExit:
                raise
            except GameError as e:
                logger.debug_exception(e)
                config.device.check_current_focus()
            except EmulatorError as e:
                logger.debug_exception(e)
                restart_emulator()
            except Exception as e:
                logger.debug_exception(e)

    return retry_wrapper


def bytes_to_str(b: bytes) -> str:
    for encoding in ["utf-8", "gbk"]:
        try:
            return b.decode(encoding)
        except UnicodeDecodeError:
            pass
    return str(b)


@dataclass
class DataLDPlayerInfo:
    # Emulator instance index, starting from 0
    index: int
    # Instance name
    name: str
    # Handle of top window
    topWnd: int
    # Handle of bind window
    bndWnd: int
    # If instance is running, 1 for True, 0 for False
    sysboot: int
    # PID of the instance process, or -1 if instance is not running
    playerpid: int
    # PID of the vbox process, or -1 if instance is not running
    vboxpid: int
    # Resolution
    width: int
    height: int
    dpi: int

    def __post_init__(self):
        self.index = int(self.index)
        self.name = bytes_to_str(self.name)
        self.topWnd = int(self.topWnd)
        self.bndWnd = int(self.bndWnd)
        self.sysboot = int(self.sysboot)
        self.playerpid = int(self.playerpid)
        self.vboxpid = int(self.vboxpid)
        self.width = int(self.width)
        self.height = int(self.height)
        self.dpi = int(self.dpi)


class LDConsole:
    def __init__(self, ld_folder: str):
        """
        Args:
            ld_folder: Installation path of MuMu12, e.g. E:/ProgramFiles/LDPlayer9
                which should have a `ldconsole.exe` in it.
        """
        self.ld_console = os.path.abspath(
            os.path.join(config.conf.emulator.emulator_folder, "./ldconsole.exe")
        )

    def subprocess_run(self, cmd, timeout=10):
        """
        Args:
            cmd (list):
            timeout (int):

        Returns:
            bytes:
        """
        cmd = [self.ld_console] + cmd
        logger.info(f"Execute: {cmd}")

        try:
            process = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=False)
        except FileNotFoundError as e:
            logger.warning(f"warning when calling {cmd}, {str(e)}")
            raise LDOpenGLIncompatible("ld_folder does not have ldconsole.exe")
        try:
            stdout, stderr = process.communicate(timeout=timeout)
        except subprocess.TimeoutExpired:
            process.kill()
            stdout, stderr = process.communicate()
            logger.warning(
                f"TimeoutExpired when calling {cmd}, stdout={stdout}, stderr={stderr}"
            )
        return stdout

    def list2(self) -> t.List[DataLDPlayerInfo]:
        """
        > ldconsole.exe list2
        0,雷电模拟器,28053900,42935798,1,59776,36816,1280,720,240
        1,雷电模拟器-1,0,0,0,-1,-1,1280,720,240
        """
        out = []
        data = self.subprocess_run(["list2"])
        for row in data.strip().split(b"\n"):
            info = row.strip().split(b",")
            info = DataLDPlayerInfo(*info)
            out.append(info)
        return out


class IScreenShotClass:
    def __init__(self, ptr):
        self.ptr = ptr

        # Define in class since ctypes.WINFUNCTYPE is windows only
        cap_type = ctypes.WINFUNCTYPE(ctypes.c_void_p)
        release_type = ctypes.WINFUNCTYPE(None)
        self.class_cap = cap_type(1, "IScreenShotClass_Cap")
        # Keep reference count
        # so __del__ won't have an empty IScreenShotClass_Cap
        self.class_release = release_type(2, "IScreenShotClass_Release")

    def cap(self):
        return self.class_cap(self.ptr)

    def __del__(self):
        self.class_release(self.ptr)


class LDOpenGLImpl:
    def __init__(self, ld_folder: str, instance_id: int):
        """
        Args:
            ld_folder: Installation path of MuMu12, e.g. E:/ProgramFiles/LDPlayer9
            instance_id: Emulator instance ID, starting from 0
        """
        ldopengl_dll = os.path.abspath(os.path.join(ld_folder, "./ldopengl64.dll"))
        logger.info(
            f"LDOpenGL init, "
            f"ld_folder={ld_folder}, "
            f"ldopengl_dll={ldopengl_dll}, "
            f"instance_id={instance_id}"
        )
        # Load dll
        try:
            self.lib = ctypes.WinDLL(ldopengl_dll)
        except OSError as e:
            logger.error(e)
            if not os.path.exists(ldopengl_dll):
                raise LDOpenGLIncompatible(
                    f"ldopengl_dll={ldopengl_dll} does not exist, "
                    f"ldopengl requires LDPlayer >= 9.0.78, please check your version"
                )
            else:
                raise LDOpenGLIncompatible(
                    f"ldopengl_dll={ldopengl_dll} exist, " f"but cannot be loaded"
                )
        # Get info after loading DLL, so DLL existence can act as a version check
        self.console = LDConsole(ld_folder)
        self.info = self.get_player_info_by_index(instance_id)

        self.lib.CreateScreenShotInstance.restype = ctypes.c_void_p

        # Get screenshot instance
        instance_ptr = ctypes.c_void_p(
            self.lib.CreateScreenShotInstance(instance_id, self.info.playerpid)
        )
        self.screenshot_instance = IScreenShotClass(instance_ptr)

    def get_player_info_by_index(self, instance_id: int):
        """
        Args:
            instance_id:

        Returns:
            DataLDPlayerInfo:

        Raises:
            ImageCaptureError:
        """
        for info in self.console.list2():
            if info.index == instance_id:
                logger.info(f"Match LDPlayer instance: {info}")
                if not info.sysboot:
                    raise ImageCaptureError(
                        "Trying to connect LDPlayer instance but emulator is not running"
                    )
                return info
        raise ImageCaptureError(f"No LDPlayer instance with index {instance_id}")

    @retry_ldopengl
    def screenshot(self):
        """
        Returns:
            np.ndarray: Image array in BGR color space
                Note that image is upside down
        """
        width, height = self.info.width, self.info.height

        img_ptr = self.screenshot_instance.cap()
        # ValueError: NULL pointer access
        if img_ptr is None:
            raise ImageCaptureError("Empty image pointer")

        img = ctypes.cast(
            img_ptr, ctypes.POINTER(ctypes.c_ubyte * (height * width * 3))
        ).contents

        image = np.ctypeslib.as_array(img).reshape((height, width, 3))
        return image


class LDOpenGL:
    def __init__(self):
        try:
            self.emulator_folder = config.conf.emulator.emulator_folder
            self.instanse_index = int(config.conf.emulator.index)

            self.ldopengl = LDOpenGLImpl(
                ld_folder=self.emulator_folder,
                instance_id=self.instanse_index,
            )

        except Exception as e:
            logger.error(e)
            logger.error("Unable to initialize LDOpenGL")
            raise LDOpenGLIncompatible

    def screenshot_ldopengl(self):
        image = self.ldopengl.screenshot()

        image = cv2.flip(image, 0)
        cv2.cvtColor(image, cv2.COLOR_BGR2RGB, dst=image)
        return image


if __name__ == "__main__":
    plt.imshow(LDOpenGL().screenshot_ldopengl())
    plt.show()
