"""
地形导入导出模块
提供完整的地形数据导入导出功能，支持多种格式和批量操作
"""

import os
import json
import math
import numpy as np
from datetime import datetime
from panda3d.core import PNMImage, Filename, PNMFileTypeRegistry, PNMFileType

class TerrainIO:
    """
    地形导入导出类
    提供完整的地形数据导入导出功能
    """
    
    def __init__(self, world):
        self.world = world
        self.supported_formats = ['.png', '.jpg', '.jpeg', '.bmp', '.tga', '.raw', '.ter']
        
    def import_heightmap(self, heightmap_path, scale=(1, 1, 1), resolution=None):
        """
        从高度图导入地形
        支持多种图像格式的高度图
        """
        if not os.path.exists(heightmap_path):
            print(f"高度图文件不存在: {heightmap_path}")
            return None
            
        try:
            # 检查文件格式
            file_ext = os.path.splitext(heightmap_path)[1].lower()
            if file_ext not in self.supported_formats:
                print(f"不支持的高度图格式: {file_ext}")
                return None
                
            # 使用现有的地形管理器创建地形
            if hasattr(self.world, 'terrain_manager'):
                terrain = self.world.terrain_manager.createTerrainFromHeightMap(heightmap_path, scale)
                return terrain
            else:
                print("地形管理器不可用")
                return None
                
        except Exception as e:
            print(f"导入高度图时出错: {e}")
            return None
            
    def export_heightmap(self, terrain_info, output_path, format='png', bit_depth=8):
        """
        导出地形为高度图
        支持多种格式和位深度
        """
        if not terrain_info:
            print("地形信息无效")
            return False
            
        try:
            # 从地形信息中获取高度图数据
            if 'heightfield' in terrain_info:
                heightfield = terrain_info['heightfield']
                if heightfield:
                    # 确保输出目录存在
                    output_dir = os.path.dirname(output_path)
                    if not os.path.exists(output_dir):
                        os.makedirs(output_dir)
                        
                    # 根据位深度调整高度图
                    if bit_depth == 16:
                        # 转换为16位
                        temp_image = PNMImage(heightfield.getXSize(), heightfield.getYSize(), 1)  # 灰度图
                        temp_image.setMaxval(65535)  # 16位最大值
                        
                        for x in range(heightfield.getXSize()):
                            for y in range(heightfield.getYSize()):
                                # 将0-1范围的值转换为0-65535范围
                                value = int(heightfield.getRed(x, y) * 65535)
                                temp_image.setGray(x, y, value)
                                
                        heightfield = temp_image
                        
                    # 获取文件扩展名
                    file_ext = os.path.splitext(output_path)[1].lower()
                    if not file_ext:
                        output_path += f".{format}"
                        
                    # 写入高度图文件
                    success = heightfield.write(Filename.fromOsSpecific(output_path))
                    if success:
                        print(f"高度图导出成功: {output_path} (位深度: {bit_depth})")
                        return True
                    else:
                        print(f"高度图导出失败: {output_path}")
                        return False
                else:
                    print("地形中没有高度图数据")
                    return False
            else:
                print("地形信息中没有高度图数据")
                return False
        except Exception as e:
            print(f"导出高度图时出错: {e}")
            return False
            
    def import_raw_heightmap(self, raw_path, width, height, bit_depth=8, scale=(1, 1, 1)):
        """
        导入RAW格式的高度图
        """
        try:
            if not os.path.exists(raw_path):
                print(f"RAW文件不存在: {raw_path}")
                return None
                
            # 读取RAW数据
            with open(raw_path, 'rb') as f:
                if bit_depth == 8:
                    data = np.fromfile(f, dtype=np.uint8)
                elif bit_depth == 16:
                    data = np.fromfile(f, dtype=np.uint16)
                else:
                    print(f"不支持的位深度: {bit_depth}")
                    return None
                    
            # 检查数据大小
            if len(data) != width * height:
                print(f"RAW数据大小不匹配: 期望 {width * height}, 实际 {len(data)}")
                return None
                
            # 创建PNM图像
            heightfield = PNMImage(width, height)
            
            # 填充高度数据
            if bit_depth == 8:
                max_val = 255.0
            else:  # 16位
                max_val = 65535.0
                
            for y in range(height):
                for x in range(width):
                    # RAW文件通常是行优先存储
                    index = y * width + x
                    value = data[index] / max_val
                    heightfield.setRed(x, y, value)
                    heightfield.setGreen(x, y, value)
                    heightfield.setBlue(x, y, value)
                    
            # 保存到临时文件
            temp_path = f"/tmp/temp_heightmap_{int(datetime.now().timestamp())}.png"
            heightfield.write(Filename.fromOsSpecific(temp_path))
            
            # 使用现有方法创建地形
            return self.import_heightmap(temp_path, scale)
            
        except Exception as e:
            print(f"导入RAW高度图时出错: {e}")
            return None
            
    def export_raw_heightmap(self, terrain_info, output_path, bit_depth=16):
        """
        导出地形为RAW格式的高度图
        """
        if not terrain_info or 'heightfield' not in terrain_info:
            print("地形信息无效或没有高度图数据")
            return False
            
        try:
            heightfield = terrain_info['heightfield']
            width = heightfield.getXSize()
            height = heightfield.getYSize()
            
            # 确保输出目录存在
            output_dir = os.path.dirname(output_path)
            if not os.path.exists(output_dir):
                os.makedirs(output_dir)
                
            # 创建数据数组
            if bit_depth == 8:
                data = np.zeros(width * height, dtype=np.uint8)
                max_val = 255.0
            elif bit_depth == 16:
                data = np.zeros(width * height, dtype=np.uint16)
                max_val = 65535.0
            else:
                print(f"不支持的位深度: {bit_depth}")
                return False
                
            # 填充数据
            for y in range(height):
                for x in range(width):
                    value = heightfield.getRed(x, y) * max_val
                    index = y * width + x
                    data[index] = int(value)
                    
            # 写入RAW文件
            with open(output_path, 'wb') as f:
                data.tofile(f)
                
            print(f"RAW高度图导出成功: {output_path} ({width}x{height}, {bit_depth}位)")
            return True
            
        except Exception as e:
            print(f"导出RAW高度图时出错: {e}")
            return False
            
    def import_terrain_data(self, terrain_data_path):
        """
        导入地形数据文件（JSON格式）
        """
        try:
            if not os.path.exists(terrain_data_path):
                print(f"地形数据文件不存在: {terrain_data_path}")
                return None
                
            with open(terrain_data_path, 'r', encoding='utf-8') as f:
                terrain_data = json.load(f)
                
            # 解析地形数据并创建地形
            terrain_type = terrain_data.get('type', 'heightmap')
            
            if terrain_type == 'heightmap' and 'heightmap_path' in terrain_data:
                heightmap_path = terrain_data['heightmap_path']
                # 检查高度图路径是否为相对路径
                if not os.path.isabs(heightmap_path):
                    # 转换为绝对路径
                    data_dir = os.path.dirname(terrain_data_path)
                    heightmap_path = os.path.join(data_dir, heightmap_path)
                    
                scale = tuple(terrain_data.get('scale', (1, 1, 1)))
                resolution = terrain_data.get('resolution')
                
                return self.import_heightmap(heightmap_path, scale, resolution)
                
            elif terrain_type == 'procedural':
                # 创建程序化地形
                return self._create_procedural_terrain(terrain_data)
                
            else:
                print(f"不支持的地形类型: {terrain_type}")
                return None
                
        except Exception as e:
            print(f"导入地形数据时出错: {e}")
            return None
            
    def export_terrain_data(self, terrain_info, output_path, include_heightmap=True):
        """
        导出地形数据文件（JSON格式）
        """
        if not terrain_info:
            print("地形信息无效")
            return False
            
        try:
            # 收集地形信息
            terrain_data = {
                'name': terrain_info.get('name', f'terrain_{int(datetime.now().timestamp())}'),
                'scale': terrain_info.get('scale', (1, 1, 1)),
                'type': 'heightmap' if terrain_info.get('heightmap') else 'procedural',
                'created_at': datetime.now().isoformat(),
                'resolution': {
                    'width': terrain_info['heightfield'].getXSize() if terrain_info.get('heightfield') else 0,
                    'height': terrain_info['heightfield'].getYSize() if terrain_info.get('heightfield') else 0
                },
                'metadata': terrain_info.get('metadata', {})
            }
            
            # 添加高度图信息
            if terrain_info.get('heightmap'):
                heightmap_path = terrain_info['heightmap']
                terrain_data['heightmap_path'] = os.path.basename(heightmap_path)
                
            # 确保输出目录存在
            output_dir = os.path.dirname(output_path)
            if not os.path.exists(output_dir):
                os.makedirs(output_dir)
                
            # 写入地形数据文件
            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(terrain_data, f, indent=2, ensure_ascii=False)
                
            print(f"地形数据导出成功: {output_path}")
            
            # 如果需要，也导出高度图
            if include_heightmap and terrain_info.get('heightmap'):
                heightmap_output = output_path.replace('.json', '_heightmap.png')
                return self.export_heightmap(terrain_info, heightmap_output)
                
            return True
            
        except Exception as e:
            print(f"导出地形数据时出错: {e}")
            return False
            
    def batch_import(self, import_configs):
        """
        批量导入地形
        import_configs: 导入配置列表，每个配置包含路径和参数
        """
        imported_terrains = []
        
        for config in import_configs:
            path = config.get('path')
            if not path:
                continue
                
            # 根据文件扩展名确定导入方法
            file_ext = os.path.splitext(path)[1].lower()
            
            if file_ext in ['.json']:
                # 导入地形数据文件
                terrain = self.import_terrain_data(path)
            elif file_ext in ['.raw']:
                # 导入RAW格式
                width = config.get('width', 129)
                height = config.get('height', 129)
                bit_depth = config.get('bit_depth', 8)
                scale = config.get('scale', (1, 1, 1))
                terrain = self.import_raw_heightmap(path, width, height, bit_depth, scale)
            else:
                # 导入图像格式高度图
                scale = config.get('scale', (1, 1, 1))
                resolution = config.get('resolution')
                terrain = self.import_heightmap(path, scale, resolution)
                
            if terrain:
                imported_terrains.append(terrain)
                
        print(f"批量导入完成，成功导入 {len(imported_terrains)} 个地形")
        return imported_terrains
        
    def batch_export(self, terrain_infos, output_dir, export_format='json', include_heightmap=True):
        """
        批量导出地形
        """
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
            
        exported_count = 0
        
        for i, terrain_info in enumerate(terrain_infos):
            if export_format.lower() == 'json':
                output_path = os.path.join(output_dir, f"terrain_{i}_data.json")
                if self.export_terrain_data(terrain_info, output_path, include_heightmap):
                    exported_count += 1
            else:
                # 导出为高度图
                output_path = os.path.join(output_dir, f"terrain_{i}_heightmap.png")
                if self.export_heightmap(terrain_info, output_path):
                    exported_count += 1
                    
        print(f"批量导出完成，成功导出 {exported_count} 个地形")
        return exported_count
        
    def _create_procedural_terrain(self, terrain_data):
        """
        创建程序化地形
        """
        try:
            if not hasattr(self.world, 'terrain_manager'):
                print("地形管理器不可用")
                return None
                
            # 获取参数
            scale = tuple(terrain_data.get('scale', (1, 1, 1)))
            resolution = terrain_data.get('resolution', 129)
            algorithm = terrain_data.get('algorithm', 'perlin')
            
            # 创建平面地形作为基础
            terrain = self.world.terrain_manager.createFlatTerrain(
                size=(scale[0], scale[1]), 
                resolution=resolution
            )
            
            if terrain:
                # 应用程序化算法修改高度
                self._apply_procedural_algorithm(terrain, algorithm, terrain_data)
                
            return terrain
            
        except Exception as e:
            print(f"创建程序化地形时出错: {e}")
            return None
            
    def _apply_procedural_algorithm(self, terrain, algorithm, params):
        """
        应用程序化算法生成地形
        """
        print(f"应用程序化算法: {algorithm}")
        # 这里可以实现各种程序化地形生成算法
        # 如Perlin噪声、Voronoi图、分形等
        
    def import_terragen_terrain(self, ter_path, scale=(1, 1, 1)):
        """
        导入Terragen格式的地形文件
        """
        try:
            if not os.path.exists(ter_path):
                print(f"Terragen文件不存在: {ter_path}")
                return None
                
            # 解析Terragen文件格式
            terrain_data = self._parse_terragen_file(ter_path)
            if not terrain_data:
                return None
                
            # 创建高度图
            width = terrain_data['width']
            height = terrain_data['height']
            heightfield = PNMImage(width, height)
            
            # 填充高度数据
            heights = terrain_data['heights']
            for y in range(height):
                for x in range(width):
                    index = y * width + x
                    if index < len(heights):
                        # 将高度值归一化到0-1范围
                        normalized_height = (heights[index] - terrain_data['min_height']) / (terrain_data['max_height'] - terrain_data['min_height'])
                        heightfield.setRed(x, y, normalized_height)
                        heightfield.setGreen(x, y, normalized_height)
                        heightfield.setBlue(x, y, normalized_height)
                        
            # 保存到临时文件
            temp_path = f"/tmp/terragen_heightmap_{int(datetime.now().timestamp())}.png"
            heightfield.write(Filename.fromOsSpecific(temp_path))
            
            # 使用现有方法创建地形
            return self.import_heightmap(temp_path, scale)
            
        except Exception as e:
            print(f"导入Terragen地形时出错: {e}")
            return None
            
    def _parse_terragen_file(self, ter_path):
        """
        解析Terragen文件
        """
        try:
            with open(ter_path, 'rb') as f:
                # 读取文件头
                header = f.read(16)
                if header[:8] != b'TERRAGEN':
                    print("无效的Terragen文件")
                    return None
                    
                # 解析版本和参数
                version = int.from_bytes(header[8:12], byteorder='little')
                width = int.from_bytes(header[12:16], byteorder='little')
                
                # 读取高度数据
                heights = []
                min_height = float('inf')
                max_height = float('-inf')
                
                while True:
                    chunk_header = f.read(4)
                    if not chunk_header:
                        break
                        
                    chunk_type = chunk_header
                    chunk_size = int.from_bytes(f.read(4), byteorder='little')
                    
                    if chunk_type == b'ALTW':
                        # 高度数据块
                        for _ in range(chunk_size // 2):
                            height = int.from_bytes(f.read(2), byteorder='little', signed=True)
                            heights.append(height)
                            min_height = min(min_height, height)
                            max_height = max(max_height, height)
                    else:
                        # 跳过其他块
                        f.seek(chunk_size, 1)
                        
                return {
                    'width': width,
                    'height': width,  # Terragen通常是方形
                    'heights': heights,
                    'min_height': min_height,
                    'max_height': max_height
                }
                
        except Exception as e:
            print(f"解析Terragen文件时出错: {e}")
            return None
            
    def export_terrain_mesh(self, terrain_info, output_path, format='obj'):
        """
        导出地形为网格文件（OBJ等格式）
        """
        print(f"导出地形网格为 {format} 格式")
        # 这里可以实现将地形导出为网格文件的功能
        return False