import sys
from typing import Optional, Callable, TypeVar, Union, TYPE_CHECKING

from mapproxy.cache.tile import TileCollection
from mapproxy.grid import TileCoord
from mapproxy.grid.meta_grid import MetaTile
from mapproxy.cache.tile import Tile
from mapproxy.image import BaseImageResult
from mapproxy.image.merge import merge_images
from mapproxy.image.tile import TileSplitter
from mapproxy.layer import BlankImageError
from mapproxy.layer.map_layer import MapLayer
from mapproxy.query import MapQuery
from mapproxy.source import SourceError
from mapproxy.util import async_
from mapproxy.util.coverage import Coverage
from mapproxy.util.py import reraise_exception, reraise
if TYPE_CHECKING:
    from mapproxy.cache.tile_manager import TileManager

# TypeVar for _create_threaded method
TTile = TypeVar('TTile', bound=Union[Tile, MetaTile])


class TileCreator:
    def __init__(self, tile_mgr: 'TileManager', dimensions=None, image_merger=None, bulk_meta_tiles=False):
        self.cache = tile_mgr.cache
        self.sources = tile_mgr.sources
        self.grid = tile_mgr.grid
        self.meta_grid = tile_mgr.meta_grid
        self.bulk_meta_tiles = bulk_meta_tiles
        self.tile_mgr = tile_mgr
        self.dimensions = dimensions
        self.image_merger = image_merger

    def is_cached(self, tile: Union[Tile, TileCoord], dimensions=None) -> bool:
        """
        Return True if the tile is cached.
        """
        return self.tile_mgr.is_cached(tile, dimensions=dimensions)

    def is_stale(self, tile: Tile) -> bool:
        """
        Return True if the tile exists in cache and is expired.
        """
        return self.tile_mgr.is_stale(tile)

    def create_tiles(self, tiles: list[Tile]) -> list[Tile]:
        if not self.sources:
            return []
        if not self.meta_grid:
            created_tiles = self._create_single_tiles(tiles)
        elif self.tile_mgr.minimize_meta_requests and len(tiles) > 1:
            # use minimal requests only for multiple tile requests (ie not for TMS)
            meta_tile = self.meta_grid.minimal_meta_tile([t.coord for t in tiles])
            created_tiles = self._create_meta_tile(meta_tile)
        else:
            meta_tiles = []
            meta_bboxes = set()
            for tile in tiles:
                assert tile.coord is not None
                meta_tile = self.meta_grid.meta_tile(tile.coord)
                if meta_tile.bbox not in meta_bboxes:
                    meta_tiles.append(meta_tile)
                    meta_bboxes.add(meta_tile.bbox)

            created_tiles = self._create_meta_tiles(meta_tiles)

        return created_tiles

    def _create_single_tiles(self, tiles: list[Tile]) -> list[Tile]:
        if self.tile_mgr.concurrent_tile_creators > 1 and len(tiles) > 1:
            return self._create_threaded(self._create_single_tile, tiles)

        created_tiles = []
        for tile in tiles:
            created_tiles.extend(self._create_single_tile(tile))
        return created_tiles

    def _create_threaded(self, create_func: Callable[[TTile], list[Tile]], tiles: list[TTile]) -> list[Tile]:
        result = []
        async_pool = async_.Pool(self.tile_mgr.concurrent_tile_creators)
        for new_tiles in async_pool.imap(create_func, tiles):
            result.extend(new_tiles)
        return result

    def _create_single_tile(self, tile: Tile, dimensions=None) -> list[Tile]:
        assert tile.coord is not None
        tile_bbox = self.grid.tile_bbox(tile.coord)
        query = MapQuery(tile_bbox, self.grid.tile_size, self.grid.srs,
                         self.tile_mgr.request_format, dimensions=self.dimensions)
        with self.tile_mgr.lock(tile):
            if not self.is_cached(tile, dimensions=dimensions):
                image_result = None
                try:
                    image_result = self._query_sources(query)
                # if source is not available, try to serve tile in cache
                except SourceError as e:
                    if self.is_stale(tile):
                        self.cache.load_tile(tile)
                    else:
                        raise reraise_exception(e, sys.exc_info())
                if not image_result:
                    return []
                if image_result.authorize_stale and self.is_stale(tile):
                    # The configuration authorises blank tiles generated by the error_handler
                    # to be replaced by stale tiles from cache.
                    self.cache.load_tile(tile)
                    return [tile]
                if self.tile_mgr.image_opts != image_result.image_opts:
                    # call as_buffer to force conversion into cache format
                    image_result.as_buffer(self.tile_mgr.image_opts)
                image_result.image_opts = self.tile_mgr.image_opts
                tile.image_result = image_result
                tile.cacheable = image_result.cacheable
                tile = self.tile_mgr.apply_tile_filter(tile)
                if image_result.cacheable:
                    self.cache.store_tile(tile)
            else:
                self.cache.load_tile(tile)
        return [tile]

    def _query_sources(self, query: MapQuery) -> Optional[BaseImageResult]:
        """
        Query all sources and return the results as a single ImageResult.
        Multiple sources will be merged into a single image.
        """

        # directly return get_map without merge if ...
        if (len(self.sources) == 1 and
                not self.image_merger and  # no special image_merger (like BandMerger)
                not (self.sources[0].coverage and  # no clipping coverage
                     self.sources[0].coverage.clip and
                     self.sources[0].coverage.intersects(query.bbox, query.srs))):
            try:
                return self.sources[0].get_map(query)
            except BlankImageError:
                return None

        def get_map_from_source(source: MapLayer) -> tuple[Optional[BaseImageResult], Optional[Coverage]]:
            try:
                img = source.get_map(query)
            except BlankImageError:
                return None, None
            else:
                return img, source.coverage

        layers: list[tuple[BaseImageResult, Coverage]] = []
        for layer in async_.imap(get_map_from_source, self.sources):
            if layer[0] is not None:
                layers.append(layer)

        return merge_images(layers, size=query.size, bbox=query.bbox, bbox_srs=query.srs,
                            image_opts=self.tile_mgr.image_opts, merger=self.image_merger)

    def _create_meta_tiles(self, meta_tiles: list[MetaTile]) -> list[Tile]:
        if self.bulk_meta_tiles:
            created_tiles = []
            for meta_tile in meta_tiles:
                created_tiles.extend(self._create_bulk_meta_tile(meta_tile))
            return created_tiles

        if self.tile_mgr.concurrent_tile_creators > 1 and len(meta_tiles) > 1:
            return self._create_threaded(self._create_meta_tile, meta_tiles)

        created_tiles = []
        for meta_tile in meta_tiles:
            created_tiles.extend(self._create_meta_tile(meta_tile))
        return created_tiles

    def _create_meta_tile(self, meta_tile: MetaTile) -> list[Tile]:
        """
        _create_meta_tile queries a single meta tile and splits it into
        tiles.
        """
        tile_size = self.grid.tile_size
        query = MapQuery(meta_tile.bbox, meta_tile.size, self.grid.srs, self.tile_mgr.request_format,
                         dimensions=self.dimensions)
        main_tile = Tile(meta_tile.main_tile_coord)
        with self.tile_mgr.lock(main_tile):
            if not all(self.is_cached(t, dimensions=self.dimensions) for t in meta_tile.tiles if t is not None):
                meta_tile_image = self._query_sources(query)
                if not meta_tile_image:
                    return []
                splitted_tiles = split_meta_tiles(meta_tile_image, meta_tile.tile_patterns,
                                                  tile_size, self.tile_mgr.image_opts)
                splitted_tiles = [self.tile_mgr.apply_tile_filter(t) for t in splitted_tiles]
                if meta_tile_image.cacheable:
                    self.cache.store_tiles(splitted_tiles, dimensions=self.dimensions)
                return splitted_tiles
            # else
        tiles = TileCollection(meta_tile.tiles)
        self.cache.load_tiles(tiles, dimensions=self.dimensions)
        return tiles.tiles

    def _create_bulk_meta_tile(self, meta_tile):
        """
        _create_bulk_meta_tile queries each tile of the meta tile in parallel
        (using concurrent_tile_creators).
        """
        tile_size = self.grid.tile_size
        main_tile = Tile(meta_tile.main_tile_coord)
        with self.tile_mgr.lock(main_tile):
            if not all(self.is_cached(t, dimensions=self.dimensions) for t in meta_tile.tiles if t is not None):
                async_pool = async_.Pool(self.tile_mgr.concurrent_tile_creators)

                def query_tile(coord):
                    try:
                        query = MapQuery(
                            self.grid.tile_bbox(coord), tile_size, self.grid.srs, self.tile_mgr.request_format,
                            dimensions=self.dimensions)
                        tile_image = self._query_sources(query)
                        if tile_image is None:
                            return None

                        if self.tile_mgr.image_opts != tile_image.image_opts:
                            # call as_buffer to force conversion into cache format
                            tile_image.as_buffer(self.tile_mgr.image_opts)

                        tile = Tile(coord, cacheable=tile_image.cacheable)
                        tile.image_result = tile_image
                        tile = self.tile_mgr.apply_tile_filter(tile)
                    except BlankImageError:
                        return None
                    else:
                        return tile

                tiles = []
                for tile_task in async_pool.imap(query_tile,
                                                 [t for t in meta_tile.tiles if t is not None],
                                                 use_result_objects=True,
                                                 ):
                    if tile_task.exception is None:
                        tile = tile_task.result
                        if tile is not None:
                            tiles.append(tile)
                    else:
                        ex = tile_task.exception
                        async_pool.shutdown(True)
                        raise reraise(ex)

                self.cache.store_tiles([t for t in tiles if t.cacheable], dimensions=self.dimensions)
                return tiles

            # else
        tiles = [Tile(coord) for coord in meta_tile.tiles]
        tile_collection = TileCollection(tiles)
        self.cache.load_tiles(tile_collection, dimensions=self.dimensions)
        return tile_collection


def split_meta_tiles(meta_tile: BaseImageResult, tiles: list[tuple[Optional[TileCoord], tuple[int, int]]],
                     tile_size: tuple[int, int], image_opts):
    try:
        # TODO png8
        # if not self.transparent and format == 'png':
        #     format = 'png8'
        splitter = TileSplitter(meta_tile, image_opts)
    except IOError:
        # TODO
        raise
    split_tiles = []
    for tile in tiles:
        tile_coord, crop_coord = tile
        if tile_coord is None:
            continue
        data = splitter.get_tile(crop_coord, tile_size)
        new_tile = Tile(tile_coord, cacheable=meta_tile.cacheable)
        new_tile.image_result = data
        split_tiles.append(new_tile)
    return split_tiles
