from io import BytesIO

import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont

from konabot.common.path import ASSETS_PATH, FONTS_PATH


def text_to_transparent_image(text, font_size=40, padding=0, text_color=(0, 0, 0)):
    """
    将文本转换为带透明背景的图像，图像大小刚好包含文本
    """
    # 创建临时图像来计算文本尺寸
    temp_image = Image.new('RGB', (1, 1), (255, 255, 255))
    temp_draw = ImageDraw.Draw(temp_image)
    
    font = ImageFont.truetype(FONTS_PATH / "montserrat.otf", font_size)
    
    # 获取文本边界框
    bbox = temp_draw.textbbox((0, 0), text, font=font)
    text_width = bbox[2] - bbox[0]
    text_height = bbox[3] - bbox[1]
    
    # 计算图像大小（文本大小 + 内边距）
    image_width = int(text_width + 2 * padding)
    image_height = int(text_height + 2 * padding)
    
    # 创建RGBA模式的空白图像（带透明通道）
    image = Image.new('RGBA', (image_width, image_height), (0, 0, 0, 0))
    draw = ImageDraw.Draw(image)
    
    # 绘制文本（考虑内边距）
    x = padding - bbox[0]  # 调整起始位置
    y = padding - bbox[1]
    
    # 设置文本颜色（带透明度）
    if len(text_color) == 3:
        text_color = text_color + (255,)  # 添加完全不透明的alpha值
    
    draw.text((x, y), text, fill=text_color, font=font)
    
    # 转换为OpenCV格式（BGRA）
    image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGBA2BGRA)
    return image_cv

def perspective_transform(image, target, corners):
    """
    对图像进行透视变换（保持透明通道）
    target: 画布
    corners: 四个角点的坐标，顺序为 [左上, 右上, 右下, 左下]
    """
    height, width = image.shape[:2]
    
    # 源点（原始图像的四个角）
    src_points = np.array([
        [0, 0],           # 左上
        [width-1, 0],     # 右上
        [width-1, height-1], # 右下
        [0, height-1]     # 左下
    ], dtype=np.float32)
    
    # 目标点（变换后的四个角）
    dst_points = np.array(corners, dtype=np.float32)
    
    # 计算透视变换矩阵
    matrix = cv2.getPerspectiveTransform(src_points, dst_points)
    
    # 获取画布大小
    target_height, target_width = target.shape[:2]

    # 应用透视变换（保持所有通道，包括alpha）
    transformed = cv2.warpPerspective(image, matrix, (target_width, target_height), flags=cv2.INTER_LINEAR)

    return transformed, matrix

def blend_with_transparency(background, foreground, position):
    """
    将带透明通道的前景图像合成到背景图像上
    position: 前景图像在背景图像上的位置 (x, y)
    """
    bg = background.copy()
    
    # 如果背景没有alpha通道，添加一个
    if bg.shape[2] == 3:
        bg = cv2.cvtColor(bg, cv2.COLOR_BGR2BGRA)
        bg[:, :, 3] = 255  # 完全不透明
    
    x, y = position
    fg_height, fg_width = foreground.shape[:2]
    bg_height, bg_width = bg.shape[:2]
    
    # 确保位置在图像范围内
    x = max(0, min(x, bg_width - fg_width))
    y = max(0, min(y, bg_height - fg_height))
    
    # 提取前景的alpha通道并归一化
    alpha_foreground = foreground[:, :, 3] / 255.0
    
    # 对于每个颜色通道进行合成
    for c in range(3):
        bg_region = bg[y:y+fg_height, x:x+fg_width, c]
        fg_region = foreground[:, :, c]
        
        # alpha混合公式
        bg[y:y+fg_height, x:x+fg_width, c] = (
            alpha_foreground * fg_region + 
            (1 - alpha_foreground) * bg_region
        )
    
    # 更新背景的alpha通道（如果需要）
    bg_alpha_region = bg[y:y+fg_height, x:x+fg_width, 3]
    bg[y:y+fg_height, x:x+fg_width, 3] = np.maximum(bg_alpha_region, foreground[:, :, 3])
    
    return bg

def precise_blend_with_perspective(background, foreground, corners):
    """
    精确合成：根据四个角点将前景图像透视合成到背景上
    """
    # 创建与背景相同大小的空白图像
    bg_height, bg_width = background.shape[:2]
    
    # 如果背景没有alpha通道，转换为BGRA
    if background.shape[2] == 3:
        background_bgra = cv2.cvtColor(background, cv2.COLOR_BGR2BGRA)
    else:
        background_bgra = background.copy()
    
    # 创建与背景相同大小的前景图层
    foreground_layer = np.zeros((bg_height, bg_width, 4), dtype=np.uint8)
    
    # 计算前景图像在背景中的边界框
    min_x = int(min(corners[:, 0]))
    max_x = int(max(corners[:, 0]))
    min_y = int(min(corners[:, 1]))
    max_y = int(max(corners[:, 1]))
    
    # 将变换后的前景图像放置到对应位置
    fg_height, fg_width = foreground.shape[:2]
    if min_y + fg_height <= bg_height and min_x + fg_width <= bg_width:
        foreground_layer[min_y:min_y+fg_height, min_x:min_x+fg_width] = foreground
    
    # 创建掩码（只在前景有内容的地方合成）
    mask = (foreground_layer[:, :, 3] > 0)
    
    # 合成图像
    result = background_bgra.copy()
    for c in range(3):
        result[:, :, c][mask] = foreground_layer[:, :, c][mask]
    result[:, :, 3][mask] = foreground_layer[:, :, 3][mask]
    
    return result

