"""Tests for robomotion.lmo — Store, helpers, and module-level functions."""
import json
import os
import shutil
import tempfile

import pytest

from robomotion.lmo import (
    MAGIC,
    THRESHOLD,
    Store,
    _hash_ref,
    _json_type,
    is_blob_ref,
)


# ── fixtures ──


@pytest.fixture
def tmp_dir():
    """Temporary directory cleaned up after test."""
    d = tempfile.mkdtemp()
    yield d
    shutil.rmtree(d, ignore_errors=True)


@pytest.fixture
def store(tmp_dir):
    """Store with rel_path set, ready for read/write."""
    s = Store(tmp_dir)
    s.set_rel_path("test/path")
    return s


def _large_string(n=THRESHOLD):
    """Generate a string whose JSON serialization is >= n bytes."""
    return "x" * n


def _large_payload(key="big", n=THRESHOLD):
    """Return a JSON-encoded dict with one field >= THRESHOLD."""
    val = _large_string(n)
    return json.dumps({key: val}, separators=(",", ":"), ensure_ascii=False).encode("utf-8")


# ── is_blob_ref ──


class TestIsBlobRef:
    def test_valid_blob_ref(self):
        ref = {"__magic": MAGIC, "__ref": "xxh3:abc123"}
        assert is_blob_ref(ref) is True

    def test_missing_magic(self):
        assert is_blob_ref({"__ref": "xxh3:abc123"}) is False

    def test_wrong_magic(self):
        assert is_blob_ref({"__magic": 0, "__ref": "xxh3:abc123"}) is False

    def test_missing_ref(self):
        assert is_blob_ref({"__magic": MAGIC}) is False

    def test_empty_ref(self):
        assert is_blob_ref({"__magic": MAGIC, "__ref": ""}) is False

    def test_not_a_dict(self):
        assert is_blob_ref("not a dict") is False
        assert is_blob_ref(42) is False
        assert is_blob_ref(None) is False
        assert is_blob_ref([]) is False


# ── _hash_ref ──


class TestHashRef:
    def test_prefix(self):
        ref = _hash_ref(b"hello")
        assert ref.startswith("xxh3:")

    def test_hex_length(self):
        ref = _hash_ref(b"hello")
        hex_part = ref.removeprefix("xxh3:")
        assert len(hex_part) == 32  # 128 bits = 32 hex chars

    def test_deterministic(self):
        assert _hash_ref(b"test data") == _hash_ref(b"test data")

    def test_different_data_different_hash(self):
        assert _hash_ref(b"aaa") != _hash_ref(b"bbb")


# ── _json_type ──


class TestJsonType:
    def test_array(self):
        assert _json_type([1, 2]) == "array"

    def test_object(self):
        assert _json_type({"a": 1}) == "object"

    def test_string(self):
        assert _json_type("hello") == "string"

    def test_boolean(self):
        assert _json_type(True) == "boolean"
        assert _json_type(False) == "boolean"

    def test_number_int(self):
        assert _json_type(42) == "number"

    def test_number_float(self):
        assert _json_type(3.14) == "number"

    def test_none(self):
        assert _json_type(None) == ""


# ── Store.__init__ / set_rel_path ──


class TestStoreInit:
    def test_init_no_dirs_created(self, tmp_dir):
        Store(tmp_dir)
        # No store/blobs directory should exist yet
        assert not os.path.exists(os.path.join(tmp_dir, "store"))

    def test_set_rel_path_creates_blobs_dir(self, tmp_dir):
        s = Store(tmp_dir)
        s.set_rel_path("abc/def")
        expected = os.path.join(tmp_dir, "store", "abc/def", "blobs")
        assert os.path.isdir(expected)

    def test_rel_path_initially_empty(self, tmp_dir):
        s = Store(tmp_dir)
        assert s.rel_path == ""


# ── put_blob / get_blob ──


