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

import numpy as np
from PIL import Image

from mapgen import viewer_export as VE


class ViewerExportTest(unittest.TestCase):
    def test_write_all(self):
        tmp = Path(tempfile.mkdtemp())
        try:
            rgb = np.zeros((20, 40, 3), np.uint8)
            rgb[:, :20] = (200, 30, 30)
            VE.write_all(tmp / "viewer", [
                {"id": "a", "name": "A", "rgb": rgb, "legend": None},
                {"id": "b", "name": "B", "rgb": rgb,
                 "legend": VE.categorical_legend(["x", "y"], np.array([[255, 0, 0], [0, 0, 255]]))},
                {"id": "c", "name": "C", "rgb": rgb,
                 "legend": VE.continuous_legend("m", -10.0, 10.0, [[0, 0, 0], [128, 128, 128], [255, 255, 255]])},
            ])
            idx = json.loads((tmp / "viewer" / "layers.json").read_text())
            self.assertEqual([L["id"] for L in idx], ["a", "b", "c"])
            self.assertEqual(idx[1]["legend"]["items"][1], {"name": "y", "color": "#0000ff"})
            self.assertEqual(idx[2]["legend"]["stops"], [[-10.0, "#000000"], [0.0, "#808080"], [10.0, "#ffffff"]])
            im = Image.open(tmp / "viewer" / "a.jpg")
            self.assertEqual((im.size, im.format), ((40, 20), "JPEG"))
        finally:
            shutil.rmtree(tmp)