# -*- coding: utf-8 -*-
"""
Export Handlers for Advanced Map Downloader

This module provides functions to export rendered maps to various
georeferenced formats: GeoTIFF, PNG+World, JPEG+World, BMP+World.
"""

import os
from qgis.PyQt.QtCore import QSize
from qgis.PyQt.QtGui import QImage
from qgis.core import (
    QgsCoordinateReferenceSystem,
    QgsMapSettings,
    QgsMapRendererParallelJob,
    QgsRectangle
)
from osgeo import gdal, osr
import numpy as np


def render_map(layer, extent, width_px, height_px, crs='EPSG:3857'):
    """
    Render a map layer to a QImage.
    
    :param layer: QgsMapLayer to render
    :param extent: QgsRectangle extent in the specified CRS
    :param width_px: Output width in pixels
    :param height_px: Output height in pixels
    :param crs: CRS of the extent (default EPSG:3857)
    :returns: QImage
    """
    crs_obj = QgsCoordinateReferenceSystem(crs)
    
    # Set up map settings
    map_settings = QgsMapSettings()
    map_settings.setLayers([layer])
    map_settings.setExtent(extent)
    map_settings.setOutputSize(QSize(width_px, height_px))
    map_settings.setDestinationCrs(crs_obj)
    
    # Create and render the image
    image = QImage(width_px, height_px, QImage.Format_ARGB32_Premultiplied)
    image.fill(0)
    
    render = QgsMapRendererParallelJob(map_settings)
    render.start()
    render.waitForFinished()
    
    return render.renderedImage()


def export_geotiff(image, extent, gsd, output_path):
    """
    Export an image to a georeferenced GeoTIFF file using GDAL.
    
    :param image: QImage to export
    :param extent: QgsRectangle in EPSG:3857
    :param gsd: Ground sampling distance (meters/pixel)
    :param output_path: Output file path
    """
    width = image.width()
    height = image.height()
    
    # Convert QImage to numpy array
    image = image.convertToFormat(QImage.Format_RGBA8888)
    ptr = image.bits()
    ptr.setsize(image.byteCount())
    arr = np.array(ptr).reshape(height, width, 4)
    
    # Create GeoTIFF using GDAL
    driver = gdal.GetDriverByName('GTiff')
    dataset = driver.Create(
        output_path, width, height, 4, gdal.GDT_Byte,
        options=['COMPRESS=LZW', 'TILED=YES']
    )
    
    if dataset is None:
        raise Exception("Failed to create GeoTIFF file")
    
    # Set geotransform
    # [top_left_x, pixel_width, rotation, top_left_y, rotation, -pixel_height]
    geotransform = [
        extent.xMinimum(),  # top-left x
        gsd,                # pixel width (west-east)
        0,                  # rotation
        extent.yMaximum(),  # top-left y
        0,                  # rotation
        -gsd                # pixel height (negative for north-up)
    ]
    dataset.SetGeoTransform(geotransform)
    
    # Set projection (EPSG:3857)
    srs = osr.SpatialReference()
    srs.ImportFromEPSG(3857)
    dataset.SetProjection(srs.ExportToWkt())
    
    # Write the image data (RGBA bands)
    for i in range(4):
        band = dataset.GetRasterBand(i + 1)
        band.WriteArray(arr[:, :, i])
        band.FlushCache()
    
    # Close the dataset
    dataset = None


def write_world_file(output_path, extent, width_px, height_px, extension):
    """
    Write a world file for georeferencing.
    
    World files contain 6 lines:
    1. Pixel size in x direction (meters/pixel)
    2. Rotation about y axis (usually 0)
    3. Rotation about x axis (usually 0)
    4. Pixel size in y direction (negative, meters/pixel)
    5. X coordinate of the center of the upper left pixel
    6. Y coordinate of the center of the upper left pixel
    
    :param output_path: Base path (without world file extension)
    :param extent: QgsRectangle in EPSG:3857
    :param width_px: Width in pixels
    :param height_px: Height in pixels
    :param extension: World file extension (e.g., 'pgw', 'jgw', 'bpw')
    """
    # Calculate pixel size
    pixel_width = extent.width() / width_px
    pixel_height = extent.height() / height_px
    
    # Center of upper-left pixel
    center_x = extent.xMinimum() + (pixel_width / 2)
    center_y = extent.yMaximum() - (pixel_height / 2)
    
    # World file path
    base_path = os.path.splitext(output_path)[0]
    world_path = f"{base_path}.{extension}"
    
    # Write world file
    with open(world_path, 'w') as f:
        f.write(f"{pixel_width:.12f}\n")  # Pixel size in x
        f.write("0\n")                      # Rotation about y
        f.write("0\n")                      # Rotation about x
        f.write(f"{-pixel_height:.12f}\n")  # Pixel size in y (negative)
        f.write(f"{center_x:.6f}\n")        # X of upper-left center
        f.write(f"{center_y:.6f}\n")        # Y of upper-left center