class TestBlobIO:
    def test_put_and_get_roundtrip(self, store):
        data = b'{"key":"value"}'
        ref = store.put_blob(data)
        assert ref.startswith("xxh3:")
        got = store.get_blob(ref, store.rel_path)
        assert got == data

    def test_put_blob_dedup(self, store):
        data = b"duplicate data"
        ref1 = store.put_blob(data)
        ref2 = store.put_blob(data)
        assert ref1 == ref2

    def test_blob_is_compressed(self, store):
        data = b"a" * 10000
        ref = store.put_blob(data)
        blob_file = store._blob_path_local(ref)
        raw_size = os.path.getsize(blob_file)
        # zstd should compress repetitive data significantly
        assert raw_size < len(data)

    def test_get_blob_file_not_found(self, store):
        with pytest.raises(FileNotFoundError):
            store.get_blob("xxh3:0000000000000000000000000000dead", store.rel_path)

    def test_get_blob_with_different_rel_path(self, tmp_dir):
        """Write blob via one rel_path, read via that rel_path from another store."""
        s1 = Store(tmp_dir)
        s1.set_rel_path("path/a")
        data = b"cross-path data"
        ref = s1.put_blob(data)

        s2 = Store(tmp_dir)
        s2.set_rel_path("path/b")
        # Read from s1's path
        got = s2.get_blob(ref, "path/a")
        assert got == data


# ── blob_path helpers ──


class TestBlobPath:
    def test_blob_path_local_uses_rel_path(self, store):
        ref = "xxh3:aabbccdd11223344aabbccdd11223344"
        p = store._blob_path_local(ref)
        assert "store/test/path/blobs/aa/" in p

    def test_blob_path_for_uses_given_path(self, store):
        ref = "xxh3:aabbccdd11223344aabbccdd11223344"
        p = store._blob_path_for(ref, "other/path")
        assert "store/other/path/blobs/aa/" in p

    def test_blob_path_prefix_split(self, store):
        """First 2 hex chars are directory, rest is filename."""
        ref = "xxh3:aabbccdd11223344aabbccdd11223344"
        p = store._blob_path_local(ref)
        assert p.endswith("aa/bbccdd11223344aabbccdd11223344")


# ── pack ──


class TestPack:
    def test_pack_no_op_when_rel_path_empty(self, tmp_dir):
        s = Store(tmp_dir)
        payload = _large_payload()
        result = s.pack(payload)
        assert result is payload  # identity check — unchanged

    def test_pack_no_op_below_threshold(self, store):
        payload = json.dumps({"small": "val"}).encode("utf-8")
        result = store.pack(payload)
        assert result is payload

    def test_pack_extracts_large_string_field(self, store):
        big = _large_string()
        payload = json.dumps({"big": big, "id": 1}, separators=(",", ":")).encode("utf-8")
        result = store.pack(payload)
        msg = json.loads(result)
        assert is_blob_ref(msg["big"])
        assert msg["id"] == 1  # small field untouched
        # BlobRef metadata
        assert msg["big"]["__type"] == "string"
        assert msg["big"]["__path"] == store.rel_path

    def test_pack_extracts_large_array_field(self, store):
        big_array = list(range(1500))
        payload = json.dumps({"arr": big_array}, separators=(",", ":")).encode("utf-8")
        assert len(payload) >= THRESHOLD
        result = store.pack(payload)
        msg = json.loads(result)
        assert is_blob_ref(msg["arr"])
        assert msg["arr"]["__type"] == "array"
        assert msg["arr"]["__len"] == 1500

    def test_pack_leaves_existing_blobref_intact(self, store):
        existing_ref = {"__magic": MAGIC, "__ref": "xxh3:existing", "__size": 5000, "__path": "old", "__type": "string"}
        payload = json.dumps({"ref": existing_ref}, separators=(",", ":")).encode("utf-8")
        result = store.pack(payload)
        # Nothing to extract, should be unchanged
        assert result is payload

    def test_pack_invalid_json(self, store):
        payload = b"not json at all"
        assert store.pack(payload) is payload

    def test_pack_non_dict_json(self, store):
        payload = json.dumps([1, 2, 3]).encode("utf-8")
        assert store.pack(payload) is payload

    def test_pack_recursive_nested_object(self, store):
        """Recursive pack extracts the large leaf inside a nested object."""
        big_string = _large_string()
        big_child = {"data": big_string}
        small_child = {"data": "small"}
        payload = json.dumps(
            {"parent": {"big": big_child, "small": small_child}},
            separators=(",", ":"),
        ).encode("utf-8")
        result = store.pack(payload)
        msg = json.loads(result)
        parent = msg["parent"]
        # The large string inside big_child is extracted, not big_child itself
        assert is_blob_ref(parent["big"]["data"])
        assert parent["big"]["data"]["__type"] == "string"
        # The small child should remain inline
        assert parent["small"] == small_child

    def test_pack_large_object_no_large_children(self, store):
        """A large object where no single child is >= THRESHOLD is extracted as a whole blob."""
        # Many small fields that together exceed threshold
        obj = {f"k{i}": f"val{'x' * 100}" for i in range(50)}
        payload = json.dumps({"bulk": obj}, separators=(",", ":")).encode("utf-8")
        assert len(payload) >= THRESHOLD
        result = store.pack(payload)
        msg = json.loads(result)
        assert is_blob_ref(msg["bulk"])
        assert msg["bulk"]["__type"] == "object"


