import json
import os
import shutil

import xxhash
import zstandard as zstd

MAGIC = 20260301
THRESHOLD = 4096  # 4KB


class Store:
    """Content-addressed, zstd-compressed blob store.
    Mirrors Go's lmo.Store — all operations are methods on this class.
    """

    def __init__(self, config_dir):
        self.config_dir = config_dir
        self.rel_path = ""
        self._compressor = zstd.ZstdCompressor()
        self._decompressor = zstd.ZstdDecompressor()

    def set_rel_path(self, rel_path):
        """Set the relative path and create the blobs directory."""
        self.rel_path = rel_path
        blob_dir = os.path.join(self.config_dir, "store", rel_path, "blobs")
        os.makedirs(blob_dir, exist_ok=True)

    # ── blob I/O ──

    def put_blob(self, data):
        """Store data as a zstd-compressed blob. Returns its XXH3-128 ref.
        Skips write if blob already exists (content-addressed dedup)."""
        ref = _hash_ref(data)
        p = self._blob_path_local(ref)
        if os.path.exists(p):
            return ref
        os.makedirs(os.path.dirname(p), exist_ok=True)
        compressed = self._compressor.compress(data)
        with open(p, "wb") as f:
            f.write(compressed)
        return ref

    def get_blob(self, ref, rel_path):
        """Read and decompress a blob identified by its ref and rel_path."""
        p = self._blob_path_for(ref, rel_path)
        with open(p, "rb") as f:
            compressed = f.read()
        return self._decompressor.decompress(compressed)

    def _blob_path_local(self, ref):
        """Compute filesystem path for writes using self.rel_path."""
        h = ref.removeprefix("xxh3:")
        return os.path.join(self.config_dir, "store", self.rel_path, "blobs", h[:2], h[2:])

    def _blob_path_for(self, ref, rel_path):
        """Compute filesystem path for reads using given rel_path."""
        h = ref.removeprefix("xxh3:")
        return os.path.join(self.config_dir, "store", rel_path, "blobs", h[:2], h[2:])

    # ── Pack ──

    def pack(self, payload):
        """Walk the JSON payload, extract fields >= Threshold as blobs,
        replace them with BlobRef markers. Returns modified bytes or
        the original if nothing was extracted."""
        if self.rel_path == "":
            return payload

        try:
            msg = json.loads(payload)
        except (json.JSONDecodeError, UnicodeDecodeError):
            return payload

        if not isinstance(msg, dict):
            return payload

        modified = False
        for key, value in list(msg.items()):
            new_val, changed = self._extract_field(value)
            if changed:
                msg[key] = new_val
                modified = True

        if not modified:
            return payload

        return json.dumps(msg, separators=(",", ":"), ensure_ascii=False).encode("utf-8")

    def _extract_field(self, value):
        """Process a single JSON value. Returns (new_value, changed).
        Objects: recurse into children. Large arrays/strings/numbers: extract.
        Existing BlobRefs: keep as-is."""
        if is_blob_ref(value):
            return value, False

        raw = json.dumps(value, separators=(",", ":"), ensure_ascii=False)
        raw_bytes = raw.encode("utf-8")

        # Object: recurse into children
        if isinstance(value, dict):
            return self._extract_object(value, raw_bytes)

        # Array or scalar: extract if large
        if len(raw_bytes) >= THRESHOLD:
            ref = self.put_blob(raw_bytes)
            return self._build_blob_ref(ref, raw_bytes, value), True

        return value, False

    def _extract_object(self, obj, raw_bytes):
        """Recurse into a JSON object, extracting large leaves."""
        if len(raw_bytes) < THRESHOLD:
            return obj, False

        modified = False
        out = {}
        for key, child in obj.items():
            new_val, changed = self._extract_field(child)
            if changed:
                out[key] = new_val
                modified = True
            else:
                out[key] = child

        if not modified:
            # Object is large but no individual children were large enough.
            # Extract the whole object as a single blob.
            ref = self.put_blob(raw_bytes)
            return self._build_blob_ref(ref, raw_bytes, obj), True

        return out, True

    # ── Resolve ──

    def resolve(self, data, key):
        """Lazy single-field resolution. Tries direct path first; if it hits
        a BlobRef, resolves it and continues with remaining path segments."""
        try:
            msg = json.loads(data)
        except (json.JSONDecodeError, UnicodeDecodeError):
            return None

        if not isinstance(msg, dict):
            return None

        parts = key.split(".")
        current = msg
        for part in parts:
            if isinstance(current, dict):
                if is_blob_ref(current):
                    current = self._resolve_ref(current)
                    if part in current:
                        current = current[part]
                    else:
                        return None
                elif part in current:
                    current = current[part]
                else:
                    return None
            elif isinstance(current, list):
                try:
                    current = current[int(part)]
                except (ValueError, IndexError):
                    return None
            else:
                return None

        if is_blob_ref(current):
            current = self._resolve_ref(current)

        return current

    def resolve_all(self, data):
        """Eagerly resolve every BlobRef in the payload. Returns modified
        bytes, or the original if no BlobRefs are found."""
        try:
            msg = json.loads(data)
        except (json.JSONDecodeError, UnicodeDecodeError):
            return data

        if not isinstance(msg, dict):
            return data

        modified = False
        for key, value in list(msg.items()):
            new_val, changed = self._resolve_value(value)
            if changed:
                msg[key] = new_val
                modified = True

        if not modified:
            return data

        return json.dumps(msg, separators=(",", ":"), ensure_ascii=False).encode("utf-8")

    def _resolve_value(self, value):
        """Resolve a single value. Returns (new_value, changed).
        BlobRefs -> decompressed data. Nested objects -> recurse."""
        if is_blob_ref(value):
            return self._resolve_ref(value), True

        if isinstance(value, dict):
            modified = False
            out = {}
            for k, v in value.items():
                new_val, changed = self._resolve_value(v)
                if changed:
                    out[k] = new_val
                    modified = True
                else:
                    out[k] = v
            if modified:
                return out, True
            return value, False

        return value, False

    def _resolve_ref(self, blob_ref):
        """Fetch and decompress a BlobRef, return parsed JSON value."""
        ref = blob_ref["__ref"]
        rel_path = blob_ref.get("__path", self.rel_path)
        data = self.get_blob(ref, rel_path)
        return json.loads(data)

    # ── BlobRef helpers ──

    def _build_blob_ref(self, ref, raw_bytes, value):
        """Build a BlobRef marker dict."""
        br = {
            "__ref": ref,
            "__magic": MAGIC,
            "__size": len(raw_bytes),
            "__path": self.rel_path,
            "__type": _json_type(value),
        }
        t = br["__type"]
        if t == "array":
            br["__len"] = len(value)
        elif t == "string":
            br["__len"] = len(value)
        return br

    # ── lifecycle ──

    def purge(self):
        """Remove the store directory for the current rel_path."""
        path = os.path.join(self.config_dir, "store", self.rel_path)
        shutil.rmtree(path, ignore_errors=True)

    def close(self):
        """Release resources."""
        pass


