import json
import re
import shutil
import tempfile
import unittest
from pathlib import Path

import numpy as np
from PIL import Image

from mapgen import pipeline as P
from mapgen import render as RN
from tests.helpers import small_grid
from mapgen.sphere import east_north
from mapgen.testing import FIXTURE_TOML

from mapgen.testing import TECT, small_world  # noqa: F401  (other tests import them from here)


class RenderTest(unittest.TestCase):
    def test_encode_roundtrip(self):
        v = np.array([-11000.0, 0.0, 8848.3])
        np.testing.assert_allclose(RN.decode(RN.encode(v, 0.5, -12000.0), 0.5, -12000.0), v, atol=0.25)

    def test_full_pipeline_small(self):
        tmp = Path(tempfile.mkdtemp())
        try:
            small_world(tmp)
            ctx = P.build(tmp, 2, log=lambda m: None)
            out = tmp / "out" / "r2"
            meta = json.loads((out / "fields.json").read_text())
            im = Image.open(out / "raster" / meta["continuous"]["elevation"]["file"])
            self.assertEqual(im.size, (256, 128))
            raw = np.array(im)
            e = meta["continuous"]["elevation"]
            z = RN.decode(raw, e["scale"], e["offset"])
            self.assertTrue(-11500 < z.min() < 0 < z.max() < 13000)
            cells = np.load(out / "cells.npz")
            self.assertEqual(len(cells["g_ids"]), ctx.grid.n)
            self.assertIn("holdridge", cells.files)
            self.assertTrue((tmp / "previews" / "r2" / "contact_sheet.png").exists())
            self.assertTrue((out / "geo" / "coast.geojson").exists())
            for f in ("proj_equal_earth.png", "proj_mollweide.png", "globes_sheet.png", "globe_north_pole.png"):
                self.assertTrue((tmp / "previews" / "r2" / f).exists(), f)
            vidx = json.loads((out / "viewer" / "layers.json").read_text())
            self.assertEqual([L["id"] for L in vidx], ["relief", "biomes", "elevation", "temperature", "rainfall",
                                                       "seasonality", "landform", "ground", "ice", "deposits", "plates", "o2",
                                                       "gravity", "pressure", "fire", "seabed", "minerals",
                                                       "bottom_temp", "sediment", "vent_potential", "currents", "sst", "productivity"])
            self.assertEqual(Image.open(out / "viewer" / "biomes.jpg").size, (256, 128))
            self.assertEqual(len(vidx[1]["legend"]["items"]), 38)
            cm = json.loads((out / "cells_meta.json").read_text())
            self.assertEqual(cm["res"], 2)
            self.assertEqual(len(cm["legends"]["holdridge"]), 38)
        finally:
            shutil.rmtree(tmp)

    def test_deterministic(self):
        outs = []
        for _ in range(2):
            tmp = Path(tempfile.mkdtemp())
            try:
                small_world(tmp)
                outs.append(P.build(tmp, 2, stop="erosion", log=lambda m: None).data["elevation_eroded_m"])
            finally:
                shutil.rmtree(tmp)
        np.testing.assert_array_equal(outs[0], outs[1])

    def test_viewer_has_the_sea_floor_and_zone_layers(self):
        from mapgen.testing import built_world
        root = built_world()
        ids = [L["id"] for L in json.loads((root / "out" / "r2" / "viewer" / "layers.json").read_text())]
        for k in ("seabed", "minerals", "bottom_temp", "sediment", "vent_potential", "pressure", "fire"):
            self.assertIn(k, ids)


class DevWidthTest(unittest.TestCase):
    def test_dev_builds_use_dev_raster_width(self):
        tmp = Path(tempfile.mkdtemp())
        try:
            small_world(tmp)
            w = tmp / "config" / "world.toml"
            w.write_text(w.read_text().replace("dev_raster_width = 256", "dev_raster_width = 128"))
            P.build(tmp, 2, log=lambda m: None)                      # res 2 != res_final → dev width
            meta = json.loads((tmp / "out" / "r2" / "fields.json").read_text())
            self.assertEqual(meta["width"], 128)
        finally:
            shutil.rmtree(tmp)


