from functools import lru_cache
from pathlib import Path
from typing import Literal, Optional, Tuple

import cv2
import numpy as np
from PIL import Image

from mower.utils import typealias as tp
from mower.utils.log import logger
from mower.utils.path import get_path


def bytes2img(
    img_data: bytes,
    gray: bool = False,
    bg: tuple[int, int, int, int] | str = "BLACK",
) -> tp.Image | tp.GrayImage:
    """解码图片

    Args:
        img_data (bytes): 图片的数据
        gray (bool, optional): 返回灰图. Defaults to False.
        bg (tuple[int, int, int, int] | str, optional): 读取透明PNG时设置的背景颜色. Defaults to "BLACK".

    Returns:
        tp.Image | tp.GrayImage: 图片
    """
    img = cv2.imdecode(np.frombuffer(img_data, np.uint8), cv2.IMREAD_UNCHANGED)
    if len(img.shape) == 2:  # 灰图
        if gray:
            return img
        else:
            return cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
    elif img.shape[2] == 4:  # 彩色PNG
        pim = Image.fromarray(img)
        pbg = Image.new("RGBA", pim.size, bg)
        pbg.paste(pim, (0, 0), pim)
        if gray:
            return np.array(pbg.convert("L"))
        else:
            return cv2.cvtColor(np.array(pbg.convert("RGB")), cv2.COLOR_BGR2RGB)
    else:  # 彩色JPG
        if gray:
            return cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        else:
            return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)


def img2bytes(
    img: tp.Image | tp.GrayImage,
    format: Literal["png", "jpg"] = "jpg",
    quality: int = 75,
) -> bytes:
    """编码图片

    Args:
        img (tp.Image | tp.GrayImage): 图片
        format (Literal["png", "jpg"], optional): 图片格式. Defaults to "jpg".
        quality (int, optional): JPG质量. Defaults to 75.

    Returns:
        bytes: 图片数据
    """
    if len(img.shape) == 3:
        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    if format == "jpg":
        data = cv2.imencode(".jpg", img, (cv2.IMWRITE_JPEG_QUALITY, quality))[1]
    else:
        data = cv2.imencode(".png", img)[1]
    return data


def read_file(filename: str | Path) -> bytes:
    """读取图片数据，使用LRU Cache缓存

    Args:
        filename (str | Path): 文件名

    Returns:
        bytes: 图片数据
    """
    filename = Path(filename)
    try:
        filename = filename.relative_to(get_path("@install"))
    except ValueError:
        pass
    logger.debug(filename)
    return np.fromfile(filename, dtype=np.uint8)


def loadimg(
    filename: str, gray: bool = False, bg: tuple[int, int, int, int] | str = "BLACK"
) -> tp.Image | tp.GrayImage:
    """加载图片

    Args:
        filename (str): 文件名
        gray (bool, optional): 返回灰图. Defaults to False.
        bg (tuple[int, int, int, int] | str, optional): 透明PNG的背景颜色. Defaults to "BLACK".

    Returns:
        tp.Image | tp.GrayImage: 图片
    """
    return bytes2img(read_file(filename), gray, bg)


def loadres(res: tp.Res, gray: bool = False) -> tp.Image | tp.GrayImage:
    """加载resources目录下的图片

    Args:
        res (tp.Res): 资源名
        gray (bool, optional): 返回灰图. Defaults to False.

    Returns:
        tp.Image | tp.GrayImage: 图片
    """
    res = f"@install/mower/resources/{res}"
    if not res.endswith(".jpg"):
        res += ".png"
    return loadimg(get_path(res), gray)


loadres = lru_cache(maxsize=256)(loadres)  # 使用装饰器没有类型提示


def load_static(name: str, gray: bool = False) -> tp.Image | tp.GrayImage:
    """加载static目录下的图片

    Args:
        name (str): 资源名
        gray (bool, optional): 返回灰图. Defaults to False.

    Returns:
        tp.Image | tp.GrayImage: 图片
    """
    name = f"@install/mower/static/{name}"
    if not name.endswith(".jpg"):
        name += ".png"
    return loadimg(get_path(name), gray)


def load_static_dir(
    subdir: str, gray: bool = True, str2int: bool = True
) -> dict[str, tp.Image | tp.GrayImage]:
    """加载static的一个子文件夹下的所有图片

    Args:
        subdir (str): 子文件夹
        gray (bool, optional): 返回灰图. Defaults to False.
        str2int (bool, optional): 字典键名转化为数字. Defaults to False.

    Returns:
        dict[str, tp.Image | tp.GrayImage]: 图片字典
    """
    result = {}
    for p in get_path(f"@install/mower/static/{subdir}").iterdir():
        key = p.stem
        if str2int:
            try:
                key = int(key)
            except ValueError:
                pass
        result[key] = loadimg(p, gray)
    return result


def thres2(img: tp.GrayImage, thresh: int) -> tp.GrayImage:
    """二值化

    Args:
        img (tp.GrayImage): 彩图
        thresh (int): 阈值

    Returns:
        tp.GrayImage: 灰图
    """
    return cv2.threshold(img, thresh, 255, cv2.THRESH_BINARY)[1]


def scope2slice(scope: tp.Scope) -> tp.Slice:
    """((x0, y0), (x1, y1)) -> ((y0, y1), (x0, x1))"""
    if scope is None:
        return slice(None), slice(None)
    return slice(scope[0][1], scope[1][1]), slice(scope[0][0], scope[1][0])


