import time
from functools import cached_property, wraps

import cv2
import numpy as np
import requests

from mower.utils import config
from mower.utils.csleep import MowerExit
from mower.utils.device.method.adb import ADB
from mower.utils.image import bytes2img
from mower.utils.log import logger
from mower.utils.network import get_new_port, is_port_in_use
from mower.utils.path import get_path


def retry_droidcast(func):
    @wraps(func)
    def retry_wrapper(self, *args, **kwargs):
        for _ in range(3):
            try:
                return func(self, *args, **kwargs)
            except MowerExit:
                raise
            except Exception as e:
                logger.debug_exception(e)
                self.start_droidcast()

    return retry_wrapper


class DroidCast:
    display_name = "DroidCast"
    package_name = "com.rayworks.droidcast"
    apk_file = "DroidCast-debug-1.2.1.apk"

    def __init__(self):
        self.session = requests.Session()
        self.port = 0
        self.socket = None

    @cached_property
    def adb(self) -> ADB:
        return config.device.get_service("adb")

    def get_droidcast_classpath(self) -> str | None:
        # TODO: 退出时（并非结束mower线程时）关闭DroidCast进程、取消ADB转发
        out = self.adb.adb_shell(f"pm path {self.package_name}")
        if out is None:
            logger.exception("无法获取CLASSPATH")
            return None
        prefix = "package:"
        postfix = ".apk"
        beg = out.index(prefix, 0)
        end = out.rfind(postfix)
        class_path = out[beg + len(prefix) : (end + len(postfix))].strip()
        class_path = "CLASSPATH=" + class_path
        logger.info(f"成功获取CLASSPATH：{class_path}")
        return class_path

    def start_droidcast(self) -> bool:
        try:
            class_path = self.get_droidcast_classpath()
        except Exception:
            class_path = None
        if not class_path:
            logger.info(f"安装{self.display_name}")
            try:
                apk_path = get_path(f"@install/mower/vendor/{self.apk_file}")
                self.adb.adb.install(str(apk_path))
                logger.info(f"{self.display_name}安装完成，获取CLASSPATH")
            except Exception as e:
                logger.error(f"{self.display_name}安装失败: {e}")
                return False
            class_path = self.get_droidcast_classpath()
            if not class_path:
                logger.error("无法获取CLASSPATH")
                return False
        port = self.port
        if port != 0 and is_port_in_use(port):
            try:
                forward_list = self.adb.adb.forward_list()
                for forward in forward_list:
                    if (
                        forward.local == f"tcp:{port}"
                        and forward.remote == f"tcp:{port}"
                    ):
                        break
            except Exception as e:
                logger.exception(e)
        if port == 0:
            port = get_new_port()
            self.port = port
            logger.info(f"更新{self.display_name}端口为{port}")
        else:
            logger.info(f"保持{self.display_name}端口为{port}")
        self.adb.adb.forward(f"tcp:{port}", f"tcp:{port}")
        logger.info(f"ADB端口转发成功，启动{self.display_name}")
        cmd = f"{class_path} app_process / {self.package_name}.Main --port={port}"
        self.socket = self.adb.adb_shell(cmd, True)
        time.sleep(1)  # TODO: 更好地等待DroidCast启动
        return True

    def decode(self, data):
        img = bytes2img(data)
        if config.conf.droidcast.rotate:
            img = cv2.rotate(img, cv2.ROTATE_180)
        return img

    @retry_droidcast
    def capture_display(self):
        if self.port == 0:
            raise Exception(f"{self.display_name}未启动")
        url = f"http://127.0.0.1:{self.port}/screenshot"
        logger.debug(f"GET {url}")
        r = self.session.get(url)
        return self.decode(r.content)

    def kill_server(self):
        pass

    def reset_when_exit(self):
        pass


class DroidCast_raw(DroidCast):
    display_name = "DroidCast_raw"
    package_name = "ink.mol.droidcast_raw"
    apk_file = "DroidCast_raw-release-1.1.apk"

    def decode(self, data):
        # https://github.com/LmeSzinc/StarRailCopilot/blob/bf62eccaaa6f46ea9796a83175df0ab7882dd137/module/device/method/droidcast.py#L226
        arr = np.frombuffer(data, dtype=np.uint16)
        if config.conf.droidcast_raw.orientation == "portrait":
            shape = (1080, 1920)
        else:
            shape = (1920, 1080)
        arr = arr.reshape(shape)

        r = cv2.bitwise_and(arr, 0b1111100000000000)
        r = cv2.convertScaleAbs(r, alpha=0.00390625)
        m = cv2.convertScaleAbs(r, alpha=0.03125)
        cv2.add(r, m, dst=r)

        g = cv2.bitwise_and(arr, 0b0000011111100000)
        g = cv2.convertScaleAbs(g, alpha=0.125)
        m = cv2.convertScaleAbs(g, alpha=0.015625, dst=m)
        cv2.add(g, m, dst=g)

        b = cv2.bitwise_and(arr, 0b0000000000011111)
        b = cv2.convertScaleAbs(b, alpha=8)
        m = cv2.convertScaleAbs(b, alpha=0.03125, dst=m)
        cv2.add(b, m, dst=b)

        img = cv2.merge([r, g, b])

        if config.conf.droidcast_raw.rotate == 90:
            img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
        elif config.conf.droidcast_raw.rotate == 180:
            img = cv2.rotate(img, cv2.ROTATE_180)
        elif config.conf.droidcast_raw.rotate == 270:
            img = cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)

        return img
