#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Matrix风格（矩阵拼接风格）海报模板
支持将多张图片按照指定的矩阵布局拼接，带有自定义边框
"""
import base64
from io import BytesIO
from typing import List, Tuple, Optional, Dict, Any, Union
from PIL import Image, ImageDraw

from .base_template import BaseTemplate


class MatrixTemplate(BaseTemplate):
    """
    矩阵拼接模板，用于将多张图片按照指定的矩阵布局拼接成一张海报。
    支持自定义边框宽度和颜色，可以生成PNG或Fabric.js JSON格式。
    """

    def __init__(self, size: Tuple[int, int] = (1350, 1800)):
        super().__init__(size)

    def generate(self,
                 images: List[Union[str, Image.Image]],
                 size: Optional[Tuple[int, int]] = None,
                 layout: Tuple[int, int] = (2, 2),
                 border_width: Optional[Dict[str, List[int]]] = None,
                 border_color: Optional[Dict[str, List[Tuple[int, int, int]]]] = None,
                 generate_fabric_json: bool = False) -> Union[Image.Image, Dict[str, Any]]:
        """
        生成矩阵拼接海报

        Args:
            images (List[Union[str, Image.Image]]): 图片列表，可以是路径或PIL Image对象
            size (Optional[Tuple[int, int]]): 输出海报尺寸 (宽, 高)，默认使用初始化时的尺寸
            layout (Tuple[int, int]): 矩阵布局 (行, 列)，例如 (2, 3) 表示2行3列
            border_width (Optional[Dict[str, List[int]]]): 边框宽度配置
                格式: {
                    'row': [上边框, 行间1, 行间2, ..., 下边框],  # 长度为 rows+1
                    'col': [左边框, 列间1, 列间2, ..., 右边框]   # 长度为 cols+1
                }
            border_color (Optional[Dict[str, List[Tuple[int, int, int]]]]): 边框颜色配置
                格式: {
                    'row': [(R,G,B), (R,G,B), ...],  # 长度为 rows+1
                    'col': [(R,G,B), (R,G,B), ...]   # 长度为 cols+1
                }
            generate_fabric_json (bool): 是否生成Fabric.js JSON格式

        Returns:
            Union[Image.Image, Dict[str, Any]]: 
                - 如果 generate_fabric_json=False: 返回PIL Image对象
                - 如果 generate_fabric_json=True: 返回包含 'image' 和 'fabric_json' 的字典
        """
        # 使用指定的尺寸或默认尺寸
        output_size = size or self.size
        rows, cols = layout

        # 验证输入
        if not images:
            return None
        
        # 验证图片数量必须与布局匹配
        if len(images) != rows * cols:
            return None

        # 加载图片
        loaded_images = []
        for img in images[:rows * cols]:
            if isinstance(img, str):
                try:
                    loaded_img = Image.open(img)
                    if loaded_img.mode != 'RGB':
                        loaded_img = loaded_img.convert('RGB')
                    loaded_images.append(loaded_img)
                except Exception:
                    return None
            elif isinstance(img, Image.Image):
                loaded_images.append(img)
            else:
                return None

        # 设置默认边框
        if border_width is None:
            border_width = {
                'row': [0] * (rows + 1),
                'col': [0] * (cols + 1)
            }
        
        if border_color is None:
            border_color = {
                'row': [(255, 255, 255)] * (rows + 1),
                'col': [(255, 255, 255)] * (cols + 1)
            }

        # 验证边框配置
        if len(border_width['row']) != rows + 1:
            return None
        if len(border_width['col']) != cols + 1:
            return None
        if len(border_color['row']) != rows + 1:
            return None
        if len(border_color['col']) != cols + 1:
            return None

        # 第1步：计算所有单元格的位置和大小
        cells = self._calculate_cell_positions(output_size, layout, border_width)
        
        # 第2步：裁剪所有子图
        cropped_images = self._prepare_cropped_images(loaded_images, cells)
        
        # 第3步：准备边框配方（位置和颜色信息）
        border_recipe = self._prepare_border_recipe(output_size, layout, border_width, border_color)
        
        # 第4步：生成输出（只是按配方组装，不做任何计算）
        if generate_fabric_json:
            # 同时生成PNG和JSON
            poster_image = self._render_png(
                cropped_images, cells, border_recipe, output_size, border_color
            )
            fabric_json = self._render_json(
                cropped_images, cells, border_recipe, output_size, border_color
            )
            return {
                'image': poster_image,
                'fabric_json': fabric_json
            }
        else:
            # 仅生成PNG
            return self._render_png(
                cropped_images, cells, border_recipe, output_size, border_color
            )

    def _calculate_cell_positions(self,
                                   canvas_size: Tuple[int, int],
                                   layout: Tuple[int, int],
                                   border_width: Dict[str, List[int]]) -> List[Dict[str, Any]]:
        """
        计算每个单元格的位置和大小

        Args:
            canvas_size: 画布尺寸 (宽, 高)
            layout: 矩阵布局 (行, 列)
            border_width: 边框宽度配置

        Returns:
            List[Dict[str, Any]]: 每个单元格的位置信息
                [
                    {
                        'row': 行索引,
                        'col': 列索引,
                        'x': x坐标,
                        'y': y坐标,
                        'width': 宽度,
                        'height': 高度
                    },
                    ...
                ]
        """
        canvas_width, canvas_height = canvas_size
        rows, cols = layout

        # 计算可用空间
        total_row_border = sum(border_width['row'])
        total_col_border = sum(border_width['col'])
        available_width = canvas_width - total_col_border
        available_height = canvas_height - total_row_border

        # 计算每个单元格的尺寸
        cell_width = available_width // cols
        cell_height = available_height // rows

        # 计算每个单元格的位置
        cells = []
        current_y = border_width['row'][0]

        for row in range(rows):
            current_x = border_width['col'][0]
            
            for col in range(cols):
                cells.append({
                    'row': row,
                    'col': col,
                    'x': current_x,
                    'y': current_y,
                    'width': cell_width,
                    'height': cell_height
                })
                
                # 移动到下一列
                current_x += cell_width + border_width['col'][col + 1]
            
            # 移动到下一行
            current_y += cell_height + border_width['row'][row + 1]

        return cells

    def _prepare_cropped_images(self, 
                                images: List[Image.Image], 
                                cells: List[Dict[str, Any]]) -> List[Image.Image]:
        """
        准备所有裁剪后的子图

        Args:
            images: 原始图片列表
            cells: 单元格位置信息

        Returns:
            裁剪后的图片列表
        """
        cropped_images = []
        for i, cell in enumerate(cells):
            if i >= len(images):
                break
            cropped_img = self._crop_image_to_fit(
                images[i],
                (cell['width'], cell['height'])
            )
            cropped_images.append(cropped_img)
        return cropped_images

    def _prepare_border_recipe(self,
                               canvas_size: Tuple[int, int],
                               layout: Tuple[int, int],
                               border_width: Dict[str, List[int]],
                               border_color: Dict[str, List[Tuple[int, int, int]]]) -> List[Dict[str, Any]]:
        """
        准备边框配方（所有边框的位置和颜色信息）

        Args:
            canvas_size: 画布尺寸
            layout: 矩阵布局
            border_width: 边框宽度
            border_color: 边框颜色

        Returns:
            边框配方列表，每个元素包含边框的位置和颜色
        """
        canvas_width, canvas_height = canvas_size
        rows, cols = layout
        border_rects = []

        # 计算可用空间
        total_row_border = sum(border_width['row'])
        total_col_border = sum(border_width['col'])
        available_width = canvas_width - total_col_border
        available_height = canvas_height - total_row_border
        cell_width = available_width // cols
        cell_height = available_height // rows

        # 准备水平边框（行边框）
        current_y = 0
        for i in range(rows + 1):
            width = border_width['row'][i]
            if width > 0:
                border_rects.append({
                    'x': 0,
                    'y': current_y,
                    'width': canvas_width,
                    'height': width,
                    'color': border_color['row'][i]
                })
            current_y += width
            if i < rows:
                current_y += cell_height

        # 准备垂直边框（列边框）
        current_x = 0
        for i in range(cols + 1):
            width = border_width['col'][i]
            if width > 0:
                border_rects.append({
                    'x': current_x,
                    'y': 0,
                    'width': width,
                    'height': canvas_height,
                    'color': border_color['col'][i]
                })
            current_x += width
            if i < cols:
                current_x += cell_width

        return border_rects

    def _crop_image_to_fit(self, image: Image.Image, target_size: Tuple[int, int]) -> Image.Image:
        """
        裁剪图片以适应目标尺寸（参考 resize_image 中的 crop 算法）

        Args:
            image: 原始图片
            target_size: 目标尺寸 (宽, 高)

        Returns:
            裁剪后的图片
        """
        target_width, target_height = target_size
        width, height = image.size

        # 计算缩放比例
        ratio_w = target_width / width
        ratio_h = target_height / height

        # 使用较大的比例进行缩放，确保图片完全覆盖目标区域
        ratio = max(ratio_w, ratio_h)
        new_width = int(width * ratio)
        new_height = int(height * ratio)
        resized = image.resize((new_width, new_height), Image.LANCZOS)

        # 计算裁剪区域（居中裁剪）
        left = (new_width - target_width) // 2
        top = (new_height - target_height) // 2
        right = left + target_width
        bottom = top + target_height

        return resized.crop((left, top, right, bottom))

    def _render_png(self,
                    cropped_images: List[Image.Image],
                    cells: List[Dict[str, Any]],
                    border_recipe: List[Dict[str, Any]],
                    canvas_size: Tuple[int, int],
                    border_color: Dict[str, List[Tuple[int, int, int]]]) -> Image.Image:
        """
        渲染PNG格式海报（纯组装，不做任何计算）

        Args:
            cropped_images: 已裁剪好的图片列表
            cells: 单元格位置信息
            border_recipe: 边框配方（位置和颜色）
            canvas_size: 画布尺寸
            border_color: 边框颜色（用于背景）

        Returns:
            拼接后的海报图片
        """
        # 创建画布
        canvas = Image.new('RGB', canvas_size, border_color['row'][0])
        draw = ImageDraw.Draw(canvas)

        # 绘制边框（按配方）
        for border in border_recipe:
            draw.rectangle(
                [(border['x'], border['y']), 
                 (border['x'] + border['width'], border['y'] + border['height'])],
                fill=border['color']
            )

        # 粘贴图片（按配方）
        for img, cell in zip(cropped_images, cells):
            canvas.paste(img, (cell['x'], cell['y']))

        return canvas

    def _render_json(self,
                     cropped_images: List[Image.Image],
                     cells: List[Dict[str, Any]],
                     border_recipe: List[Dict[str, Any]],
                     canvas_size: Tuple[int, int],
                     border_color: Dict[str, List[Tuple[int, int, int]]]) -> Dict[str, Any]:
        """
        渲染Fabric.js JSON格式（纯组装，不做任何计算）

        Args:
            cropped_images: 已裁剪好的图片列表
            cells: 单元格位置信息
            border_recipe: 边框配方（位置和颜色）
            canvas_size: 画布尺寸
            border_color: 边框颜色（用于背景）

        Returns:
            Fabric.js JSON对象
        """
        canvas_width, canvas_height = canvas_size
        fabric_objects = []

        # 添加背景矩形
        fabric_objects.append({
            'type': 'rect',
            'left': 0,
            'top': 0,
            'width': canvas_width,
            'height': canvas_height,
            'fill': self._rgb_to_hex(border_color['row'][0]),
            'selectable': False,
            'evented': False
        })

        # 添加边框矩形（按配方）
        for border in border_recipe:
            fabric_objects.append({
                'type': 'rect',
                'left': border['x'],
                'top': border['y'],
                'width': border['width'],
                'height': border['height'],
                'fill': self._rgb_to_hex(border['color']),
                'selectable': False,
                'evented': False
            })

        # 添加图片对象（按配方）
        for img, cell in zip(cropped_images, cells):
            img_base64 = self._image_to_base64(img)
            fabric_objects.append({
                'type': 'image',
                'left': cell['x'],
                'top': cell['y'],
                'width': cell['width'],
                'height': cell['height'],
                'src': f"data:image/png;base64,{img_base64}",
                'scaleX': 1,
                'scaleY': 1,
                'cropX': 0,
                'cropY': 0,
                'selectable': True,
                'evented': True
            })

        # 构建完整的Fabric.js JSON
        fabric_json = {
            'version': '5.3.0',
            'objects': fabric_objects,
            'background': self._rgb_to_hex(border_color['row'][0]),
            'width': canvas_width,
            'height': canvas_height
        }

        return fabric_json

    def _rgb_to_hex(self, rgb: Tuple[int, int, int]) -> str:
        """将RGB颜色转换为十六进制格式"""
        return f"#{rgb[0]:02x}{rgb[1]:02x}{rgb[2]:02x}"

    def _image_to_base64(self, image: Image.Image) -> str:
        """将PIL Image转换为base64字符串"""
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode()
        return img_str