class RiverDrawTest(unittest.TestCase):
    def test_rivers_are_drawn_and_seam_skipped(self):
        from types import SimpleNamespace
        g = SimpleNamespace(n=4, lat=np.array([0.0, 0.0, 0.0, 0.0]), lon=np.array([-90.0, 0.0, 90.0, 179.0]))
        recv = np.array([1, 2, 2, 0])
        img = np.zeros((20, 40, 3), np.uint8)
        out = RN.draw_rivers(img, g, recv, np.array([True, True, False, True]), np.array([2, 2, 0, 1], np.int8))
        row = out[9:11].max(axis=0)
        self.assertTrue(np.all(row[11:29, 2] > 100))           # blue along lon −90..90
        self.assertEqual(int(out[:, 32:, 2].sum()), 0)          # the 179 → −90 segment wraps the seam: not drawn


class RegionalWidthTest(unittest.TestCase):
    def test_finer_than_final_uses_full_width(self):
        tmp = Path(tempfile.mkdtemp())
        try:
            small_world(tmp)
            w = tmp / "config" / "world.toml"
            w.write_text(w.read_text().replace("res_final = 5", "res_final = 1").replace("res_dev = 4", "res_dev = 1")
                         .replace("dev_raster_width = 256", "dev_raster_width = 128"))
            P.build(tmp, 2, log=lambda m: None)                  # res 2 > res_final: regional/finer build
            meta = json.loads((tmp / "out" / "r2" / "fields.json").read_text())
            self.assertEqual(meta["width"], 256)
        finally:
            shutil.rmtree(tmp)


class PixelCoastTest(unittest.TestCase):
    def test_pixel_land_rule(self):
        ocean_k = np.array([[[True, True, True], [False, False, False], [True, False, False], [True, False, False]]])
        z = np.array([[50.0, -30.0, 10.0, -10.0]])
        np.testing.assert_array_equal(RN.pixel_land(ocean_k, z), [[False, True, True, False]])


class ReliefRiverOrderTest(unittest.TestCase):
    def test_first_order_streams_hidden(self):
        from types import SimpleNamespace
        g = SimpleNamespace(n=3, lat=np.zeros(3), lon=np.array([-90.0, 0.0, 90.0]))
        img = np.zeros((20, 40, 3), np.uint8)
        out = RN.draw_rivers(img, g, np.array([1, 2, 2]), np.array([True, True, False]), np.array([1, 1, 0], np.int8))
        self.assertEqual(int(out.sum()), 0)


class RasterRangeTest(unittest.TestCase):
    def test_continuous_rasters_hold_every_configurable_value(self):
        from mapgen import render as RD
        from mapgen.config import ZONE_FIELDS
        top = lambda name: RD.CONTINUOUS[name][1] * 65535 + RD.CONTINUOUS[name][2]
        g_lo = ZONE_FIELDS["gravity_g"][0]
        need = {"pressure": ZONE_FIELDS["pressure_bar"][1], "o2_fraction": ZONE_FIELDS["o2_fraction"][1],
                "fire": ZONE_FIELDS["fire_reactivity"][1], "gravity": ZONE_FIELDS["gravity_g"][1],
                "po2": ZONE_FIELDS["o2_fraction"][1] * ZONE_FIELDS["pressure_bar"][1], "plant_height": 5.0 / g_lo}
        for name, v in need.items():
            self.assertGreaterEqual(top(name), v, name)


class CurrentArrowsTest(unittest.TestCase):
    def test_eastward_current_draws_horizontal_arrows_at_sea_only(self):
        g = small_grid(2)
        e, _ = east_north(g.xyz)
        ocean = g.lat < 0
        base = np.zeros((128, 256, 3), np.uint8)
        out = RN.draw_currents(base.copy(), g, 0.5 * e, ocean)
        ys, xs = np.nonzero(out.any(axis=2))
        self.assertGreater(len(ys), 50)
        self.assertTrue(np.all(ys >= 60))                          # north half (land) untouched (row 64 = equator)
        still = RN.draw_currents(base.copy(), g, np.zeros((g.n, 3)), ocean)
        self.assertFalse(still.any())                              # no current, no arrows

    def test_new_layers_registered(self):
        for name in ("sst", "productivity", "current_speed", "upwelling"):
            self.assertIn(name, RN.CONTINUOUS)
            self.assertIn(name, RN.RAMPS)
        ids = [v[0] for v in RN.VIEWER_LAYERS]
        for vid in ("currents", "sst", "productivity"):
            self.assertIn(vid, ids)