def cropimg(img: tp.Image, scope: tp.Scope) -> tp.Image:
    """crop image"""
    return img[scope2slice(scope)]


def cmatch(
    img1: tp.Image, img2: tp.Image, thresh: int = 10, draw: bool = False
) -> bool:
    "比较平均色"
    h, w, _ = img1.shape
    ca = cv2.mean(img1)[:3]
    cb = cv2.mean(img2)[:3]
    diff = np.array(ca).astype(int) - np.array(cb).astype(int)
    diff = np.max(np.maximum(diff, 0)) - np.min(np.minimum(diff, 0))

    if draw:
        board = np.zeros([h + 5, w * 2, 3], dtype=np.uint8)
        board[:h, :w, :] = img1
        board[h:, :w, :] = ca
        board[:h, w:, :] = img2
        board[h:, w:, :] = cb

        from matplotlib import pyplot as plt

        logger.debug(f"{ca=} {cb=} {diff=}")
        plt.imshow(board)
        plt.show()

    return diff <= thresh


def template(
    img1: tp.Image, img2: tp.Image, scope: Optional[tp.Scope] = None, draw: bool = False
) -> Tuple[float, tp.Scope]:
    "在img1中模板匹配img2"
    try:
        h, w, _ = img2.shape
        if scope:
            x, y = scope[0]
            img1 = cropimg(img1, scope)
        else:
            x, y = (0, 0)
        result = cv2.matchTemplate(img1, img2, cv2.TM_CCOEFF_NORMED)
        min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
        score = max_val
        p1 = (max_loc[0] + x, max_loc[1] + y)
        p2 = (p1[0] + w, p1[1] + h)
        scope = p1, p2
        if draw:
            from matplotlib import pyplot as plt

            disp = img1.copy()
            cv2.rectangle(
                disp, max_loc, (max_loc[0] + w, max_loc[1] + h), (0, 255, 0), 3
            )
            plt.imshow(disp)
            plt.show()
        return score, scope
    except Exception as e:
        logger.exception(e)
        return 0, ((0, 0), (0, 0))


def diff_ratio(
    img1: tp.GrayImage,
    img2: tp.GrayImage,
    thresh: int = 0,
    ratio: float = 0.05,
    draw: bool = False,
) -> bool:
    """计算两张灰图之间不同的像素所占比例

    Args:
        img1 (tp.GrayImage): 一张灰图
        img2 (tp.GrayImage): 另一张灰图
        thresh (int, optional): 认为像素有差别的阈值. Defaults to 0.
        ratio (float, optional): 判定有差别的比例阈值. Defaults to 0.05.
        draw (bool, optional): 画出有差异的像素. Defaults to False.

    Returns:
        bool: 两张灰图是否有差异
    """
    h, w = img1.shape
    diff = cv2.absdiff(img1, img2)
    thres = thres2(diff, thresh)
    result = cv2.countNonZero(thres) / w / h

    if draw:
        from matplotlib import pyplot as plt

        logger.debug(f"{result=}")
        plt.imshow(diff, cmap="gray", vmin=0, vmax=255)
        plt.show()
        plt.imshow(thres, cmap="gray", vmin=0, vmax=255)
        plt.show()

    return result > ratio


def crop2content(img: tp.GrayImage) -> tp.GrayImage:
    try:
        contours, _ = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        rect = [cv2.boundingRect(c) for c in contours]
        left = min(x for x, y, w, h in rect)
        right = max([x + w for x, y, w, h in rect])
        top = min([y for x, y, w, h in rect])
        bottom = max([y + h for x, y, w, h in rect])
        return cropimg(img, ((left, top), (right, bottom)))
    except Exception as e:
        logger.exception(e)
        return img


def horizontal_stitch(images: list) -> np.ndarray:
    """
    水平方向拼接多张图像，利用重叠区域对齐。
    """
    # 初始化最终拼接结果为第一张图像
    stitched_image = images[0]
    for i in range(1, len(images)):
        # 当前图像和已拼接的图像
        prev_img = stitched_image
        curr_img = images[i]
        # 转为灰度图，用于计算重叠区域
        gray_prev = cv2.cvtColor(prev_img, cv2.COLOR_RGB2GRAY)
        gray_curr = cv2.cvtColor(curr_img, cv2.COLOR_RGB2GRAY)
        # 匹配 prev_img 的右侧区域和 curr_img 的左侧区域
        match_width = min(800, gray_prev.shape[1])  # 假设重叠区域宽度最大为 800 像素
        prev_region = gray_prev[:, -match_width:]  # 取 prev_img 的右侧区域
        result = cv2.matchTemplate(gray_curr, prev_region, method=cv2.TM_CCOEFF_NORMED)
        _, _, _, max_loc = cv2.minMaxLoc(result)
        overlap_x = max_loc[0]  # 获取匹配的 x 坐标
        # 校验重叠区域的合理性
        if overlap_x >= 0 and overlap_x < curr_img.shape[1]:
            # 拼接：保留 curr_img 中非重叠部分
            non_overlap_part = curr_img[:, overlap_x + match_width :]
            stitched_image = np.hstack((prev_img, non_overlap_part))
        else:
            logger.info("未检测到有效的重叠区域，直接拼接当前图像")
            stitched_image = np.hstack((prev_img, curr_img))
    return stitched_image
