#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Deep Zoom 瓦片集生成工具
用于将高分辨率PNG图像转换为Deep Zoom格式，适用于三体项目的大型历史图表展示

使用方法:
  python deepzoom_generator.py --input <input_image.png> --output <output_dir> --tile_size <tile_size> --overlap <overlap>

参数说明:
  --input: 输入的PNG图像文件路径
  --output: 输出的Deep Zoom目录路径
  --tile_size: 瓦片大小，默认512
  --overlap: 瓦片重叠像素，默认1
  --format: 输出瓦片格式，支持jpg或png，默认jpg
  --quality: JPEG图像质量(1-100)，默认90

示例:
  python deepzoom_generator.py --input "三体结构3.drawio.png" --output deepzoom_output
"""

import os
import argparse
import math
from PIL import Image
from xml.dom import minidom
import logging
from tqdm import tqdm

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class DeepZoomGenerator:
    """Deep Zoom 瓦片集生成器类"""
    
    def __init__(self, input_image_path, output_dir, tile_size=512, overlap=1, 
                 output_format='jpg', quality=90):
        """
        初始化DeepZoomGenerator
        
        参数:
            input_image_path: 输入图像路径
            output_dir: 输出目录路径
            tile_size: 瓦片大小
            overlap: 瓦片重叠像素
            output_format: 输出格式(jpg或png)
            quality: JPEG质量
        """
        self.input_image_path = input_image_path
        self.output_dir = output_dir
        self.tile_size = tile_size
        self.overlap = overlap
        self.output_format = output_format.lower()
        self.quality = quality
        
        # 验证参数
        self._validate_params()
        
        # 创建输出目录
        self._create_output_dirs()
        
        # 加载图像
        self.image = self._load_image()
        self.width, self.height = self.image.size
        
        # 计算金字塔层级
        self.levels = self._calculate_levels()
        
        logger.info(f"输入图像: {input_image_path}")
        logger.info(f"图像尺寸: {self.width}x{self.height}")
        logger.info(f"输出目录: {output_dir}")
        logger.info(f"瓦片大小: {tile_size}, 重叠像素: {overlap}")
        logger.info(f"输出格式: {output_format}")
        logger.info(f"金字塔层级: {self.levels}")
    
    def _validate_params(self):
        """验证输入参数"""
        # 检查输入图像是否存在
        if not os.path.exists(self.input_image_path):
            raise FileNotFoundError(f"输入图像文件不存在: {self.input_image_path}")
        
        # 检查输出格式
        if self.output_format not in ['jpg', 'png']:
            raise ValueError(f"不支持的输出格式: {self.output_format}，仅支持jpg和png")
        
        # 检查质量参数
        if not (1 <= self.quality <= 100):
            raise ValueError(f"JPEG质量必须在1-100之间: {self.quality}")
        
        # 检查瓦片大小
        if self.tile_size <= 0:
            raise ValueError(f"瓦片大小必须大于0: {self.tile_size}")
        
        # 检查重叠像素
        if self.overlap < 0:
            raise ValueError(f"重叠像素不能为负数: {self.overlap}")
    
    def _create_output_dirs(self):
        """创建输出目录结构"""
        # 创建主输出目录
        os.makedirs(self.output_dir, exist_ok=True)
        
        # 提取基本文件名(不含扩展名)
        base_name = os.path.splitext(os.path.basename(self.input_image_path))[0]
        
        # 设置DZI文件名和瓦片目录
        self.dzi_filename = f"{base_name}.dzi"
        self.tiles_dir = f"{base_name}_files"
        self.tiles_dir_path = os.path.join(self.output_dir, self.tiles_dir)
        
        # 创建瓦片目录
        os.makedirs(self.tiles_dir_path, exist_ok=True)
    
    def _load_image(self):
        """加载输入图像"""
        try:
            image = Image.open(self.input_image_path)
            # 确保图像为RGB模式
            if image.mode != 'RGB':
                image = image.convert('RGB')
            return image
        except Exception as e:
            raise IOError(f"无法加载图像: {e}")
    
    def _calculate_levels(self):
        """计算金字塔层级数量"""
        # 计算最大维度
        max_dim = max(self.width, self.height)
        # 计算需要的层级数，确保最小维度至少为1
        levels = math.floor(math.log2(max_dim)) + 1
        return levels
    
    def _create_dzi_file(self):
        """创建DZI XML文件"""
        # 创建XML文档
        doc = minidom.getDOMImplementation().createDocument(None, 'Image', None)
        root = doc.documentElement
        root.setAttribute('xmlns', 'http://schemas.microsoft.com/deepzoom/2008')
        root.setAttribute('Format', self.output_format)
        root.setAttribute('Overlap', str(self.overlap))
        root.setAttribute('TileSize', str(self.tile_size))
        
        # 创建Size元素
        size_element = doc.createElement('Size')
        size_element.setAttribute('Height', str(self.height))
        size_element.setAttribute('Width', str(self.width))
        root.appendChild(size_element)
        
        # 保存XML文件
        dzi_file_path = os.path.join(self.output_dir, self.dzi_filename)
        with open(dzi_file_path, 'w', encoding='utf-8') as f:
            root.writexml(f, indent='  ', addindent='  ', newl='\n')
        
        logger.info(f"创建DZI文件: {dzi_file_path}")
    
    def _generate_tiles(self):
        """生成所有层级的瓦片"""
        current_image = self.image.copy()
        current_width, current_height = current_image.size
        
        # 从最高分辨率到最低分辨率生成瓦片
        for level in range(self.levels):
            # 创建当前层级的目录
            level_dir = os.path.join(self.tiles_dir_path, str(level))
            os.makedirs(level_dir, exist_ok=True)
            
            # 计算当前层级的瓦片数量
            tiles_x = max(1, math.ceil((current_width + 2 * self.overlap) / self.tile_size))
            tiles_y = max(1, math.ceil((current_height + 2 * self.overlap) / self.tile_size))
            
            logger.info(f"生成层级 {level} 的瓦片: {tiles_x}x{tiles_y}")
            
            # 使用tqdm创建进度条
            total_tiles = tiles_x * tiles_y
            with tqdm(total=total_tiles, desc=f"层级 {level}", unit="tile") as pbar:
                # 生成每个瓦片
                for y in range(tiles_y):
                    for x in range(tiles_x):
                        self._generate_single_tile(current_image, level, x, y, level_dir)
                        pbar.update(1)
            
            # 如果不是最后一层，缩小图像到下一层
            if level < self.levels - 1:
                new_width = max(1, current_width // 2)
                new_height = max(1, current_height // 2)
                current_image = current_image.resize((new_width, new_height), Image.Resampling.LANCZOS)
                current_width, current_height = current_image.size
    
    def _generate_single_tile(self, image, level, tile_x, tile_y, level_dir):
        """生成单个瓦片"""
        width, height = image.size
        
        # 计算瓦片在原图中的位置
        tile_size_no_overlap = self.tile_size - 2 * self.overlap
        start_x = max(0, tile_x * tile_size_no_overlap - self.overlap)
        start_y = max(0, tile_y * tile_size_no_overlap - self.overlap)
        
        # 计算瓦片的实际大小
        end_x = min(width, start_x + self.tile_size)
        end_y = min(height, start_y + self.tile_size)
        actual_width = end_x - start_x
        actual_height = end_y - start_y
        
        # 创建一个新的瓦片图像(空白背景)
        tile = Image.new('RGB', (self.tile_size, self.tile_size), color=(255, 255, 255))
        
        # 从原图中裁剪瓦片区域
        tile_region = image.crop((start_x, start_y, end_x, end_y))
        
        # 将裁剪的区域粘贴到瓦片上
        tile.paste(tile_region, (0, 0))
        
        # 保存瓦片
        tile_filename = os.path.join(level_dir, f"{tile_x}_{tile_y}.{self.output_format}")
        
        if self.output_format == 'jpg':
            tile.save(tile_filename, 'JPEG', quality=self.quality, optimize=True)
        else:
            tile.save(tile_filename, 'PNG', optimize=True)
    
    def generate(self):
        """生成完整的Deep Zoom瓦片集"""
        logger.info("开始生成Deep Zoom瓦片集...")
        
        # 创建DZI文件
        self._create_dzi_file()
        
        # 生成瓦片
        self._generate_tiles()
        
        logger.info("Deep Zoom瓦片集生成完成!")
        logger.info(f"DZI文件: {os.path.join(self.output_dir, self.dzi_filename)}")
        logger.info(f"瓦片目录: {self.tiles_dir_path}")


def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='Deep Zoom瓦片集生成工具')
    parser.add_argument('--input', '-i', required=True, help='输入的PNG图像文件路径')
    parser.add_argument('--output', '-o', required=True, help='输出的Deep Zoom目录路径')
    parser.add_argument('--tile_size', '-t', type=int, default=512, help='瓦片大小，默认512')
    parser.add_argument('--overlap', '-l', type=int, default=1, help='瓦片重叠像素，默认1')
    parser.add_argument('--format', '-f', default='jpg', choices=['jpg', 'png'], help='输出瓦片格式，默认jpg')
    parser.add_argument('--quality', '-q', type=int, default=90, help='JPEG图像质量(1-100)，默认90')
    return parser.parse_args()


def main():
    """主函数"""
    args = parse_args()
    
    try:
        # 创建DeepZoomGenerator实例
        generator = DeepZoomGenerator(
            input_image_path=args.input,
            output_dir=args.output,
            tile_size=args.tile_size,
            overlap=args.overlap,
            output_format=args.format,
            quality=args.quality
        )
        
        # 生成Deep Zoom瓦片集
        generator.generate()
        
    except Exception as e:
        logger.error(f"生成Deep Zoom瓦片集时出错: {e}")
        raise


if __name__ == '__main__':
    main()