import shutil
import tempfile
import tracemalloc
import unittest
from pathlib import Path

import numpy as np

from mapgen import pipeline as P
from mapgen import render as RN
from mapgen.testing import small_world


def _files(root: Path) -> dict:
    return {p.relative_to(root).as_posix(): p.read_bytes()
            for p in sorted(root.rglob("*")) if p.is_file() and "cache" not in p.parts}


class LowMemoryTest(unittest.TestCase):
    def test_low_memory_build_writes_identical_files(self):
        # oracle: the normal-mode build of the same world; every output byte must match
        outs = []
        for low in (False, True):
            tmp = Path(tempfile.mkdtemp())
            try:
                small_world(tmp)
                P.build(tmp, 2, log=lambda m: None, low_memory=low)
                outs.append({**_files(tmp / "out"), **_files(tmp / "previews")})
            finally:
                shutil.rmtree(tmp)
        self.assertEqual(sorted(outs[0]), sorted(outs[1]))
        for k in outs[0]:
            self.assertEqual(outs[0][k], outs[1][k], k)

    def test_chunked_relief_matches_and_uses_less_memory(self):
        rng = np.random.default_rng(0)
        H, W = 512, 1024
        z = rng.normal(0, 2000, (H, W))
        hs = rng.random((H, W))
        zone = rng.integers(0, 38, (H, W)).astype(np.int16)
        ground = rng.integers(0, 9, (H, W)).astype(np.int8)
        ice = rng.integers(0, 5, (H, W)).astype(np.int8)
        lake = rng.random((H, W)) < 0.05
        peaks = []
        outs = []
        for low in (False, True):
            tracemalloc.start()
            outs.append(RN.relief_rgb(z, hs, zone, ground, ice, lake, low_memory=low))
            peaks.append(tracemalloc.get_traced_memory()[1])
            tracemalloc.stop()
        np.testing.assert_array_equal(outs[0], outs[1])
        self.assertLess(peaks[1], 0.4 * peaks[0])                     # measured ≈ 0.28


class ProjectionMemoryTest(unittest.TestCase):
    def test_sampling_converts_one_channel_at_a_time(self):
        from mapgen import projections as PJ
        rng = np.random.default_rng(1)
        img = rng.integers(0, 256, (512, 1024, 3), dtype=np.uint8)
        lat = rng.uniform(-90, 90, (64, 64))
        lon = rng.uniform(-180, 180, (64, 64))
        whole = img.astype(np.float64)                        # oracle: the whole-image float conversion
        want = np.stack([PJ.sample_equirect(whole[..., k], lat, lon) for k in range(3)], axis=-1)
        del whole
        tracemalloc.start()
        got = PJ._sample_rgb(img, lat, lon)
        peak = tracemalloc.get_traced_memory()[1]
        tracemalloc.stop()
        np.testing.assert_array_equal(got, want)
        self.assertLess(peak, 0.5 * img.size * 8)               # no full-size float64 copy


class HillshadeTest(unittest.TestCase):
    def test_chunked_hillshade_matches_and_uses_less_memory(self):
        z = np.random.default_rng(2).normal(0, 1500, (700, 1400))        # 700 rows: chunk edges fall mid-raster
        peaks, outs = [], []
        for low in (False, True):
            tracemalloc.start()
            outs.append(RN.hillshade(z, 6371.0, low_memory=low))
            peaks.append(tracemalloc.get_traced_memory()[1])
            tracemalloc.stop()
        np.testing.assert_array_equal(outs[0], outs[1])
        self.assertLess(peaks[1], 0.5 * peaks[0])