# -*- 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, QEventLoop, QTimer
from qgis.PyQt.QtGui import QImage, QColor
from qgis.PyQt.QtWidgets import QApplication
from qgis.core import (
    QgsCoordinateReferenceSystem,
    QgsMapSettings,
    QgsMapRendererCustomPainterJob
)
from osgeo import gdal, osr
import numpy as np


def _is_image_blank(image):
    """
    Check if an image is blank (all transparent or single color).

    :param image: QImage to check
    :returns: True if image appears blank
    """
    if image is None or image.isNull():
        return True

    # Sample a few pixels to check if the image has content
    width = image.width()
    height = image.height()

    if width <= 0 or height <= 0:
        return True

    # Check center and corner pixels
    sample_points = [
        (width // 2, height // 2),  # center
        (width // 4, height // 4),  # top-left quarter
        (3 * width // 4, 3 * height // 4),  # bottom-right quarter
    ]

    first_pixel = None
    for x, y in sample_points:
        if 0 <= x < width and 0 <= y < height:
            pixel = image.pixelColor(x, y)
            if first_pixel is None:
                first_pixel = pixel
            # If we find a non-transparent pixel that differs, image has
            # content
            if pixel.alpha() > 0 and pixel != first_pixel:
                return False

    # Check if all sampled pixels are fully transparent
    if first_pixel and first_pixel.alpha() == 0:
        return True

    return False


def render_map(layer, extent, width_px, height_px,
               crs='EPSG:3857', max_retries=5, retry_delay_ms=1000):
    """
    Render a map layer to a QImage with retry support for XYZ tile layers.

    XYZ tile layers need time to fetch tiles from the network. This function
    implements a retry mechanism that allows tiles to load between attempts.

    :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)
    :param max_retries: Maximum number of render attempts (default 5)
    :param retry_delay_ms: Delay between retries in milliseconds (default 1000)
    :returns: QImage
    :raises: Exception if rendering fails after all retries
    """
    from qgis.PyQt.QtGui import QPainter

    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)

    rendered_image = None
    layer_name = layer.name() if layer else "Unknown"

    for attempt in range(max_retries):
        # Create fresh image for each attempt
        image = QImage(
            QSize(
                width_px,
                height_px),
            QImage.Format.Format_ARGB32_Premultiplied)
        image.fill(QColor(0, 0, 0, 0))  # Transparent background

        # Use custom painter job for more control
        painter = QPainter(image)
        render_job = QgsMapRendererCustomPainterJob(map_settings, painter)
        render_job.start()
        render_job.waitForFinished()
        painter.end()

        # Check if we got valid content
        if not _is_image_blank(image):
            rendered_image = image
            break

        # If this wasn't the last attempt, wait and let tiles load
        if attempt < max_retries - 1:
            # Process events to allow tile loading
            loop = QEventLoop()
            QTimer.singleShot(retry_delay_ms, loop.quit)
            loop.exec()

            # Also process application events
            QApplication.processEvents()

    # Final validation
    if rendered_image is None or _is_image_blank(rendered_image):
        raise Exception(
            f"Failed to render map from layer '{layer_name}' after {max_retries} attempts. "
            "This may occur if the map source (e.g., OpenTopoMap) failed to load tiles. "
            "Try a different map source, check your internet connection, or try zooming "
            "to the area first in QGIS to pre-cache the tiles."
        )

    return rendered_image


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
    """
    # Validate image
    if image is None or image.isNull():
        raise Exception("Invalid image: Image is None or null")

    width = image.width()
    height = image.height()

    if width <= 0 or height <= 0:
        raise Exception(f"Invalid image dimensions: {width}x{height}")

    # Convert QImage to numpy array
    image = image.convertToFormat(QImage.Format.Format_RGBA8888)
    ptr = image.bits()

    # Check if bits() returned valid data
    if ptr is None:
        raise Exception("Failed to get image data: image.bits() returned None")

    # Use sizeInBytes() if available (newer QGIS), fallback to byteCount()
    try:
        byte_size = image.sizeInBytes()
    except AttributeError:
        byte_size = image.byteCount()

    ptr.setsize(byte_size)
    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.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.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)'
}
