from datetime import datetime
from pathlib import Path
from typing import Optional

import colormatcher
import cv2
import numpy as np
from maa.controller import AdbController, Controller
from maa.define import MaaAdbScreencapMethodEnum
from maa.toolkit import Toolkit
from roimage import Roi, Roimage

# 初始化设备参数
# device_serial = "127.0.0.1:16384"
device_serial = None
adb_screencap_method = MaaAdbScreencapMethodEnum.Default

# 初始窗口大小 (width, height)
# window_size = (720, 1280) # 竖屏
window_size = (1280, 720)  # 横屏

# 右键拖曳
move_start = (0, 0)
move_start_roi = (0, 0)


# 初始化 Roi
std_roimage: Optional[Roimage] = Roimage(window_size[0], window_size[1])  # 标准化截图
win_roimage: Roimage = Roimage(
    0, 0, 0, 0, std_roimage
)  # 相对 std_roimage ，窗口显示的区域
crop_list: list[Roimage] = []  # 相对 std_roimage ，需要裁剪的区域

# 左键裁剪ROI区域
crop_start: Optional[Roimage] = None  # 相对 win_roimage ，正在裁剪的区域
crop_end = Roimage(0, 0, 0, 0, win_roimage).getRoiInRoot()

# 初始化参数
win_name = "image"
trackbars_name = "trackbars"
file_name = "image"
src_path = Path("./src")
files = [f.name for f in src_path.glob("*.png")] if src_path.exists() else []

# 显示 Roi
trackbars_img = np.ones((100, 400, 3), dtype=np.uint8) * 255

# -----------------------------------------------


# 截图参数
def set_screenshot_target_side(c: Controller):
    # c.set_screenshot_target_long_side(1280)
    c.set_screenshot_target_short_side(720)


# ROI 放大方法
def amplify(rect: list[int]) -> list[int]:
    x, y, w, h = rect
    return [x - 50, y - 50, w + 100, h + 100]


# 颜色匹配方法
def match_color(image) -> tuple[int, int, list[tuple[list[int]]]]:
    # framework ColorMatch common method:
    #   4.COLOR_BGR2RGB 40.COLOR_BGR2HSV 6.COLOR_BGR2GRAY
    # colormatcher method:
    #   1.Simple 2.RGBDistance
    method = cv2.COLOR_BGR2RGB
    reverse = cv2.COLOR_RGB2BGR  # 不需要逆转时为None，用于显示 'MainColors' 窗口
    cluster_colors = colormatcher.kmeansClusterColors(image, method)
    return method, reverse, colormatcher.RGBDistance(cluster_colors)


# 获取 ADB 设备列表
def get_adb_devices_list() -> Optional[Controller]:
    print("MaaToolkit search in progress...\n")
    device_list = Toolkit.find_adb_devices()
    if len(device_list):
        for i, d in enumerate(device_list):
            print(f"{i:>3} | {d.address:>21} | {d.name}")

        if device_input := input("Please select the device (ENTER to pass): "):
            i = int(device_input)
            if 0 <= i < len(device_list):
                device_serial = device_list[i].address
                return AdbController(
                    adb_path=device_list[i].adb_path,
                    address=device_serial,
                    screencap_methods=adb_screencap_method,
                )


# OpenCV 鼠标回调
# -events 鼠标事件（如按下鼠标左键，释放鼠标左键，鼠标移动等）
# -x x坐标
# -y y坐标
# -flags params 其他参数
def mouse(event, x, y, flags, param) -> None:
    global crop_end
    crop_end = Roimage(0, 0, x, y, win_roimage).getRoiInRoot()
    crop(event, x, y, flags, param)
    zoom(event, x, y, flags, param)
    move(event, x, y, flags, param)
    show_roi(crop_end)


def show_roi(roi: Roi):
    trackbars_img.fill(255)
    cv2.putText(
        trackbars_img,
        f"{roi.rectangle}",
        (0, 60),
        cv2.FONT_HERSHEY_SIMPLEX,
        1,
        (0, 0, 0),
        2,
    )
    cv2.imshow(trackbars_name, trackbars_img)


