# -*- coding: utf-8 -*-
"""
合并两张图片（通常是格子遮罩图和文字遮罩图）
"""

import sys
import cv2
import numpy as np
from pathlib import Path

# Windows编码修复
if sys.platform == 'win32':
    import io
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')
    sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace')


def merge_images(background_path, overlay_path, output_path):
    """
    将透明PNG图片叠加到背景图片上
    
    参数:
        background_path: 背景图片路径（textGreenPanelImgPath，底层）
        overlay_path: 叠加图片路径（panelImgPath，顶层，支持透明通道）
        output_path: 输出图片路径
    """
    background_path = Path(background_path)
    overlay_path = Path(overlay_path)
    output_path = Path(output_path)
    
    # 确保输出目录存在
    output_path.parent.mkdir(parents=True, exist_ok=True)
    
    try:
        # 读取背景图片（支持中文路径）
        print(f"[INFO] 读取背景图片: {background_path.name}")
        bg_data = np.fromfile(str(background_path), dtype=np.uint8)
        background = cv2.imdecode(bg_data, cv2.IMREAD_COLOR)
        if background is None:
            raise ValueError(f"无法读取背景图片: {background_path}")
        
        print(f"[INFO] 背景图片尺寸: {background.shape[1]}x{background.shape[0]}")
        
        # 读取叠加图片（支持透明通道）
        print(f"[INFO] 读取叠加图片: {overlay_path.name}")
        overlay_data = np.fromfile(str(overlay_path), dtype=np.uint8)
        overlay = cv2.imdecode(overlay_data, cv2.IMREAD_UNCHANGED)
        if overlay is None:
            raise ValueError(f"无法读取叠加图片: {overlay_path}")
        
        print(f"[INFO] 叠加图片尺寸: {overlay.shape[1]}x{overlay.shape[0]}")
        print(f"[INFO] 叠加图片通道数: {overlay.shape[2] if len(overlay.shape) > 2 else 1}")
        
        # 确保两张图片尺寸一致
        if background.shape[:2] != overlay.shape[:2]:
            print(f"[WARN] 图片尺寸不匹配，调整叠加图片尺寸")
            overlay = cv2.resize(overlay, (background.shape[1], background.shape[0]))
            print(f"[INFO] 调整后叠加图片尺寸: {overlay.shape[1]}x{overlay.shape[0]}")
        
        # 处理叠加图片的透明通道
        if len(overlay.shape) == 3 and overlay.shape[2] == 4:
            # RGBA格式，有透明通道
            print(f"[INFO] 检测到RGBA透明图片，进行Alpha混合")
            
            # 分离RGB和Alpha通道
            overlay_rgb = overlay[:, :, :3]
            alpha = overlay[:, :, 3] / 255.0  # 归一化到0-1
            
            # 扩展alpha通道到3个维度
            alpha_3ch = np.stack([alpha, alpha, alpha], axis=2)
            
            # Alpha混合：result = background * (1 - alpha) + overlay * alpha
            # 透明区域(alpha=0)显示背景，不透明区域(alpha=1)显示叠加图片
            result = background * (1 - alpha_3ch) + overlay_rgb * alpha_3ch
            result = result.astype(np.uint8)
            
            print(f"[INFO] Alpha混合完成，透明区域保留背景")
            
        elif len(overlay.shape) == 3 and overlay.shape[2] == 3:
            # RGB格式，没有透明通道，将白色区域视为透明
            print(f"[INFO] 检测到RGB图片，将白色区域视为透明")
            
            # 创建mask：白色区域(255,255,255)为透明，其他为不透明
            white_threshold = 250  # 接近白色的阈值
            white_mask = np.all(overlay >= white_threshold, axis=2)
            alpha = (~white_mask).astype(float)  # 非白色区域alpha=1，白色区域alpha=0
            
            # 扩展alpha到3个维度
            alpha_3ch = np.stack([alpha, alpha, alpha], axis=2)
            
            # Alpha混合
            result = background * (1 - alpha_3ch) + overlay * alpha_3ch
            result = result.astype(np.uint8)
            
            print(f"[INFO] 白色区域视为透明处理完成")
            
        else:
            # 灰度图或其他格式，简单叠加
            print(f"[INFO] 其他格式图片，使用简单叠加")
            if len(overlay.shape) == 2:
                overlay = cv2.cvtColor(overlay, cv2.COLOR_GRAY2BGR)
            
            # 对于灰度图，将接近白色的区域视为透明
            if len(overlay.shape) == 3:
                gray_overlay = cv2.cvtColor(overlay, cv2.COLOR_BGR2GRAY)
                white_mask = gray_overlay >= 250
                alpha = (~white_mask).astype(float)
                alpha_3ch = np.stack([alpha, alpha, alpha], axis=2)
                result = background * (1 - alpha_3ch) + overlay * alpha_3ch
                result = result.astype(np.uint8)
            else:
                result = cv2.addWeighted(background, 0.7, overlay, 0.3, 0)
        
        print(f"[INFO] 合并完成，结果图片尺寸: {result.shape[1]}x{result.shape[0]}")
        
        # 保存结果图片（支持中文路径）
        print(f"[INFO] 正在保存合并图片到: {output_path.name}")
        success, encoded_img = cv2.imencode('.png', result)
        if success:
            encoded_img.tofile(str(output_path))
            print(f"[SUCCESS] 已保存合并图片: {output_path.name}")
        else:
            raise RuntimeError(f"图片编码失败: {output_path}")
            
    except Exception as e:
        print(f"[ERROR] 详细错误信息: {str(e)}")
        import traceback
        traceback.print_exc()
        raise RuntimeError(f"合并图片失败: {str(e)}")


if __name__ == '__main__':
    if len(sys.argv) < 4:
        print("用法: python merge_images.py <背景图片路径> <叠加图片路径> <输出图片路径>")
        sys.exit(1)
    
    background_path = sys.argv[1]
    overlay_path = sys.argv[2]
    output_path = sys.argv[3]
    
    try:
        merge_images(background_path, overlay_path, output_path)
    except Exception as e:
        print(f"[ERROR] 合并失败: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)