# ── resolve ──


class TestResolve:
    def test_resolve_direct_field(self, store):
        payload = json.dumps({"name": "Alice"}).encode("utf-8")
        assert store.resolve(payload, "name") == "Alice"

    def test_resolve_blob_ref_field(self, store):
        big = _large_string()
        original = json.dumps({"big": big, "id": 1}, separators=(",", ":")).encode("utf-8")
        packed = store.pack(original)
        # Resolve the extracted field
        result = store.resolve(packed, "big")
        assert result == big

    def test_resolve_nested_path(self, store):
        payload = json.dumps({"a": {"b": {"c": 42}}}).encode("utf-8")
        assert store.resolve(payload, "a.b.c") == 42

    def test_resolve_missing_key(self, store):
        payload = json.dumps({"a": 1}).encode("utf-8")
        assert store.resolve(payload, "nonexistent") is None

    def test_resolve_array_index(self, store):
        payload = json.dumps({"items": ["zero", "one", "two"]}).encode("utf-8")
        assert store.resolve(payload, "items.1") == "one"

    def test_resolve_invalid_json(self, store):
        assert store.resolve(b"not json", "key") is None

    def test_resolve_through_blobref(self, store):
        """BlobRef at intermediate path level — resolve should walk through it."""
        inner = {"x": 99, "y": 100}
        inner_json = json.dumps(inner, separators=(",", ":")).encode("utf-8")
        # Manually store the inner object as a blob
        ref = store.put_blob(inner_json)
        blob_ref = store._build_blob_ref(ref, inner_json, inner)
        packed = json.dumps({"nested": blob_ref}).encode("utf-8")
        assert store.resolve(packed, "nested.x") == 99


# ── resolve_all ──


class TestResolveAll:
    def test_resolve_all_roundtrip(self, store):
        big = _large_string()
        original = {"big": big, "id": 1}
        payload = json.dumps(original, separators=(",", ":")).encode("utf-8")
        packed = store.pack(payload)
        assert packed != payload  # something was extracted
        resolved = store.resolve_all(packed)
        assert json.loads(resolved) == original

    def test_resolve_all_no_blobrefs(self, store):
        payload = json.dumps({"a": 1}).encode("utf-8")
        assert store.resolve_all(payload) is payload

    def test_resolve_all_nested_blobrefs(self, store):
        big1 = _large_string()
        big2 = _large_string(THRESHOLD + 500)
        original = {"f1": big1, "f2": big2, "small": "ok"}
        payload = json.dumps(original, separators=(",", ":")).encode("utf-8")
        packed = store.pack(payload)
        resolved = store.resolve_all(packed)
        assert json.loads(resolved) == original

    def test_resolve_all_invalid_json(self, store):
        data = b"not json"
        assert store.resolve_all(data) is data


# ── _resolve_ref uses __path ──