# 计算绘图四边形坐标
# -rectPts 相对于 std_roimage 的两点坐标 ((left,top),(right,bottom))
def count_draw_coordinate(rect_pts):
    z = win_roimage.zoom
    x, y = win_roimage.point
    return (
        (int(rect_pts[0][0] * z - x), int(rect_pts[0][1] * z - y)),
        (int(rect_pts[1][0] * z - x), int(rect_pts[1][1] * z - y)),
    )


# 绘图
# -rois 相对于 std_roimage 的 crop_roimage 列表
def draw(rois: list[Roimage]) -> None:
    img = win_roimage.image.copy()
    for roi in rois:
        pt1, pt2 = count_draw_coordinate(roi.rectanglePoints)
        cv2.rectangle(img, pt1, pt2, (0, 255, 0), 2)
    cv2.imshow(win_name, img)


def crop(event, x, y, flags, param) -> None:
    global crop_start, crop_end
    if event == cv2.EVENT_LBUTTONDOWN:  # 按下左键
        # 记录（x,y）坐标
        crop_start = Roimage(0, 0, x, y, win_roimage)
    elif crop_start is not None and (
        event == cv2.EVENT_LBUTTONUP
        or (event == cv2.EVENT_MOUSEMOVE and (flags & cv2.EVENT_FLAG_LBUTTON))
    ):  # 释放左键 或 按住左键拖曳
        # 修正（x,y）坐标
        w, h = win_roimage.size
        x = max(0, min(w - 1, x))
        y = max(0, min(h - 1, y))
        # 记录（x,y）坐标
        crop_end = crop_start.getCropRoi(x, y).getRoiInRoot()
        # 绘图
        rois: list[Roimage] = crop_list.copy()
        rois.append(crop_end)
        draw(rois)
        # 保存 crop Roi
        if event == cv2.EVENT_LBUTTONUP:
            crop_list.append(crop_end.copy(std_roimage))
            crop_start = None


# 计算缩放倍数
# -flag 鼠标滚轮上移或下移的标识
# -zoom 缩放倍数
# -step 缩放系数，滚轮每步缩放0.1
def count_zoom(flag, zoom: float, step: float = 0.1):
    if flag > 0:  # 滚轮上移
        zoom += step
        if zoom > 3:  # 最多只能放大到 g_image 三倍大
            zoom = 3
    else:  # 滚轮下移
        zoom -= step
        if zoom < 1:  # 最多只能缩小到 g_image 大小
            zoom = 1
    return zoom


# 滚轮放大
def zoom(event, x, y, flags, param) -> None:
    global win_roimage
    if event == cv2.EVENT_MOUSEWHEEL:  # 滚轮
        z = count_zoom(flags, win_roimage.zoom)
        wx, wy, w, h = win_roimage.rectangle
        x = int((wx + x) * z / win_roimage.zoom - x)
        y = int((wy + y) * z / win_roimage.zoom - y)
        win_roimage = Roimage(w, h, x, y, win_roimage.parent, z)
        draw(crop_list)


# 计算移动后的坐标
# -pt0 win_roimage 的原始坐标
# -pt1 鼠标按下右键时的坐标
# -pt2 鼠标当前坐标
def count_move_coordinate(pt0, pt1, pt2):
    return pt0[0] + pt1[0] - pt2[0], pt0[1] + pt1[1] - pt2[1]


def move(event, x, y, flags, param) -> None:
    global move_start, move_start_roi, win_roimage
    if event == cv2.EVENT_RBUTTONDOWN:
        move_start = (x, y)
        move_start_roi = win_roimage.point
    elif event == cv2.EVENT_RBUTTONUP or (
        event == cv2.EVENT_MOUSEMOVE and (flags & cv2.EVENT_FLAG_RBUTTON)
    ):
        x, y = count_move_coordinate(move_start_roi, move_start, (x, y))
        win_roimage = Roimage(
            win_roimage.width, win_roimage.height, x, y, std_roimage, win_roimage.zoom
        )
        draw(crop_list)


# 轨迹条回调
# -pos 轨迹条位置
def trackbar_change(pos) -> None:
    pos = pos / 100  # get a scaling factor from trackbar pos
    w = int(std_roimage.width * pos)  # scale w
    h = int(std_roimage.height * pos)  # scale h
    cv2.resizeWindow(win_name, w, h)  # resize a window


def readfile(file: str):
    file_path = src_path / file
    print("src:", file_path.absolute())
    return cv2.imread(str(file_path))


