import argparse
import sys
from argparse import Namespace
from math import ceil, log
from pathlib import Path

import numpy as np
import numpy.typing as npt
from PIL import Image
from tqdm import tqdm


def pad_image_to_fit_tile_size_and_make_transparent(
    source: npt.NDArray, tile_size: int
) -> npt.NDArray:
    """
    Pad an image to fit a tile size with transparent pixels. Make the image transparent if it was not already.
    :param source: the source image
    :param tile_size: the size of the tiles the image will be split into
    :return: An image whose width and height are multiples of the tile size
    """

    height, width, _ = source.shape
    missing_in_y = -height % tile_size
    missing_in_x = -width % tile_size

    # Make the image transparent if it was not already
    if source.shape[2] == 3:
        transparency_layer = np.full((height, width, 1), 255, dtype=source.dtype)
        source = np.concatenate([source, transparency_layer], axis=2)  # (H, W, 4)

    # right padding
    right_pad = np.full((height, missing_in_x, 4), [0, 0, 0, 0], dtype=source.dtype)
    source = np.concatenate((source, right_pad), axis=1)

    # bottom padding
    bottom_pad = np.full(
        (missing_in_y, width + missing_in_x, 4), [0, 0, 0, 0], dtype=source.dtype
    )
    return np.concatenate((source, bottom_pad), axis=0)


def get_tiles(source: npt.NDArray, tile_size: int) -> npt.NDArray:
    """
    Generate square tiles from an image
    If the image width or height is not a multiple of tile_size the right side tiles and bottom tiles will be padded with black pixels
    :param source: The source image
    :param tile_size: The size of the tile in pixels
    :return: An array with shape (nb_tiles_y, nb_tiles_x, tile_size, tile_size [, pixel size (3 if rgb)] )
    """
    source = pad_image_to_fit_tile_size_and_make_transparent(source, tile_size)
    height, width, nb_colors = source.shape
    nb_tiles_x = width // tile_size
    nb_tiles_y = height // tile_size
    return source.reshape(
        (nb_tiles_y, tile_size, nb_tiles_x, tile_size, nb_colors)
    ).swapaxes(1, 2)


def save_to_directory(
    tiled_image: npt.NDArray,
    target_directory: Path,
    image_format: str,
    progress_bar: tqdm | None = None,
) -> None:
    """
    Given a tiled map obtained via get_tiles, save it to a directory that can later be used by mapping libraries such as leaflet
    Resulting directory structure:
    └── target_directory
        ├── x0
        │   ├── y0.png
        │   └── y1.png
        └── x1
            ├── y0.png
            └── y1.png
    :param tiled_image: The tiled image to save
    :param target_directory: The target directory where the tiles will be saved. Must not exist.
    :param image_format: The image format of the tiles. Tested for .webp and .png
    :param progress_bar: Optional, if a progress bar is given it will be ticked for each image saved
    """
    target_directory.mkdir()

    # Tiled_image format is row, col, tile (y,x)
    # Tiling system usually work with col, row, tile (x,y)
    # Let's swap x and y
    tiled_image_x_then_y = tiled_image.swapaxes(0, 1)
    for x_index, column_of_tiles in enumerate(tiled_image_x_then_y):
        subdirectory = target_directory / str(x_index)
        subdirectory.mkdir()
        for y_index, tile in enumerate(column_of_tiles):
            Image.fromarray(tile).save(subdirectory / f"{y_index}{image_format}")
            if progress_bar:
                progress_bar.update(1)


def get_maximum_zoom_level(source_map_shape: tuple, tile_size: int) -> int:
    """
    Compute the maximum zoom level to use with a given map
    Knowing that the maximum zoom level should contain the map in its original size
    :param source_map_shape: the shape of the source map
    :param tile_size: the tile size
    """
    # At zoom level z, the world is split into 2^z columns (x = 0 … 2^z - 1) and 2^z rows (y = 0 … 2^z - 1).
    # If the map (in full size) should be split into only 1 column and 1 row because the tile size is lower than the original then the max zoom level should be 0
    # It it should be split into 2 columns or 2 rows then the max zoom level should be 1
    # And so on
    largest_side_length_in_px = max(source_map_shape)
    nb_tiles_max = ceil(largest_side_length_in_px / tile_size)
    return ceil(log(nb_tiles_max) / log(2))


def resize_image_for_zoom_level(
    source_map: npt.NDArray, zoom_level: int, maximum_zoom_level: int
) -> npt.NDArray:
    """
    Resize an image for a zoom level, assuming the maximum zoom level is the zoom at which the image will not be resized
    At maximum_zoom_level - 1, each image side length in pixels will be divided by 2
    At maximum_zoom_level -2, each image side will be divided by 4
    And so on
    :param source_map: the source map image
    :param zoom_level: the zoom level the image must be resized to
    :param maximum_zoom_level: The maximum_zoom_level for the map
    :return: the image resized to match the zoom level
    """
    resize_factor = 2 ** (maximum_zoom_level - zoom_level)
    if resize_factor == 1:
        return source_map

    original_height, original_width, _ = source_map.shape
    target_height = original_height // resize_factor
    target_width = original_width // resize_factor
    return np.asarray(
        Image.fromarray(source_map).resize((target_width, target_height)), dtype="uint8"
    )


def get_tiles_for_all_zoom_levels(
    source_map: npt.NDArray, tile_size: int
) -> list[npt.NDArray]:
    """
    Generate the tiles for a source map for all zoom levels
    :param source_map: The map we want to generate tiles from
    :param tile_size: each tile size
    :return: An array containing all tiles for all required zoom_levels, of shape (nb_zooms, nb_tiles_y, nb_tiles_x, tile_size, tile_size [, pixel_size (3 if RGB)]
    """
    maximum_zoom_level = get_maximum_zoom_level(source_map.shape, tile_size)
    res = []
    for zoom_level in tqdm(
        range(maximum_zoom_level + 1), desc="Generating map for each zoom level"
    ):
        image_resized_for_zoom_level = resize_image_for_zoom_level(
            source_map, zoom_level, maximum_zoom_level
        )
        res.append(get_tiles(image_resized_for_zoom_level, tile_size))
    return res


def parse_args() -> Namespace:
    parser = argparse.ArgumentParser(description="Create map tiles from a map source.")
    parser.add_argument(
        "map_source",
        help="Path or identifier of the map source to generate tiles from",
    )
    parser.add_argument(
        "--tile-size", "-t", help="Size of a tile", default=256, type=int
    )
    parser.add_argument(
        "--output-directory",
        "-o",
        help="Output directory where to store the tiles, must not exist",
        required=True,
    )
    parser.add_argument(
        "--image-format",
        "-f",
        help="Output image format. Tested for .png and .webp",
        default=".webp",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    map_source: str = args.map_source
    map_image = np.asarray(Image.open(map_source), dtype="uint8")
    tiled_maps_by_zoom_level = get_tiles_for_all_zoom_levels(map_image, args.tile_size)
    Path.mkdir(args.output_directory)
    nb_images_to_save = sum(
        tiled_map.shape[0] * tiled_map.shape[1]
        for tiled_map in tiled_maps_by_zoom_level
    )
    with tqdm(desc="Saving generated tiles", total=nb_images_to_save) as progress_bar:
        for zoom_level, tiled_map in enumerate(tiled_maps_by_zoom_level):
            save_to_directory(
                tiled_map,
                Path(args.output_directory) / str(zoom_level),
                image_format=args.image_format,
                progress_bar=progress_bar,
            )


if __name__ == "__main__":
    main()
