#!/usr/bin/env python3
"""
Dubai Metaverse - GIS to Unreal Terrain Conversion Script
Converts elevation data (DEM/GeoTIFF) to Unreal Engine terrain format
"""

import os
import sys
import argparse
from typing import Tuple, Optional

try:
    import rasterio
    import numpy as np
except ImportError:
    print("Error: Required packages not installed.")
    print("Install with: pip install rasterio numpy")
    sys.exit(1)


def load_elevation_data(dem_file: str) -> Tuple[np.ndarray, dict]:
    """
    Load elevation data from GeoTIFF file.
    
    Returns:
        elevation_data: NumPy array of elevation values
        metadata: Dictionary with geospatial metadata
    """
    try:
        with rasterio.open(dem_file) as src:
            elevation_data = src.read(1)  # Read first band
            metadata = {
                'width': src.width,
                'height': src.height,
                'crs': src.crs,
                'transform': src.transform,
                'bounds': src.bounds
            }
            return elevation_data, metadata
    except Exception as e:
        print(f"Error loading elevation data: {e}")
        sys.exit(1)


def normalize_elevation(elevation_data: np.ndarray, min_elev: float, max_elev: float) -> np.ndarray:
    """
    Normalize elevation data to 0-1 range for Unreal Engine.
    
    Unreal Engine uses 0-1 normalized height values.
    """
    # Clip to min/max range
    elevation_data = np.clip(elevation_data, min_elev, max_elev)
    
    # Normalize to 0-1
    normalized = (elevation_data - min_elev) / (max_elev - min_elev)
    
    return normalized


def export_heightmap(normalized_data: np.ndarray, output_file: str, format: str = 'raw'):
    """
    Export normalized elevation data as heightmap.
    
    Formats:
    - 'raw': Raw binary format (16-bit)
    - 'png': PNG format (16-bit grayscale)
    """
    # Convert to 16-bit integer (0-65535)
    heightmap = (normalized_data * 65535).astype(np.uint16)
    
    if format == 'raw':
        heightmap.tofile(output_file)
        print(f"✓ Exported heightmap to {output_file} (RAW format)")
    elif format == 'png':
        try:
            from PIL import Image
            # Convert to 16-bit PNG
            img = Image.fromarray(heightmap, mode='I;16')
            img.save(output_file)
            print(f"✓ Exported heightmap to {output_file} (PNG format)")
        except ImportError:
            print("Warning: PIL not installed, falling back to RAW format")
            heightmap.tofile(output_file)
            print(f"✓ Exported heightmap to {output_file} (RAW format)")
    else:
        print(f"Error: Unknown format '{format}'")
        sys.exit(1)


def main():
    parser = argparse.ArgumentParser(description='Convert GIS elevation data to Unreal terrain')
    parser.add_argument('input', help='Input DEM/GeoTIFF file')
    parser.add_argument('--output', '-o', default='data/processed/terrain_heightmap.raw',
                       help='Output heightmap file')
    parser.add_argument('--format', choices=['raw', 'png'], default='raw',
                       help='Output format (raw or png)')
    parser.add_argument('--min-elev', type=float, help='Minimum elevation (meters)')
    parser.add_argument('--max-elev', type=float, help='Maximum elevation (meters)')
    
    args = parser.parse_args()
    
    if not os.path.exists(args.input):
        print(f"Error: Input file not found: {args.input}")
        sys.exit(1)
    
    # Create output directory if it doesn't exist
    os.makedirs(os.path.dirname(args.output), exist_ok=True)
    
    print("==========================================")
    print("Dubai Metaverse - GIS to Unreal Terrain")
    print("==========================================")
    print("")
    print(f"Loading elevation data: {args.input}")
    
    elevation_data, metadata = load_elevation_data(args.input)
    
    print(f"Data dimensions: {metadata['width']} x {metadata['height']}")
    print(f"Bounds: {metadata['bounds']}")
    print("")
    
    # Determine elevation range
    if args.min_elev is None:
        min_elev = float(np.nanmin(elevation_data))
    else:
        min_elev = args.min_elev
    
    if args.max_elev is None:
        max_elev = float(np.nanmax(elevation_data))
    else:
        max_elev = args.max_elev
    
    print(f"Elevation range: {min_elev:.2f}m to {max_elev:.2f}m")
    print("")
    
    print("Normalizing elevation data...")
    normalized = normalize_elevation(elevation_data, min_elev, max_elev)
    
    print("Exporting heightmap...")
    export_heightmap(normalized, args.output, args.format)
    
    print("")
    print("==========================================")
    print("Conversion Complete")
    print("==========================================")
    print("")
    print("Next steps:")
    print("1. Import heightmap to Unreal Engine")
    print("2. Create landscape from heightmap")
    print("3. Adjust landscape material and settings")
    print("")
    print(f"Heightmap file: {args.output}")
    print(f"Dimensions: {metadata['width']} x {metadata['height']}")
    print("")


if __name__ == '__main__':
    main()