# map/tests/test_servecache.py
import os
import tempfile
import time
import unittest
from pathlib import Path
from unittest import mock

import numpy as np

import servecache as SC


class ServeCacheTest(unittest.TestCase):
    def setUp(self):
        self.tmp = Path(tempfile.mkdtemp())
        self.arrays = {"a": np.arange(10, dtype=np.float32), "ids": np.array([3, 1, 2], np.uint64),
                       "m": np.array([True, False]), "xyz": np.random.default_rng(1).normal(size=(5, 3))}

    def test_round_trip_keeps_values_and_dtypes(self):
        d = SC.cache_dir(self.tmp, "abc")
        self.assertEqual(d.name, f"serve-abc-v{SC.VERSION}")
        SC.write(d, self.arrays, {"note": "x"})
        arrays, meta = SC.read(d)
        self.assertEqual(meta["note"], "x")
        for k, v in self.arrays.items():
            self.assertIsInstance(arrays[k], np.memmap)
            self.assertEqual(arrays[k].dtype, v.dtype)
            np.testing.assert_array_equal(arrays[k], v)

    def test_read_loads_more_arrays_than_the_soft_file_limit(self):
        import resource
        soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
        if hard != resource.RLIM_INFINITY and hard < 400:
            self.skipTest("hard open-file limit too low")
        d = SC.cache_dir(self.tmp, "many")
        SC.write(d, {f"f{i}": np.arange(3) + i for i in range(300)}, {})
        got = None
        try:
            resource.setrlimit(resource.RLIMIT_NOFILE, (128, hard))     # each memory map keeps a file open
            got = SC.read(d)
            self.assertIsNotNone(got, "a complete cache must not read as damaged")
            self.assertEqual(int(got[0]["f299"][2]), 301)
        finally:
            del got
            resource.setrlimit(resource.RLIMIT_NOFILE, (soft, hard))

    def test_read_ignores_a_folder_without_done(self):
        d = SC.cache_dir(self.tmp, "abc")
        SC.write(d, self.arrays, {})
        (d / "done").unlink()
        self.assertIsNone(SC.read(d))
        self.assertIsNone(SC.read(self.tmp / "missing"))

    def test_failed_write_leaves_no_final_folder(self):
        d = SC.cache_dir(self.tmp, "abc")
        with mock.patch.object(SC, "save_npy", side_effect=OSError("disk full")):
            with self.assertRaises(OSError):
                SC.write(d, self.arrays, {})
        self.assertFalse(d.exists())
        self.assertEqual(list(self.tmp.iterdir()), [], "the temporary folder is removed too")

    def test_load_or_build_builds_once(self):
        d = SC.cache_dir(self.tmp, "abc")
        calls = []
        def build():
            calls.append(1)
            return self.arrays, {"n": 1}
        a1, _ = SC.load_or_build(d, build)
        a2, meta = SC.load_or_build(d, build)
        self.assertEqual(len(calls), 1)
        self.assertEqual(meta["n"], 1)
        self.assertIsInstance(a1["a"], np.memmap, "even the first load hands back the written maps")
        np.testing.assert_array_equal(a2["xyz"], self.arrays["xyz"])

    def test_unwritable_falls_back_to_ram(self):
        d = SC.cache_dir(self.tmp, "abc")
        logs = []
        with mock.patch.object(SC, "write", side_effect=OSError("read-only file system")):
            arrays, _ = SC.load_or_build(d, lambda: (self.arrays, {}), log=logs.append)
        self.assertIs(arrays["a"], self.arrays["a"])
        self.assertEqual(len(logs), 1)
        self.assertIn("read-only file system", logs[0])

    def test_prune_removes_only_unkept_serve_folders(self):
        keep = SC.cache_dir(self.tmp, "new")
        old = SC.cache_dir(self.tmp, "old")
        for d in (keep, old):
            SC.write(d, self.arrays, {})
        area = self.tmp / "0123abcd"
        area.mkdir()
        (area / "cells.npz").write_bytes(b"x")
        tiles = self.tmp / "v4-abc"
        tiles.mkdir()
        gone = SC.prune(self.tmp, [keep])
        self.assertEqual(gone, [old])
        self.assertTrue(keep.exists() and area.exists() and tiles.exists())

    def test_prune_keeps_young_temporary_folders_and_removes_old_ones(self):
        young = self.tmp / f"serve-x-v{SC.VERSION}.tmp-1"
        stale = self.tmp / f"serve-y-v{SC.VERSION}.tmp-2"
        young.mkdir()
        stale.mkdir()
        t = time.time() - 7200
        os.utime(stale, (t, t))
        SC.prune(self.tmp, [])
        self.assertTrue(young.exists(), "may belong to a server still writing it")
        self.assertFalse(stale.exists())

    def test_prune_removes_a_young_temporary_folder_of_a_process_that_is_gone(self):
        import subprocess
        pr = subprocess.Popen(["true"])
        pr.wait()
        orphan = self.tmp / f"serve-z-v{SC.VERSION}.tmp-{pr.pid}"
        orphan.mkdir()
        SC.prune(self.tmp, [])
        self.assertFalse(orphan.exists(), "its writer died: nobody finishes it")

    def test_prune_skips_a_folder_another_process_removed_meanwhile(self):
        gone = self.tmp / f"serve-q-v{SC.VERSION}.tmp-1"
        with mock.patch.object(Path, "iterdir", lambda _: iter([gone])), \
                mock.patch.object(Path, "is_dir", lambda _: True):
            self.assertEqual(SC.prune(self.tmp, []), [])

    def test_another_server_finishing_first_is_no_failure(self):
        import errno
        d = SC.cache_dir(self.tmp, "race")
        real = os.rename

        def other_first(src, dst):
            with mock.patch.object(SC.os, "rename", real):
                SC.write(Path(dst), self.arrays, {})
            raise OSError(errno.ENOTEMPTY, "Directory not empty")
        logs = []
        with mock.patch.object(SC.os, "rename", other_first):
            arrays, _ = SC.load_or_build(d, lambda: (self.arrays, {}), log=logs.append)
        self.assertEqual(logs, [], "no RAM fallback")
        self.assertIsInstance(arrays["a"], np.memmap)
        self.assertFalse([p for p in self.tmp.iterdir() if ".tmp-" in p.name])

    def test_a_damaged_cache_is_rebuilt(self):
        for damage in ("empty", "truncated"):
            d = SC.cache_dir(self.tmp, damage)
            SC.write(d, self.arrays, {})
            f = d / "xyz.npy"
            f.write_bytes(b"" if damage == "empty" else f.read_bytes()[:200])   # e.g. power lost after the rename
            logs = []
            arrays, _ = SC.load_or_build(d, lambda: (self.arrays, {}), log=logs.append)
            self.assertIsInstance(arrays["xyz"], np.memmap, damage)
            np.testing.assert_array_equal(arrays["xyz"], self.arrays["xyz"])
            self.assertIsNotNone(SC.read(d), f"{damage}: healed for the next start")