# 截图
def screenshot() -> Optional[np.ndarray]:
    if not controller:
        return None
    print("Screenshot in progress...")
    controller.post_screencap().wait()
    image = controller.cached_image
    print("Screenshot completed.")
    return image


# 缩放图片到目标短边
def resize_to_target_short_side(
    image: np.ndarray, target_short_side: int = 720
) -> np.ndarray:
    height, width = image.shape[:2]
    short_side = min(width, height)
    if short_side == target_short_side:
        return image  # 不需要缩放
    scale = target_short_side / short_side
    new_width = int(round(width * scale))
    new_height = int(round(height * scale))
    print(f"Resizing image from {width}x{height} to {new_width}x{new_height}")
    return cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_AREA)


# 获取标准化的 Roimage
def get_std_roimage() -> Optional[Roimage]:
    global file_name
    image: Optional[np.ndarray] = None
    if len(files):
        file_name = files.pop(0)
        image = readfile(file_name)
        file_name = file_name.split(".")[0]
        if image is not None:
            image = resize_to_target_short_side(image, 720)
    elif controller:
        image = screenshot()
        file_name = datetime.now().strftime("%H%M%S")  # '%Y%m%d%H%M%S'
    if image is None:
        return None
    height, width, _ = image.shape
    roimage = Roimage(width, height)
    roimage.image = image
    return roimage


# 获取放大后的 Roi 四边形
# -roi: 需要放大的 Roi
def get_amplified_roi_rectangle(roi: Roimage) -> list[int]:
    x, y, w, h = amplify(roi.rectangle)
    return Roimage(w, h, x, y, roi.parent).rectangle


def is_window_visible(win_name: str, destroyed: bool = False) -> bool:
    if destroyed:
        return False
    try:
        return cv2.getWindowProperty(win_name, cv2.WND_PROP_VISIBLE) > 0
    except cv2.error:
        return False