def export_png_with_world(image, extent, output_path):
    """
    Export an image to PNG with a world file (.pgw).
    
    :param image: QImage to export
    :param extent: QgsRectangle in EPSG:3857
    :param output_path: Output file path (should end with .png)
    """
    # Save PNG
    if not output_path.lower().endswith('.png'):
        output_path = os.path.splitext(output_path)[0] + '.png'
    
    image.save(output_path, 'PNG')
    
    # Write world file
    write_world_file(output_path, extent, image.width(), image.height(), 'pgw')


def export_jpeg_with_world(image, extent, output_path):
    """
    Export an image to JPEG with a world file (.jgw).
    
    :param image: QImage to export
    :param extent: QgsRectangle in EPSG:3857
    :param output_path: Output file path (should end with .jpg or .jpeg)
    """
    # Save JPEG
    if not (output_path.lower().endswith('.jpg') or output_path.lower().endswith('.jpeg')):
        output_path = os.path.splitext(output_path)[0] + '.jpg'
    
    # Convert to RGB (JPEG doesn't support alpha)
    image_rgb = image.convertToFormat(QImage.Format_RGB888)
    image_rgb.save(output_path, 'JPEG', 95)  # Quality 95
    
    # Write world file
    write_world_file(output_path, extent, image.width(), image.height(), 'jgw')


def export_bmp_with_world(image, extent, output_path):
    """
    Export an image to BMP with a world file (.bpw).
    
    :param image: QImage to export
    :param extent: QgsRectangle in EPSG:3857
    :param output_path: Output file path (should end with .bmp)
    """
    # Save BMP
    if not output_path.lower().endswith('.bmp'):
        output_path = os.path.splitext(output_path)[0] + '.bmp'
    
    # Convert to RGB (BMP typically doesn't use alpha)
    image_rgb = image.convertToFormat(QImage.Format_RGB888)
    image_rgb.save(output_path, 'BMP')
    
    # Write world file
    write_world_file(output_path, extent, image.width(), image.height(), 'bpw')


def export_map(image, extent, gsd, output_path, export_format='geotiff'):
    """
    Unified export function that handles all formats.
    
    :param image: QImage to export
    :param extent: QgsRectangle in EPSG:3857
    :param gsd: Ground sampling distance (meters/pixel)
    :param output_path: Output file path
    :param export_format: One of 'geotiff', 'png_world', 'jpeg_world', 'bmp_world'
    """
    if export_format == 'geotiff':
        # Ensure .tif extension
        if not output_path.lower().endswith(('.tif', '.tiff')):
            output_path = os.path.splitext(output_path)[0] + '.tif'
        export_geotiff(image, extent, gsd, output_path)
    
    elif export_format == 'png_world':
        export_png_with_world(image, extent, output_path)
    
    elif export_format == 'jpeg_world':
        export_jpeg_with_world(image, extent, output_path)
    
    elif export_format == 'bmp_world':
        export_bmp_with_world(image, extent, output_path)
    
    else:
        raise ValueError(f"Unknown export format: {export_format}")
    
    return output_path


# Export format choices for UI
EXPORT_FORMATS = {
    'GeoTIFF (.tif)': 'geotiff',
    'PNG + World File (.png + .pgw)': 'png_world',
    'JPEG + World File (.jpg + .jgw)': 'jpeg_world',
    'BMP + World File (.bmp + .bpw)': 'bmp_world'
}

# File filters for save dialog
FILE_FILTERS = {
    'geotiff': 'GeoTIFF Files (*.tif *.tiff)',
    'png_world': 'PNG Files (*.png)',
    'jpeg_world': 'JPEG Files (*.jpg *.jpeg)',
    'bmp_world': 'BMP Files (*.bmp)'
}