if __name__ == "__main__":
    unittest.main()


class SaveNpyTest(unittest.TestCase):
    def test_same_bytes_as_np_save(self):
        rng = np.random.default_rng(0)
        arrays = [rng.normal(size=3_000_000), rng.random((1000, 3)), np.zeros(0), np.array(3.5),
                  rng.integers(0, 9, 50_000).astype(np.int8), rng.random(10_000) < 0.5,
                  np.arange(12, dtype=np.uint64).reshape(3, 4)[:, ::2], np.asfortranarray(rng.random((300, 7)))]
        with tempfile.TemporaryDirectory() as t:
            for i, a in enumerate(arrays):
                with self.subTest(i=i, dtype=str(a.dtype), shape=a.shape):
                    np.save(Path(t) / "want.npy", np.ascontiguousarray(a), allow_pickle=False)
                    SC.save_npy(Path(t) / "got.npy", a)
                    self.assertEqual((Path(t) / "got.npy").read_bytes(), (Path(t) / "want.npy").read_bytes())

    def test_written_in_plain_writes_not_tofile(self):
        """numpy's tofile() path preallocates the file; btrfs then stores it uncompressed."""
        from numpy.lib import format as F
        seen = []
        real = F.write_array
        def spy(fp, *a, **k):
            seen.append(F.isfileobj(fp))
            return real(fp, *a, **k)
        with tempfile.TemporaryDirectory() as t, mock.patch.object(F, "write_array", spy):
            SC.save_npy(Path(t) / "a.npy", np.ones(1000))
            self.assertTrue(np.array_equal(np.load(Path(t) / "a.npy"), np.ones(1000)))
        self.assertEqual(seen, [False])