if __name__ == "__main__":
    # Print Help Infos
    print(
        "Usage: python3 main.py [device serial]\n"
        "Put the images under ./src, and run this script, it will be auto converted to target size.\n"
        "Hold down the left mouse button, drag mouse to select a ROI.\n"
        "Hold down the right mouse button, drag mouse to move the image.\n"
        "Use the mouse wheel to zoom the image.\n"
        "press 'S' or 'ENTER' to save ROIs.\n"
        "press 'F' to save a full standardized screenshot.\n"
        "press 'R' to output only the ROI ranges, not save.\n"
        "press 'c' or 'C' (with connected) to output the ROI ranges and colors, not save.\n"
        "press 'Z' or 'DELETE' or 'BACKSPACE' to remove the latest ROI.\n"
        "press '0' ~ '9' to resize the window.\n"
        "press 'Q' or 'ESC' to quit.\n"
        "The cropped images will be saved in ./dst.\n"
    )
    # ---------------------------------------------------------

    # 初始化 dst 文件夹
    dst_path = Path.cwd() / "dst"
    if not dst_path.is_dir():
        dst_path.mkdir()

    # 搜索并连接设备
    controller = get_adb_devices_list()
    if not controller:
        print("ADB Devices Not Found. Reading from ./src folder...")
    if controller:
        set_screenshot_target_side(controller)
        if controller.post_connection().failed:
            print(f"Failed to connect device({device_serial}).")
            exit(1)

    # 初始化 cv2 窗口
    cv2.namedWindow(win_name, cv2.WINDOW_NORMAL)
    cv2.setWindowProperty(win_name, cv2.WND_PROP_TOPMOST, 1)
    cv2.setMouseCallback(win_name, mouse)

    cv2.namedWindow(trackbars_name, cv2.WINDOW_NORMAL)
    cv2.setWindowProperty(trackbars_name, cv2.WND_PROP_TOPMOST, 1)
    cv2.createTrackbar("Scale", trackbars_name, 100, 200, trackbar_change)

    cropping = False
    destroyedMainColors = True
    while True:
        if not cropping:
            std_roimage = get_std_roimage()
            if std_roimage is None:
                break
            win_roimage = Roimage(0, 0, 0, 0, std_roimage)
            crop_list.clear()
        draw(crop_list)

        key = cv2.waitKey(0) & 0xFF
        cropping = True
        # q Q esc
        if key in [ord("q"), ord("Q"), 27]:
            cv2.destroyAllWindows()
            exit()
        # 0
        if key == 48:
            cv2.setTrackbarPos("Scale", trackbars_name, 100)
            continue
        # 1 ~ 5
        if key in range(49, 54):
            cv2.setTrackbarPos("Scale", trackbars_name, 100 + 15 * (key - 48))
            continue
        # 6 ~ 9
        if key in range(54, 58):
            cv2.setTrackbarPos("Scale", trackbars_name, 100 - 15 * (58 - key))
            continue
        # z Z delete backspace
        if key in [ord("z"), ord("Z"), 0, 8]:
            if len(crop_list):
                crop_list.pop()
            continue

        cropping = False
        needSave = True
        needColorMatch = False
        connected = False
        mains = []
        # r R
        if key in [ord("r"), ord("R")]:
            needSave = False
        # c C
        elif key in [ord("c")]:
            needSave = False
            needColorMatch = True
            connected = False
        elif key in [ord("C")]:
            needSave = False
            needColorMatch = True
            connected = True
        # f F
        elif key in [ord("f"), ord("F")]:
            crop_list.append(Roimage(0, 0, 0, 0, std_roimage))
        # s S enter
        elif key not in [ord("s"), ord("S"), ord("\r"), ord("\n")]:
            continue

        for roi in crop_list:
            print("")
            img = roi.image

            if needSave:
                x1, y1, w1, h1 = roi.rectangle
                x2, y2, w2, h2 = get_amplified_roi_rectangle(roi)
                dst_filename: str = (
                    f"{file_name}_{x1}_{y1}_{w1}_{h1}__{x2}_{y2}_{w2}_{h2}.png"
                )
                dst_file_path = dst_path / dst_filename
                print(rf"dst: {dst_file_path.absolute()}")
                cv2.imwrite(str(dst_file_path), roi.image)

            print(
                f"original roi: {roi.rectangle}\n"
                f"amplified roi: {get_amplified_roi_rectangle(roi)}"
            )

            if needColorMatch:
                method, reverse, colors = match_color(img)
                ret = {
                    "recognition": "ColorMatch",
                    "roi": [],
                    "method": method,
                    "lower": [],
                    "upper": [],
                    "count": [],
                    "connected": connected,
                }
                mainColors = []
                for center, lower, upper in colors[0]:
                    count = colormatcher.getCount(img, lower, upper, connected, method)
                    ret["lower"].append(lower)
                    ret["upper"].append(upper)
                    ret["count"].append(count)
                    mainColor = np.zeros((80, 200, 3), dtype=np.uint8)
                    mainColor[:28, :] = lower
                    mainColor[28:58, :] = upper
                    mainColor[58:, :] = center
                    if reverse is not None:
                        mainColor = cv2.cvtColor(mainColor, reverse)
                    cv2.putText(
                        mainColor,
                        f"{lower}",
                        (0, 20),
                        cv2.FONT_HERSHEY_SIMPLEX,
                        0.75,
                        (0, 0, 0),
                        2,
                    )
                    cv2.putText(
                        mainColor,
                        f"{upper}",
                        (0, 50),
                        cv2.FONT_HERSHEY_SIMPLEX,
                        0.75,
                        (255, 255, 255),
                        2,
                    )
                    cv2.putText(
                        mainColor,
                        f"{count}",
                        (66, 76),
                        cv2.FONT_HERSHEY_SIMPLEX,
                        0.75,
                        (127, 127, 127),
                        2,
                    )
                    mainColors.append(mainColor)
                mains.append(cv2.hconcat(mainColors))
                cv2.namedWindow("MainColors", cv2.WINDOW_NORMAL)
                cv2.imshow("MainColors", cv2.vconcat(mains))
                destroyedMainColors = False
                print(
                    f"ColorMatch: {ret}".replace("'", '"')
                    .replace("False", "false")
                    .replace("True", "true")
                )
            elif is_window_visible("MainColors", destroyedMainColors):
                cv2.destroyWindow("MainColors")
                destroyedMainColors = True

            print("")

            print("Press 'Q'/'ESC' to quit or just continute.")

    print("Exiting...")
    cv2.destroyAllWindows()