def draw_line_bresenham(image, x0, y0, x1, y1, color):
    """使用Bresenham算法画线，避免间隙"""
    dx = abs(x1 - x0)
    dy = abs(y1 - y0)
    sx = 1 if x0 < x1 else -1
    sy = 1 if y0 < y1 else -1
    err = dx - dy
    
    while True:
        if 0 <= x0 < image.shape[1] and 0 <= y0 < image.shape[0]:
            image[y0, x0] = color
        
        if x0 == x1 and y0 == y1:
            break
            
        e2 = 2 * err
        if e2 > -dy:
            err -= dy
            x0 += sx
        if e2 < dx:
            err += dx
            y0 += sy

def slice_and_stretch(image, slice_lines, direction):
    '''
    image: 图像
    slice_lines: 切割线（两个点的列表），一般是倾斜45度的直线
    direction: 移动方向向量（二元数组）
    '''
    # 获取图片的尺寸
    height, width = image.shape[:2]
    # 创建一个由移动方向扩充后，更大的图片
    new_width = int(width + abs(direction[0]))
    new_height = int(height + abs(direction[1]))
    new_image = np.zeros((new_height, new_width, 4), dtype=image.dtype)
    # 先把图片放在新图的和方向相反的一侧
    offset_x = int(abs(min(0, direction[0])))
    offset_y = int(abs(min(0, direction[1])))
    new_image[offset_y:offset_y+height, offset_x:offset_x+width] = image
    # 切割线也跟着偏移
    slice_lines = [(x + offset_x, y + offset_y) for (x, y) in slice_lines]
    # 复制切割线经过的像素，沿着方向移动，实现类似拖尾的效果
    apply_trail_effect_vectorized(new_image, slice_lines, direction)
    apply_stroke_vectorized(new_image, slice_lines, direction)


    return new_image, offset_x, offset_y

def apply_trail_effect_vectorized(new_image, slice_lines, direction):
    """向量化实现拖尾效果"""
    height, width = new_image.shape[:2]
    
    # 创建坐标网格
    y_coords, x_coords = np.mgrid[0:height, 0:width]
    
    # 向量化计算点到直线的距离
    line_vec = np.array([slice_lines[1][0] - slice_lines[0][0], 
                        slice_lines[1][1] - slice_lines[0][1]])
    point_vecs = np.stack([x_coords - slice_lines[0][0], 
                          y_coords - slice_lines[0][1]], axis=-1)
    
    # 计算叉积（有向距离）
    cross_products = (line_vec[0] * point_vecs[:, :, 1] - 
                     line_vec[1] * point_vecs[:, :, 0])
    
    # 选择直线右侧的像素 (d1 > 0)
    mask = cross_products > 0
    
    # 计算目标位置
    target_x = (x_coords + direction[0]).astype(int)
    target_y = (y_coords + direction[1]).astype(int)
    
    # 创建有效位置掩码
    valid_mask = mask & (target_x >= 0) & (target_x < width) & \
                (target_y >= 0) & (target_y < height)
    
    # 批量复制像素
    new_image[target_y[valid_mask], target_x[valid_mask]] = \
        new_image[y_coords[valid_mask], x_coords[valid_mask]]

