"""图像生成工具模块

提供棋盘格和 ChArUco 纹理生成功能。
"""

from pathlib import Path

import numpy as np
from PIL import Image


def create_checkerboard_image(filepath, rows=6, cols=9, square_size_px=500):
    """创建棋盘格纹理图像"""
    width = cols * square_size_px
    height = rows * square_size_px
    img = np.ones((height, width, 3), dtype=np.uint8) * 255

    for r in range(rows):
        for c in range(cols):
            if (r + c) % 2 == 1:
                img[r * square_size_px:(r + 1) * square_size_px,
                    c * square_size_px:(c + 1) * square_size_px] = 0

    border = square_size_px
    img_with_border = np.pad(
        img,
        pad_width=((border, border), (border, border), (0, 0)),
        mode="constant",
        constant_values=255,
    )

    abs_filepath = Path(filepath).resolve()
    abs_filepath.parent.mkdir(parents=True, exist_ok=True)
    Image.fromarray(img_with_border).save(abs_filepath)
    usd_filepath = str(abs_filepath).replace("\\", "/")
    print(f"[*] 棋盘格纹理已生成: {usd_filepath}")
    return usd_filepath


def create_charuco_image(filepath, squares_x=30, squares_y=10, square_size_px=90):
    """创建 ChArUco 纹理图像"""
    width = squares_x * square_size_px
    height = squares_y * square_size_px
    abs_filepath = Path(filepath).resolve()
    abs_filepath.parent.mkdir(parents=True, exist_ok=True)

    try:
        import cv2
        dictionary = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_250)
        try:
            board = cv2.aruco.CharucoBoard((squares_x, squares_y), 1.0, 0.70, dictionary)
        except TypeError:
            board = cv2.aruco.CharucoBoard_create(squares_x, squares_y, 1.0, 0.70, dictionary)
        if hasattr(board, "generateImage"):
            img = board.generateImage((width, height), marginSize=0)
        else:
            img = board.draw((width, height), marginSize=0)
        if len(img.shape) == 2:
            img = np.repeat(img[:, :, None], 3, axis=2)
        Image.fromarray(img).save(abs_filepath)
        usd_filepath = str(abs_filepath).replace("\\", "/")
        print(f"[*] ChArUco 纹理已生成: {usd_filepath}")
        return usd_filepath
    except Exception as exc:
        print(f"[WARN] OpenCV ChArUco 生成失败，使用内置 ChArUco 风格纹理: {exc}")

    # 备用实现
    img = np.ones((height, width, 3), dtype=np.uint8) * 255
    marker_cells = 6
    marker_margin = max(3, square_size_px // 7)
    marker_size = square_size_px - 2 * marker_margin
    marker_cell_px = max(1, marker_size // marker_cells)
    marker_px = marker_cell_px * marker_cells

    def marker_bits(marker_id):
        marker = np.ones((marker_cells, marker_cells), dtype=np.uint8) * 255
        marker[0, :] = 0
        marker[-1, :] = 0
        marker[:, 0] = 0
        marker[:, -1] = 0
        state = (marker_id + 1) * 1103515245 + 12345
        for r in range(1, marker_cells - 1):
            for c in range(1, marker_cells - 1):
                state = (state * 1664525 + 1013904223 + r * 97 + c * 193) & 0xFFFFFFFF
                marker[r, c] = 0 if (state & 1) else 255
        return marker

    marker_id = 0
    for r in range(squares_y):
        for c in range(squares_x):
            y0 = r * square_size_px
            x0 = c * square_size_px
            if (r + c) % 2 == 1:
                img[y0:y0 + square_size_px, x0:x0 + square_size_px] = 0
                continue
            marker = marker_bits(marker_id)
            marker_img = np.kron(marker, np.ones((marker_cell_px, marker_cell_px), dtype=np.uint8))
            marker_img = marker_img[:marker_px, :marker_px]
            marker_rgb = np.repeat(marker_img[:, :, None], 3, axis=2)
            marker_y = y0 + (square_size_px - marker_px) // 2
            marker_x = x0 + (square_size_px - marker_px) // 2
            img[marker_y:marker_y + marker_px, marker_x:marker_x + marker_px] = marker_rgb
            marker_id += 1

    Image.fromarray(img).save(abs_filepath)
    usd_filepath = str(abs_filepath).replace("\\", "/")
    print(f"[*] ChArUco 风格纹理已生成: {usd_filepath}")
    return usd_filepath