class ByRowsTest(unittest.TestCase):
    """Rasters are coloured and encoded CHUNK rows at a time: the same bytes as in one piece."""
    def test_ramp_and_encode_same_as_whole(self):
        from unittest import mock
        import numpy as np
        from mapgen import render as R
        v = np.random.default_rng(3).normal(0, 2000, size=(300, 37))
        v[5, 5], v[7, 7] = -np.inf, np.inf
        stops = [[0, 0, 0], [10, 200, 30], [255, 255, 255]]
        with mock.patch.object(R, "CHUNK", 10**6):
            want_r, want_e = R._ramp(v, -3000.0, 4000.0, stops), R.encode(v, 0.5, -12000.0)
        with mock.patch.object(R, "CHUNK", 64):
            got_r, got_e = R._ramp(v, -3000.0, 4000.0, stops), R.encode(v, 0.5, -12000.0)
        self.assertEqual(got_r.shape, (300, 37, 3))
        self.assertEqual(got_r.dtype, np.uint8)
        self.assertEqual(got_e.dtype, np.uint16)
        self.assertTrue(np.array_equal(got_r, want_r))
        self.assertTrue(np.array_equal(got_e, want_e))
        # hand-derived: midpoint of a 0..1 ramp between 0 and 200 → 100
        self.assertEqual(R._ramp(np.full((200, 2), 0.5), 0.0, 1.0, [[0, 0, 0], [200, 200, 200]])[150, 1, 0], 100)
        self.assertEqual(int(R.encode(np.full((200, 2), 3.0), 0.5, -12000.0)[199, 0]), 24006)


class CompiledRenderTest(unittest.TestCase):
    def test_ramp_same_bytes_as_numpy_formula(self):
        import numpy as np
        from mapgen import render as R
        rng = np.random.default_rng(7)
        v = np.concatenate([rng.normal(0, 3000, 199_997), [np.inf, -np.inf, 0.0, -0.0, 4000.0, -3000.0, 500.0]])
        v = v.reshape(-1, 7)
        for lo, hi, stops in ((-3000, 4000, [[0, 0, 0], [10, 200, 30], [255, 255, 255]]),
                              (-6500.0, 0.0, [[11, 43, 90], [30, 90, 150], [143, 198, 224]]),
                              (0.1, 0.35, [[1, 2, 3], [250, 9, 77], [3, 255, 100], [255, 255, 255]]),
                              (0, 1, [[0, 0, 0], [255, 255, 255]])):
            st = np.asarray(stops, dtype=np.float64)
            t = np.clip((v - lo) / (hi - lo), 0, 1) * (len(st) - 1)          # the numpy statements, written out
            i = np.minimum(t.astype(np.int64), len(st) - 2)
            f = (t - i)[..., None]
            want = (st[i] * (1 - f) + st[i + 1] * f).astype(np.uint8)
            got = R._ramp(v, lo, hi, stops)
            self.assertEqual(got.dtype, np.uint8)
            self.assertTrue(np.array_equal(got, want), (lo, hi))

    def test_sample_cont_same_as_numpy_sum(self):
        import numpy as np
        from mapgen import render as R
        rng = np.random.default_rng(8)
        v = rng.normal(size=5000) * 10.0 ** rng.integers(-6, 6, 5000)
        for k in (1, 3, 5):
            idx = rng.integers(0, 5000, (300, 170, k)).astype(np.int32)
            w = rng.random((300, 170, k)).astype(np.float32)
            want = np.sum(v[idx] * w, axis=-1)
            self.assertTrue(np.array_equal(R.sample_cont(v, idx, w), want), k)

    def test_writer_finishes_everything_and_raises_errors(self):
        import threading
        from mapgen import render as R
        done = []
        w = R._Writer(threads=2, depth=2)
        for i in range(9):
            w(lambda i=i: done.append(i))
        w.close()
        self.assertEqual(sorted(done), list(range(9)))
        w = R._Writer()
        w(lambda: (_ for _ in ()).throw(OSError("disk full")))
        with self.assertRaises(OSError):
            w.close()
        with self.assertRaises(OSError):                 # surfaced at the latest by the next call over depth
            w = R._Writer(threads=1, depth=1)
            w(lambda: (_ for _ in ()).throw(OSError("disk full")))
            w(lambda: None)