import os
from math import floor
from typing import Optional

import pygame


class ImageHandler:
    def __init__(self, maps_dir: str, base_size: int):
        self.maps_dir: str = maps_dir
        self.base_size: int = base_size
        self.cache: dict = {}
        self.load_base_images()

    def load_base_images(self) -> None:
        cache = {}

        for path in os.listdir(self.maps_dir):
            fullpath = os.path.join(self.maps_dir, path)
            name, x, y = path.split(".")[0].split("_")
            cache[(int(x), int(y))] = pygame.image.load(fullpath).convert_alpha()

        self.cache = {
            1: cache
        }

    def get_for_pos(self, x: int, y: int, zoom: float = 1) -> Optional[pygame.Surface]:
        cache = self.cache.get(zoom, {})
        pos = (x, y)
        if pos not in cache:
            if zoom == 1:
                return None

            # Multiple maps per zone
            elif zoom < 1:
                img = pygame.Surface((self.base_size, self.base_size), pygame.SRCALPHA)
                n_maps = int(1 / zoom)
                sub_size = self.base_size * zoom
                for sub_y in range(n_maps):
                    for sub_x in range(n_maps):
                        sub_img = self.get_for_pos(x * n_maps + sub_x, y * n_maps + sub_y, 1)
                        if sub_img is not None:
                            sub_img = pygame.transform.scale_by(sub_img, zoom)
                            img.blit(sub_img, [sub_x * sub_size, sub_y * sub_size])

            # Zone is part of a map
            else:
                img = self.get_for_pos(floor(x / zoom), floor(y / zoom), 1)
                if img is not None:
                    sub_x = x % zoom
                    sub_y = y % zoom
                    sub_size = self.base_size / zoom
                    img = img.subsurface([sub_x * sub_size, sub_y * sub_size, sub_size, sub_size])
                    img = pygame.transform.scale_by(img, zoom)

            cache[pos] = img

        self.cache[zoom] = cache
        return cache[pos]