import os
import math
from pathlib import Path
from importlib.resources import files
from typing import Union, List, Tuple, Optional

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

# 尝试导入 torch
try:
    import torch

    TORCH_AVAILABLE = True
except ImportError:
    TORCH_AVAILABLE = False

# 支持的图像扩展名
SUPPORTED_IMAGE_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif', '.webp'}


def merge_images(
        images: Union[str, list, "torch.Tensor", np.ndarray],
        tags: Optional[List[str]] = None,
        col_size: Union[int, Tuple[int, int], List[int]] = None,
        num_cols: Optional[int] = None,
        save_file: Optional[str] = None,
        background_color: Tuple[int, int, int] = (255, 255, 255),
        col_border_width: int = 0,
        col_border_color: Tuple[int, int, int] = (255, 0, 0),
        font_path: Optional[str] = None,
        font_size: Optional[int] = 30,
) -> Optional[Union[np.ndarray, "torch.Tensor", Image.Image]]:
    """
    将多张图片合并在一张内，网格布局，图片以宽高等比缩放居中填充。

    :param images:
        - str: 文件夹路径 → 自动读图，tags 默认为文件名（不含后缀）
        - 其他类型（tensor / numpy / PIL list）→ tags=None 时用序号 "0", "1", ...
    :param tags: 标签列表。若为 None，则按上述规则自动生成。
    :param col_size: 单元格尺寸 (w, h) 或正方形边长 int。
    :param num_cols: 每行列数，默认为 ceil(sqrt(n))。
    :param save_file: 若指定，则保存图像且**不返回结果**（返回 None）。
    :param background_color: 背景色 (R, G, B)
    :param col_border_width: 网格线宽度
    :param col_border_color: 网格线颜色
    :param font_path: 字体路径
    :param font_size: 字体大小
    :return: 合并后的图像（PIL / numpy / tensor），除非 save_file 被指定 → 返回 None
    """

    input_type = None
    pil_images = []

    # ==================== Step 0: 处理 images 为文件夹 ====================
    if isinstance(images, str):
        image_dir = Path(images)
        if not image_dir.is_dir():
            raise ValueError(f"Provided path '{images}' is not a valid directory.")

        image_files = []
        for f in sorted(image_dir.iterdir()):
            if f.is_file() and f.suffix.lower() in SUPPORTED_IMAGE_EXTENSIONS:
                image_files.append(f)
        if not image_files:
            raise ValueError(f"No supported images in: {images}")

        for f in image_files:
            try:
                img = Image.open(f).convert('RGB')
                pil_images.append(img)
            except Exception as e:
                print(f"Warning: Skip invalid image {f}: {e}")
        if not pil_images:
            raise RuntimeError("No valid images loaded.")

        input_type = 'pil'
        # 自动标签：文件名（不含后缀）
        if tags is None:
            tags = [f.stem for f in image_files]

    else:
        # ==================== Step 1: 转换为 PIL 列表 ====================
        if isinstance(images, list):
            if not images:
                raise ValueError("Empty image list")
            first = images[0]
            if TORCH_AVAILABLE and isinstance(first, torch.Tensor):
                input_type = 'tensor'
                if first.dim() == 3:
                    for img in images:
                        assert img.dim() == 3 and img.shape[0] == 3
                        img_np = (img.permute(1, 2, 0).clamp(0, 1) * 255).byte().cpu().numpy()
                        pil_images.append(Image.fromarray(img_np, 'RGB'))
                else:
                    raise ValueError("Unsupported tensor format")
            elif isinstance(first, np.ndarray):
                input_type = 'numpy'
                for img in images:
                    assert img.ndim == 3 and img.shape[2] in (1, 3)
                    if img.shape[2] == 3:
                        rgb = img[:, :, ::-1]
                    else:
                        rgb = np.stack([img.squeeze()] * 3, axis=-1)
                    pil_images.append(Image.fromarray(rgb.astype(np.uint8), 'RGB'))
            elif isinstance(first, Image.Image):
                input_type = 'pil'
                pil_images = [img.convert('RGB') for img in images]
            else:
                raise TypeError(f"Unsupported type in list: {type(first)}")
        elif TORCH_AVAILABLE and isinstance(images, torch.Tensor):
            input_type = 'tensor'
            if images.dim() == 4 and images.shape[1] == 3:
                for img in images:
                    img_np = (img.permute(1, 2, 0).clamp(0, 1) * 255).byte().cpu().numpy()
                    pil_images.append(Image.fromarray(img_np, 'RGB'))
            else:
                raise ValueError("Tensor must be BCHW with 3 channels")
        elif isinstance(images, np.ndarray):
            input_type = 'numpy'
            if images.ndim == 4 and images.shape[3] in (1, 3):
                for img in images:
                    if img.shape[2] == 3:
                        rgb = img[:, :, ::-1]
                    else:
                        rgb = np.stack([img.squeeze()] * 3, axis=-1)
                    pil_images.append(Image.fromarray(rgb.astype(np.uint8), 'RGB'))
            else:
                raise ValueError("Numpy array must be BHWC with 1 or 3 channels")
        else:
            raise TypeError("Unsupported images type")

    n_imgs = len(pil_images)
    if n_imgs == 0:
        raise ValueError("No images to merge.")

    # ==================== Step 2: 自动生成 tags（仅当非文件夹且 tags=None） ====================
    if tags is None:
        # 注意：文件夹情况已在上一步处理，此处只处理非文件夹
        tags = [str(i) for i in range(n_imgs)]

    assert len(tags) == n_imgs, f"tags length ({len(tags)}) != image count ({n_imgs})"

    # ==================== Step 3: 布局参数 ====================
    if col_size is None:
        w0, h0 = pil_images[0].size
        col_w, col_h = w0, h0
    elif isinstance(col_size, int):
        col_w = col_h = col_size
    elif isinstance(col_size, (tuple, list)) and len(col_size) == 2:
        col_w, col_h = col_size
    else:
        raise ValueError("col_size must be int or (w, h)")

    if num_cols is None:
        num_cols = math.ceil(math.sqrt(n_imgs))
    num_rows = math.ceil(n_imgs / num_cols)

    total_w = num_cols * col_w + (num_cols - 1) * col_border_width
    total_h = num_rows * col_h + (num_rows - 1) * col_border_width

    canvas = Image.new('RGB', (total_w, total_h), background_color)
    draw = ImageDraw.Draw(canvas)

    # ==================== Step 4: 字体 ====================
    if font_path is None:
        try:
            font_path = files("aabd.image").joinpath("AlibabaPuHuiTi-3-65-Medium.ttf").as_posix()
            if not os.path.exists(font_path):
                font_path = None
        except Exception:
            font_path = None

    try:
        font = ImageFont.truetype(font_path, font_size) if font_path else ImageFont.load_default()
    except IOError:
        font = ImageFont.load_default()

    # 新增：默认半透明黑色背景色
    DEFAULT_TEXT_BG_COLOR = (0, 0, 0, 180)  # RGBA中的A表示透明度，范围是0-255

    # ==================== Step 5: 渲染图像和标签 ====================
    for idx, pil_img in enumerate(pil_images):
        row = idx // num_cols
        col = idx % num_cols
        x0 = col * (col_w + col_border_width)
        y0 = row * (col_h + col_border_width)

        # 缩放居中
        img_w, img_h = pil_img.size
        scale = min(col_w / img_w, col_h / img_h)
        new_w, new_h = int(img_w * scale), int(img_h * scale)
        resized = pil_img.resize((new_w, new_h), Image.LANCZOS)
        paste_x = x0 + (col_w - new_w) // 2
        paste_y = y0 + (col_h - new_h) // 2
        canvas.paste(resized, (paste_x, paste_y))

        # 标签（即使为空字符串也不绘制）
        if tags[idx] is not None and str(tags[idx]).strip():
            text = str(tags[idx])

            # 获取文本包围盒（以左上角为锚点）
            left, top, right, bottom = draw.textbbox((0, 0), text, font=font, anchor='lt')
            text_width = right - left
            text_height = bottom - top

            # 标签左上角位置（相对于单元格）
            label_x = x0 + 5
            label_y = y0 + 5

            # 创建半透明背景
            text_bg = Image.new('RGBA', (text_width, text_height), (0, 0, 0, 180))
            canvas.paste(text_bg, (label_x, label_y), text_bg)

            # 使用 anchor='lt' 确保文字从左上角开始绘制，与背景对齐
            draw.text((label_x, label_y), text, fill=(255, 255, 255), font=font, anchor='lt')
        # 边框
        if col_border_width > 0:
            draw.rectangle([x0, y0, x0 + col_w, y0 + col_h],
                           outline=col_border_color, width=col_border_width)

    # ==================== Step 6: 保存或返回 ====================
    if save_file:
        save_path = Path(save_file)
        save_path.parent.mkdir(parents=True, exist_ok=True)
        canvas.save(save_path)
        return None  # 不返回图像

    # 否则按输入类型返回
    if input_type == 'tensor':
        out_np = np.array(canvas)
        out_tensor = torch.from_numpy(out_np).permute(2, 0, 1).unsqueeze(0).float() / 255.0
        return out_tensor
    elif input_type == 'numpy':
        out_bgr = np.array(canvas)[:, :, ::-1]
        return np.expand_dims(out_bgr, axis=0)
    else:
        return canvas


if __name__ == '__main__':
    merge_images('imgs', save_file='merge.png', num_cols=2)
