import json
import os
import re
import tempfile
from datetime import datetime
from ftplib import FTP

import numpy as np
import pygame
from scipy.signal import convolve2d

from src.utils.minecraft.chunk import Chunk, MalformedChunk, OldChunk
from src.utils.minecraft.region import Region
from src.utils.paths import CONFIG_DIR, CACHE_DIR, get_project_path

CHUNK_SIZE = 16
CHUNKS_IN_REGION = 32
REGION_SIZE = CHUNK_SIZE * CHUNKS_IN_REGION
GRADIENT_MAT = np.array([
    [0, 0, 0],
    [0, 1, 0.5],
    [0, -0.5, -1]
])
GRADIENT_RANGE = 2
GAUSSIAN_MAT = np.array([
    [1, 4, 7, 4, 1],
    [4, 16, 26, 16, 4],
    [7, 26, 41, 26, 7],
    [4, 16, 26, 16, 4],
    [1, 4, 7, 4, 1],
])
GAUSSIAN_FACT = GAUSSIAN_MAT.sum()


class AutoMapper:
    CONFIG_PATH = os.path.join(CONFIG_DIR, "mapper.json")
    CACHE_PATH = os.path.join(CACHE_DIR, "regions.txt")
    COLORS_PATH = os.path.join(CACHE_DIR, "colors.json")
    MAPS_DIR = os.path.join(CACHE_DIR, "maps")
    MAP_SIZE = 1024

    def __init__(self):
        self.config: FTPConfig = FTPConfig(self.CONFIG_PATH)
        self.ftp: FTP = FTP(self.config.HOST)
        self.regions: list[tuple[int, int]] = []
        self.temp_dir: tempfile.TemporaryDirectory[str] = tempfile.TemporaryDirectory()
        self.colors: dict[str, tuple[float, float, float]] = {}
        self.cache: dict[tuple[int, int], int] = {}
        self.colormaps: dict[str, dict[str, tuple[float, float, float]]] = {}
        self.available_regions: list[tuple[int, int]] = []
        self.load_colors()
        self.load_cache()
        self.load_colormaps()
        self.no_color: set[str] = set()

    def __enter__(self):
        self.ftp.login(self.config.USERNAME, self.config.PASSWORD)
        self.ftp.cwd(self.config.DIR)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.ftp.close()

    def load_colors(self) -> None:
        with open(self.COLORS_PATH, "r") as f:
            self.colors = json.load(f)

    def load_cache(self) -> None:
        if os.path.exists(self.CACHE_PATH):
            with open(self.CACHE_PATH, "r") as f:
                lines = filter(lambda l: len(l.strip()) != 0, f.read().splitlines())
                self.cache = {}
                for line in lines:
                    rx, rz, ts = map(int, line.split(" "))
                    self.cache[(rx, rz)] = ts

    def load_colormaps(self) -> None:
        biome_ids = {}
        with open(get_project_path("res", "biome_ids.txt")) as f:
            lines = f.read().splitlines()
            for line in lines:
                if len(line.strip()) == 0:
                    continue
                biome_name, biome_id = line.split(" ")
                biome_ids[f"minecraft:{biome_name}"] = int(biome_id)

        self.colormaps = {}
        maps_dir = get_project_path("res", "colormaps")
        for filename in os.listdir(maps_dir):
            if not filename.endswith(".png"):
                continue
            img = pygame.image.load(os.path.join(maps_dir, filename))
            block = filename.split(".")[0]
            colormap = {}
            for biome_name, biome_id in biome_ids.items():
                colormap[biome_name] = img.get_at((biome_id, 0))[:3]
            self.colormaps[f"minecraft:{block}"] = colormap

        self.colormaps["minecraft:grass"] = self.colormaps["minecraft:grass_block"]
        self.colormaps["minecraft:tall_grass"] = self.colormaps["minecraft:grass_block"]
        self.colormaps["minecraft:vine"] = self.colormaps["minecraft:oak_leaves"]
        self.colormaps["minecraft:large_fern"] = self.colormaps["minecraft:grass_block"]
        self.colormaps["minecraft:fern"] = self.colormaps["minecraft:grass_block"]
        self.colormaps["minecraft:melon_stem"] = self.colormaps["minecraft:foliage"]
        self.colormaps["minecraft:attached_melon_stem"] = self.colormaps["minecraft:foliage"]
        self.colormaps["minecraft:mangrove_leaves"] = self.colormaps["minecraft:foliage"]
        self.colormaps["minecraft:bubble_column"] = self.colormaps["minecraft:water"]

    def save_cache(self) -> None:
        with open(self.CACHE_PATH, "w") as f:
            for (rx, rz), ts in self.cache.items():
                f.write(f"{rx} {rz} {ts}\n")

    def get_available_regions(self) -> dict[tuple[int, int], int]:
        files = self.ftp.mlsd(facts=["modify"])
        regions = {}
        self.available_regions = []
        for filename, facts in files:
            m = re.match(r"r\.(-?\d+)\.(-?\d+)\.mca", filename)
            if m:
                rx = int(m.group(1))
                rz = int(m.group(2))
                t = datetime.strptime(facts["modify"], "%Y%m%d%H%M%S")
                regions[(rx, rz)] = int(t.timestamp())
                self.available_regions.append((rx, rz))
        return regions

    def fetch_region(self, rx: int, rz: int) -> str:
        name = f"r.{rx}.{rz}.mca"
        outpath = os.path.join(self.temp_dir.name, name)
        with open(outpath, "wb") as f:
            self.ftp.retrbinary(f"RETR {name}", f.write, 1024)

        return outpath

    def map_region(self, rx: int, rz: int) -> tuple[pygame.Surface, np.array]:
        print(f"  [Fetching region ({rx},{rz})]")
        path = self.fetch_region(rx, rz)
        region = Region(path)
        print("  [Rendering]")
        surf = pygame.Surface([REGION_SIZE, REGION_SIZE], pygame.SRCALPHA)
        surf.fill((0, 0, 0, 0))
        heightmap = np.zeros((REGION_SIZE, REGION_SIZE), dtype="uint16")

        for cz in range(CHUNKS_IN_REGION):
            for cx in range(CHUNKS_IN_REGION):
                ox, oy = cx * CHUNK_SIZE, cz * CHUNK_SIZE
                chunk = region.get_chunk(rx * CHUNKS_IN_REGION + cx, rz * CHUNKS_IN_REGION + cz)
                if isinstance(chunk, Chunk):
                    hm = self.render_chunk(chunk, surf, ox, oy)
                    heightmap[oy:oy+CHUNK_SIZE, ox:ox+CHUNK_SIZE] = hm
                elif isinstance(chunk, MalformedChunk):
                    pygame.draw.rect(surf, (92, 47, 32, 200), [ox, oy, CHUNK_SIZE, CHUNK_SIZE])
                elif isinstance(chunk, OldChunk):
                    pygame.draw.rect(surf, (32, 61, 92, 200), [ox, oy, CHUNK_SIZE, CHUNK_SIZE])

        print()
        os.remove(path)

        return surf, heightmap

    def map_region_group(self, x: int, z: int) -> pygame.Surface:
        surf = pygame.Surface([self.MAP_SIZE, self.MAP_SIZE], pygame.SRCALPHA)
        surf.fill((0, 0, 0, 0))
        n_regions = self.MAP_SIZE // REGION_SIZE
        heightmap = np.zeros((self.MAP_SIZE, self.MAP_SIZE), dtype="uint16")
        for dz in range(n_regions):
            for dx in range(n_regions):
                rx, rz = x * n_regions + dx, z * n_regions + dz
                if (rx, rz) in self.available_regions:
                    region, hm = self.map_region(rx, rz)
                    ox, oy = dx * REGION_SIZE, dz * REGION_SIZE
                    surf.blit(region, [ox, oy])
                    heightmap[oy:oy+REGION_SIZE, ox:ox+REGION_SIZE] = hm
                    self.cache[(rx, rz)] = int(datetime.utcnow().timestamp())

        gradient = convolve2d(heightmap, GRADIENT_MAT, boundary="symm")[1:self.MAP_SIZE+1, 1:self.MAP_SIZE+1]
        gradient = gradient.clip(-GRADIENT_RANGE, GRADIENT_RANGE)
        gradient = 1 + gradient / 2 / GRADIENT_RANGE
        gradient = np.array([gradient, gradient, gradient, np.ones(gradient.shape)])
        gradient = np.swapaxes(gradient, 0, 2)
        surf_array = pygame.surfarray.array3d(surf)
        alpha_array = pygame.surfarray.array_alpha(surf).reshape((*surf_array.shape[0:2], 1))
        surf_array = np.concatenate((surf_array, alpha_array), 2)
        res = surf_array * gradient
        res = res.clip(0, 255)
        buf = res.transpose((1, 0, 2)).astype("uint8").tobytes(order="C")
        surf = pygame.image.frombuffer(buf, surf.get_size(), "RGBA")

        return surf

    def render_chunk(self, chunk: Chunk, surf: pygame.Surface, ox: int, oy: int) -> np.array:
        blocks, heightmap, biomes, is_empty = chunk.get_top_blocks()
        if is_empty:
            pygame.draw.rect(surf, (0, 0, 0, 0), [ox, oy, CHUNK_SIZE, CHUNK_SIZE])

        else:
            for z in range(CHUNK_SIZE):
                for x in range(CHUNK_SIZE):
                    color = self.get_color(blocks[z][x], biomes[z][x])
                    surf.set_at((ox + x, oy + z), color)
        return heightmap

    def get_color(self, block: str, biome: str) -> tuple[float, float, float, float]:
        if block in self.colormaps:
            r, g, b = self.colormaps[block][biome]
            r2, g2, b2 = self.colors.get(block, (0, 0, 0, 0))
            return min(255.0, r*r2/255), min(255.0, g*g2/255), min(255.0, b*b2/255), 255
        if block not in self.colors:
            if block not in self.no_color:
                print(f"    no color for {block}")
                self.no_color.add(block)
        return self.colors.get(block, (0, 0, 0, 0))

    def map_world(self) -> None:
        if not os.path.exists(self.MAPS_DIR):
            os.mkdir(self.MAPS_DIR)

        regions = self.get_available_regions()
        groups_to_map = set()
        n_regions = self.MAP_SIZE // REGION_SIZE
        for pos, modified_at in regions.items():
            if pos in self.cache and modified_at <= self.cache[pos]:
                continue

            gx = pos[0] // n_regions
            gz = pos[1] // n_regions
            groups_to_map.add((gx, gz))

        n_groups = len(groups_to_map)
        print(f"Groups to map: {n_groups}")
        proceed = input("Proceed ? y/[N] ")
        if proceed.strip().lower() == "y":
            groups_to_map = sorted(groups_to_map, key=lambda g: g[0]*g[0] + g[1]*g[1])
            for i, (gx, gz) in enumerate(groups_to_map):
                print(f"[Mapping group ({gx}, {gz}) ({i+1}/{n_groups})]")
                try:
                    surf = self.map_region_group(gx, gz)

                    pygame.image.save(surf, os.path.join(self.MAPS_DIR, f"map_{gx}_{gz}.png"))
                    self.save_cache()
                except Exception as e:
                    raise e


class FTPConfig:
    HOST = ""
    USERNAME = ""
    PASSWORD = ""
    DIR = ""

    def __init__(self, path: str):
        self._path: str = path

        self.load()

    def load(self) -> None:
        if os.path.exists(self._path):
            with open(self._path, "r") as f:
                config = json.load(f)

            self.HOST = config["host"]
            self.USERNAME = config["username"]
            self.PASSWORD = config["password"]
            self.DIR = config["dir"]

    def save(self) -> None:
        with open(self._path, "w") as f:
            json.dump({
                "host": self.HOST,
                "username": self.USERNAME,
                "password": self.PASSWORD,
                "dir": self.DIR
            }, f, indent=4)


if __name__ == '__main__':
    pygame.init()
    # x = -1
    # z = 0
    with AutoMapper() as mapper:
        # mapper.get_available_regions()
        # surf = mapper.map_region_group(x, z)
        # pygame.image.save(surf, f"/tmp/map_{x}_{z}.png")
        mapper.map_world()