class TestResolveRefPath:
    def test_resolve_ref_uses_blobref_path(self, tmp_dir):
        """_resolve_ref should read from BlobRef __path, not store's own rel_path."""
        s1 = Store(tmp_dir)
        s1.set_rel_path("writer/path")
        data = json.dumps("hello from writer").encode("utf-8")
        ref = s1.put_blob(data)
        blob_ref = s1._build_blob_ref(ref, data, "hello from writer")
        assert blob_ref["__path"] == "writer/path"

        # A different store with a different rel_path should still resolve via __path
        s2 = Store(tmp_dir)
        s2.set_rel_path("reader/path")
        resolved = s2._resolve_ref(blob_ref)
        assert resolved == "hello from writer"

    def test_resolve_ref_falls_back_to_store_path(self, store):
        """If BlobRef has no __path, use the store's rel_path."""
        data = json.dumps({"k": "v"}).encode("utf-8")
        ref = store.put_blob(data)
        blob_ref = {"__ref": ref, "__magic": MAGIC}  # no __path
        resolved = store._resolve_ref(blob_ref)
        assert resolved == {"k": "v"}


# ── _build_blob_ref ──


class TestBuildBlobRef:
    def test_string_blob_ref(self, store):
        value = "hello world"
        raw = json.dumps(value).encode("utf-8")
        ref = _hash_ref(raw)
        br = store._build_blob_ref(ref, raw, value)
        assert br["__ref"] == ref
        assert br["__magic"] == MAGIC
        assert br["__size"] == len(raw)
        assert br["__path"] == store.rel_path
        assert br["__type"] == "string"
        assert br["__len"] == len(value)

    def test_array_blob_ref(self, store):
        value = [1, 2, 3]
        raw = json.dumps(value).encode("utf-8")
        ref = _hash_ref(raw)
        br = store._build_blob_ref(ref, raw, value)
        assert br["__type"] == "array"
        assert br["__len"] == 3

    def test_object_blob_ref_no_len(self, store):
        value = {"a": 1}
        raw = json.dumps(value).encode("utf-8")
        ref = _hash_ref(raw)
        br = store._build_blob_ref(ref, raw, value)
        assert br["__type"] == "object"
        assert "__len" not in br

    def test_number_blob_ref(self, store):
        value = 42
        raw = json.dumps(value).encode("utf-8")
        ref = _hash_ref(raw)
        br = store._build_blob_ref(ref, raw, value)
        assert br["__type"] == "number"
        assert "__len" not in br


# ── purge ──


class TestPurge:
    def test_purge_removes_store_dir(self, tmp_dir):
        s = Store(tmp_dir)
        s.set_rel_path("to/purge")
        s.put_blob(b"some data")
        store_path = os.path.join(tmp_dir, "store", "to/purge")
        assert os.path.isdir(store_path)
        s.purge()
        assert not os.path.exists(store_path)

    def test_purge_idempotent(self, tmp_dir):
        s = Store(tmp_dir)
        s.set_rel_path("gone")
        s.purge()
        s.purge()  # should not raise


# ── pack + resolve full cycle ──


class TestPackResolveCycle:
    def test_full_cycle_mixed_payload(self, store):
        """Pack a message with mixed small/large fields, then resolve_all to get original back."""
        original = {
            "id": 42,
            "name": "test",
            "big_data": _large_string(THRESHOLD + 1000),
            "nested": {
                "small": "inline",
                "big_child": _large_string(THRESHOLD + 500),
            },
        }
        payload = json.dumps(original, separators=(",", ":")).encode("utf-8")
        packed = store.pack(payload)
        msg = json.loads(packed)

        # Small fields should be inline
        assert msg["id"] == 42
        assert msg["name"] == "test"
        # Large fields should be BlobRefs
        assert is_blob_ref(msg["big_data"])
        assert is_blob_ref(msg["nested"]["big_child"])
        assert msg["nested"]["small"] == "inline"

        # Full round-trip
        resolved = store.resolve_all(packed)
        assert json.loads(resolved) == original

    def test_pack_then_lazy_resolve_each_field(self, store):
        """Pack, then lazily resolve individual fields."""
        original = {
            "f1": _large_string(),
            "f2": "small",
            "f3": list(range(1000)),
        }
        payload = json.dumps(original, separators=(",", ":")).encode("utf-8")
        packed = store.pack(payload)

        assert store.resolve(packed, "f1") == original["f1"]
        assert store.resolve(packed, "f2") == "small"
        assert store.resolve(packed, "f3") == original["f3"]