def apply_stroke_vectorized(new_image, slice_lines, direction):
    """使用向量化操作优化笔画效果"""
    height, width = new_image.shape[:2]
    
    # 1. 找到所有非透明像素
    non_transparent = np.where(new_image[:, :, 3] > 0)
    if len(non_transparent[0]) == 0:
        return
    
    y_coords, x_coords = non_transparent
    
    # 2. 向量化计算点到直线的距离
    line_vec = np.array([slice_lines[1][0] - slice_lines[0][0], 
                        slice_lines[1][1] - slice_lines[0][1]])
    point_vecs = np.column_stack([x_coords - slice_lines[0][0], 
                                 y_coords - slice_lines[0][1]])
    
    # 计算叉积（距离）
    cross_products = (line_vec[0] * point_vecs[:, 1] - 
                     line_vec[1] * point_vecs[:, 0])
    
    # 3. 选择靠近直线的像素
    mask = np.abs(cross_products) < 1.0
    selected_y = y_coords[mask]
    selected_x = x_coords[mask]
    selected_pixels = new_image[selected_y, selected_x]
    
    if len(selected_x) == 0:
        return
    
    # 4. 预计算采样点
    length = np.sqrt(direction[0]**2 + direction[1]**2)
    if length == 0:
        return
    
    # 创建采样偏移
    dx_dy = np.array([(dx, dy) for dx in [-0.5, 0, 0.5] 
                     for dy in [-0.5, 0, 0.5]])
    
    # 5. 批量计算目标位置
    steps = max(1, int(length * 2))
    alpha = 0.7
    
    for k in range(1, steps + 1):
        # 对所有选中的像素批量计算新位置
        scale = k / steps
        
        # 为每个像素和每个采样点计算目标位置
        for dx, dy in dx_dy:
            target_x = np.round(selected_x + dx + direction[0] * scale).astype(int)
            target_y = np.round(selected_y + dy + direction[1] * scale).astype(int)
            
            # 创建有效位置掩码
            valid_mask = (target_x >= 0) & (target_x < width) & \
                        (target_y >= 0) & (target_y < height)
            
            if np.any(valid_mask):
                valid_target_x = target_x[valid_mask]
                valid_target_y = target_y[valid_mask]
                valid_source_idx = np.where(valid_mask)[0]
                
                # 批量混合像素
                source_pixels = selected_pixels[valid_source_idx]
                target_pixels = new_image[valid_target_y, valid_target_x]
                
                new_image[valid_target_y, valid_target_x] = (
                    alpha * source_pixels + (1 - alpha) * target_pixels
                )

async def generate_dice_image(number: str) -> BytesIO:
    # 将文本转换为带透明背景的图像
    text = number

    # 如果文本太长，直接返回金箍棒
    if(len(text) > 50):
        output = BytesIO()
        push_image = Image.open(ASSETS_PATH / "img" / "dice" / "stick.png")
        push_image.save(output,format='PNG')
        output.seek(0)
        return output

    text_image = text_to_transparent_image(
        text, 
        font_size=60, 
        text_color=(0, 0, 0)  # 黑色文字
    )

    # 获取长宽比
    height, width = text_image.shape[:2]
    aspect_ratio = width / height

    # 根据长宽比设置拉伸系数
    stretch_k = 1
    if aspect_ratio > 1:
        stretch_k = aspect_ratio

    # 骰子的方向
    up_direction = (51 - 16, 5 - 30)  # 右上角点 - 左上角点

    move_distance = (up_direction[0] * (stretch_k - 1), up_direction[1] * (stretch_k - 1))
    
    # 加载背景图像，保留透明通道
    background = cv2.imread(ASSETS_PATH / "img" / "dice" / "template.png", cv2.IMREAD_UNCHANGED)

    height, width = background.shape[:2]

    background, offset_x, offset_y = slice_and_stretch(background, 
                                                       [(10,10),(0,0)], 
                                                        move_distance)

    # 定义3D变换的四个角点（透视效果）
    # 顺序: [左上, 右上, 右下, 左下]
    corners = np.array([
        [16, 30],    # 左上
        [51 + move_distance[0], 5 + move_distance[1]],     # 右上（上移，创建透视）
        [88 + move_distance[0], 33 + move_distance[1]],    # 右下
        [49, 62]     # 左下（下移）
    ], dtype=np.float32)
    corners[:, 0] += offset_x
    corners[:, 1] += offset_y


    # 对文本图像进行3D变换（保持透明通道）
    transformed_text, transform_matrix = perspective_transform(text_image, background, corners)
    
    min_x = int(min(corners[:, 0]))
    min_y = int(min(corners[:, 1]))
    final_image_simple = blend_with_transparency(background, transformed_text, (min_x, min_y))

    pil_final = Image.fromarray(final_image_simple)
    # 导入一系列图像
    images: list[Image.Image] = [Image.open(ASSETS_PATH / "img" / "dice" / f"{i}.png") for i in range(1, 12)]
    images.append(pil_final)
    frame_durations = [100] * (len(images) - 1) + [100000]
    # 将导入的图像尺寸扩展为和 pil_final 相同的大小，随帧数进行扩展，然后不放大的情况下放在最中间
    if(aspect_ratio > 1):
        target_size = pil_final.size
        for i in range(len(images) - 1):
            k = i / (len(images) - 1)
            now_distance = (move_distance[0] * k, move_distance[1] * k)
            img = np.array(images[i])
            img, _, _ = slice_and_stretch(img,
                            [(10,10),(0,0)], 
                            now_distance)
            # 只扩展边界，图像本身不放大
            img_width, img_height = img.shape[1], img.shape[0]
            new_img = Image.new("RGBA", target_size, (0, 0, 0, 0))
            this_offset_x = (target_size[0] - img_width) // 2
            this_offset_y = (target_size[1] - img_height) // 2
            # new_img.paste(img, (this_offset_x, this_offset_y))
            new_img.paste(Image.fromarray(img), (this_offset_x, this_offset_y))
            images[i] = new_img
        

    # 保存为BytesIO对象
    output = BytesIO()
    images[0].save(output,
               save_all=True,
               append_images=images[1:],
               duration=frame_durations,
               format='GIF',
               loop=1)
    output.seek(0)
    # pil_final.save(output, format='PNG')
    return output