# ── module-level helpers ──

def is_blob_ref(value):
    """Check if value is a BlobRef marker dict."""
    return (
        isinstance(value, dict)
        and value.get("__magic") == MAGIC
        and bool(value.get("__ref"))
    )


def _hash_ref(data):
    """Compute XXH3-128 hash -> 'xxh3:{32hex}'. Canonical big-endian,
    matches Go's zeebo/xxh3."""
    digest = xxhash.xxh3_128_digest(data)
    return "xxh3:" + digest.hex()


def _json_type(value):
    """Map Python type to JSON type string (matches Go buildBlobRef)."""
    if isinstance(value, list):
        return "array"
    if isinstance(value, dict):
        return "object"
    if isinstance(value, str):
        return "string"
    if isinstance(value, bool):
        return "boolean"
    if isinstance(value, (int, float)):
        return "number"
    return ""


# ── module-level singleton (mirrors Go lmo_rt.go) ──

_store = None


def init_lmo():
    """Initialize the LMO store using lmo_store_path from GetRobotInfo().
    The deskbot owns the path convention — the plugin treats it as opaque."""
    global _store
    if _store is not None:
        _store.close()
        _store = None

    from robomotion.runtime import Runtime
    from robomotion.utils import get_config_dir

    robot_info = Runtime.get_robot_info()
    rel_path = str(robot_info.get("lmo_store_path", ""))

    _store = Store(get_config_dir())
    if rel_path:
        _store.set_rel_path(rel_path)


def lmo_active():
    """Return True if the LMO store is initialized."""
    return _store is not None


def pack_value(value):
    """Pack a value using recursive extraction (matches Go PackValue).
    Returns (packed_value, True) if packing occurred, (None, False) otherwise."""
    if _store is None or _store.rel_path == "":
        return None, False

    raw = json.dumps(value, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
    if len(raw) < THRESHOLD:
        return None, False

    result = _store.pack(raw)
    if result is raw:
        return None, False

    packed = json.loads(result)
    return packed, True


def resolve_blob_ref_value(blob_ref_map):
    """Resolve a BlobRef value, learning rel_path from the BlobRef if needed
    (matches Go ResolveBlobRefValue)."""
    if _store is None:
        return None

    ref = blob_ref_map.get("__ref", "")
    path = blob_ref_map.get("__path", "")

    # Learn rel_path from BlobRef if store doesn't have one yet
    if _store.rel_path == "" and path:
        _store.set_rel_path(path)

    rel_path = path if path else _store.rel_path
    data = _store.get_blob(ref, rel_path)
    return json.loads(data)


def lmo_pack(payload):
    """Pack payload bytes, extracting large fields as blobs. No-op if store is None."""
    if _store is None:
        return payload
    return _store.pack(payload)


def lmo_resolve(data, key):
    """Lazily resolve a single field. No-op if store is None."""
    if _store is None:
        return None
    return _store.resolve(data, key)


def lmo_resolve_all(data):
    """Eagerly resolve all BlobRefs. No-op if store is None."""
    if _store is None:
        return data
    return _store.resolve_all(data)


def close_lmo():
    """Close the store."""
    global _store
    if _store is not None:
        _store.close()
        _store = None
