#######################################################################
# Copyright (c) 2019-present, Blosc Development Team <blosc@blosc.org>
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause
#######################################################################

from __future__ import annotations

import copy
import os
from contextlib import contextmanager, nullcontext
from typing import TYPE_CHECKING, Any

import numpy as np

import blosc2
from blosc2.c2array import C2Array
from blosc2.msgpack_utils import msgpack_unpackb

if TYPE_CHECKING:
    from collections.abc import Iterator, KeysView

    from blosc2.schunk import SChunk

PROFILE = False  # Set to True to enable PROFILE prints in EmbedStore


class EmbedStore:
    """
    A dictionary-like container for storing NumPy/Blosc2 arrays (NDArray or SChunk) as nodes.

    For NumPy arrays, Blosc2 NDArrays (even if they live in external ``.b2nd`` files),
    and Blosc2 SChunk objects, the data is read and embedded into the store. For remote
    arrays (``C2Array``), only lightweight references (URL base and path) are stored.
    If you need a richer hierarchical container with optional external references, consider using
    `blosc2.TreeStore` or `blosc2.DictStore`.

    Parameters
    ----------
    urlpath : str or None, optional
        Path for persistent storage. Using a '.b2e' extension is recommended.
        If None, the embed store will be in memory only, which can be
        deserialized later using the :func:`blosc2.from_cframe` function.
    mode : str, optional
        File mode ('r', 'w', 'a'). Default is 'w'.
    mmap_mode : str or None, optional
        Memory mapping mode for read access. For now, only ``"r"`` is supported,
        and only when ``mode="r"``. Default is None.
    cparams : dict or None, optional
        Compression parameters for nodes and the embed store itself.
        Default is None, which uses the default Blosc2 parameters.
    dparams : dict or None, optional
        Decompression parameters for nodes and the embed store itself.
        Default is None, which uses the default Blosc2 parameters.
    storage : blosc2.Storage or None, optional
        Storage properties for the embed store.  If passed, it will override
        the `urlpath` and `mode` parameters.
    chunksize : int, optional
        Size of chunks for the backing storage. Default is 1 MiB.

    Notes
    -----
    Without file locking, EmbedStore is single-process, single-writer: the key
    map is cached in Python, so mutations made through another handle are not
    seen until the store is reopened, and concurrent writers can corrupt each
    other's entries.

    With file locking enabled on *every* handle (``locking=True`` in
    :class:`blosc2.Storage`, or the ``BLOSC_LOCKING`` environment variable),
    an on-disk EmbedStore can be shared across processes: mutations are
    atomic (the whole data-write + map-update runs under an exclusive lock)
    and every access re-syncs the key map, so readers follow keys added or
    removed by other processes.  A crash between the data write and the map
    flush can leave unreachable bytes in the container (harmless; reclaimed
    by rewriting the store).  Not supported together with ``mmap_mode``, nor
    on network filesystems (NFS).

    Examples
    --------
    >>> estore = EmbedStore(urlpath="example_estore.b2e", mode="w")
    >>> estore["/node1"] = np.array([1, 2, 3])
    >>> estore["/node2"] = blosc2.ones(2)
    >>> estore["/node3"] = blosc2.arange(3, dtype="i4", urlpath="external_node3.b2nd", mode="w")
    >>> urlpath = blosc2.URLPath("@public/examples/ds-1d.b2nd", "https://cat2.cloud/demo")
    >>> estore["/node4"] = blosc2.open(urlpath, mode="r")
    >>> print(list(estore.keys()))
    ['/node1', '/node2', '/node3', '/node4']
    >>> print(estore["/node1"][:])
    [1 2 3]

    """

    def __init__(
        self,
        urlpath: str | None = None,
        mode: str = "a",
        cparams: blosc2.CParams | None = None,
        dparams: blosc2.CParams | None = None,
        storage: blosc2.Storage | None = None,
        chunksize: int | None = 2**13,
        _from_schunk: SChunk | None = None,
        *,
        mmap_mode: str | None = None,
        meta: dict | None = None,
    ):
        """Initialize EmbedStore."""

        # For some reason, the SChunk store cannot achieve the same compression ratio as the NDArray store,
        # although it is more efficient in terms of CPU usage.
        # Let's use the SChunk store by default and continue experimenting.
        self._schunk_store = True  # put this to False to use an NDArray instead of a SChunk
        self.urlpath = urlpath
        if mmap_mode not in (None, "r"):
            raise ValueError("For EmbedStore containers, mmap_mode must be None or 'r'")
        if mmap_mode == "r" and mode != "r":
            raise ValueError("For EmbedStore containers, mmap_mode='r' requires mode='r'")
        self.mmap_mode = mmap_mode

        if _from_schunk is not None:
            self.urlpath = _from_schunk.urlpath
            self.cparams = _from_schunk.cparams
            self.dparams = _from_schunk.dparams
            self.mode = _from_schunk.mode
            self.mmap_mode = getattr(_from_schunk, "mmap_mode", None)
            self._store = _from_schunk
            self.storage = blosc2.Storage(
                contiguous=_from_schunk.contiguous,
                urlpath=_from_schunk.urlpath,
                mode=self.mode,
                mmap_mode=self.mmap_mode,
                initial_mapping_size=getattr(_from_schunk, "initial_mapping_size", None),
            )
            self.storage.meta = _from_schunk.meta
            self._set_shared(getattr(_from_schunk, "locking", False))
            self._load_metadata()
            return

        self.mode = mode
        self.cparams = cparams or blosc2.CParams()
        # self.cparams.nthreads = 1  # for debugging purposes, use only one thread
        self.dparams = dparams or blosc2.DParams()
        # self.dparams.nthreads = 1  # for debugging purposes, use only one thread
        if storage is None:
            self.storage = blosc2.Storage(
                contiguous=True,
                urlpath=urlpath,
                mode=mode,
            )
        else:
            self.storage = storage

        if mode in ("r", "a") and urlpath:
            self._store = blosc2.blosc2_ext.open(
                urlpath, mode=mode, offset=0, mmap_mode=mmap_mode, locking=self.storage.locking
            )
            self.storage.meta = self._store.meta
            self._set_shared(self.storage.locking)
            self._load_metadata()
            return

        _cparams = copy.deepcopy(self.cparams)
        _cparams.typesize = 1  # ensure typesize is set to 1 for byte storage
        _storage = self.storage
        _storage.meta = meta if meta is not None else {"b2embed": {"version": 1}}
        if self._schunk_store:
            self._store = blosc2.SChunk(
                chunksize=chunksize,
                data=None,
                cparams=_cparams,
                dparams=self.dparams,
                storage=_storage,
            )
        else:
            self._store = blosc2.zeros(
                chunksize,
                dtype=np.uint8,
                cparams=_cparams,
                dparams=self.dparams,
                storage=_storage,
            )
        self._embed_map: dict = {}
        self._current_offset = 0
        self._set_shared(self.storage.locking)
        # Flush the (empty) map right away so that another handle opened on
        # this store always finds the metadata to re-sync from
        self._save_metadata()

    @property
    def _backing_schunk(self) -> SChunk:
        """The SChunk under the store (the NDArray backing adds one indirection)."""
        return self._store if self._schunk_store else self._store.schunk

    def _set_shared(self, locking: bool) -> None:
        """Decide whether this handle may share the container with other processes.

        True when it is on-disk and file locking is enabled — explicitly or
        through the BLOSC_LOCKING environment variable (which the C library
        honors for any on-disk container opened with the default I/O).
        """
        env = os.environ.get("BLOSC_LOCKING", "")
        env_locking = env not in ("", "0") and self.mmap_mode is None
        self._shared = (bool(locking) or env_locking) and self.urlpath is not None

    def _sync_metadata(self) -> None:
        """Reload the key map if another handle changed the store.

        The raw vlmeta read polls the on-disk frame and re-syncs the C-level
        cached state (including ``change_tick``); the msgpack decode of the map
        is skipped when nothing changed.  A no-op for non-shared stores.
        """
        if not self._shared:
            return
        sc = self._backing_schunk
        try:
            raw = sc.vlmeta.get_vlmeta("estore_metadata")
        except KeyError:
            # Nothing flushed yet (freshly created store)
            return
        tick = sc.change_tick
        if tick == self._meta_tick:
            return
        metadata = msgpack_unpackb(raw)
        self._embed_map = metadata["embed_map"]
        self._current_offset = metadata["current_offset"]
        self._meta_tick = tick

    @contextmanager
    def _write_bracket(self):
        """Make a whole mutation atomic against other processes.

        Holds the exclusive frame lock (a no-op without locking) and re-syncs
        the key map before the mutation, so concurrent writers cannot start
        from the same offset or clobber each other's map entries.
        """
        with self._backing_schunk.holding_lock():
            self._sync_metadata()
            yield

    def _validate_key(self, key: str) -> None:
        """Validate node key."""
        if not isinstance(key, str):
            raise TypeError("Key must be a string.")
        if not key.startswith("/"):
            raise ValueError("Key must start with '/'.")
        if len(key) > 1 and key.endswith("/"):
            raise ValueError("Key cannot end with '/' unless it is the root key '/'.")
        if "//" in key:
            raise ValueError("Key cannot contain consecutive slashes '//'.")
        for char in (":", "\0", "\n", "\r", "\t"):
            if char in key:
                raise ValueError(f"Key cannot contain character: {char!r}")
        if key in self._embed_map:
            raise ValueError(f"Key '{key}' already exists in store.")

    def _ensure_capacity(self, needed_bytes: int) -> None:
        """Ensure backing storage has enough capacity."""
        required_size = self._current_offset + needed_bytes
        if required_size > self._store.shape[0]:
            new_size = max(required_size, int(self._store.shape[0] * 1.5))
            self._store.resize((new_size,))

    def __setitem__(
        self, key: str, value: blosc2.Array | SChunk | blosc2.ObjectArray | blosc2.BatchArray
    ) -> None:
        """Add a node to the embed store."""
        if self.mode == "r":
            raise ValueError("Cannot set items in read-only mode.")
        with self._write_bracket():
            self._validate_key(key)
            if isinstance(value, C2Array):
                self._embed_map[key] = {"urlbase": value.urlbase, "path": value.path}
            else:
                if isinstance(value, np.ndarray):
                    value = blosc2.asarray(value, cparams=self.cparams, dparams=self.dparams)
                serialized_data = value.to_cframe()
                data_len = len(serialized_data)
                if not self._schunk_store:
                    self._ensure_capacity(data_len)
                offset = self._current_offset
                if self._schunk_store:
                    self._store[offset : offset + data_len] = serialized_data
                else:
                    self._store[offset : offset + data_len] = np.frombuffer(serialized_data, dtype=np.uint8)
                self._current_offset += data_len
                self._embed_map[key] = {"offset": offset, "length": data_len}
            self._save_metadata()

    def __getitem__(self, key: str) -> blosc2.NDArray | SChunk | blosc2.ObjectArray | blosc2.BatchArray:
        """Retrieve a node from the embed store."""
        # Only a shared store needs the lock: it is the only mode where
        # _sync_metadata() does real work and where the index/data race exists.
        # holding_lock() also refreshes, which costs a filesystem poll, so the
        # non-shared path takes neither.
        with self._backing_schunk.holding_lock() if self._shared else nullcontext():
            self._sync_metadata()
            if key not in self._embed_map:
                raise KeyError(f"Key '{key}' not found in the embed store.")
            node_info = self._embed_map[key]
            urlbase = node_info.get("urlbase", None)
            if not urlbase:
                offset = node_info["offset"]
                length = node_info["length"]
                serialized_data = bytes(self._store[offset : offset + length])

        if urlbase:
            # Outside the lock: opening a C2Array involves an HTTP round trip
            return blosc2.open(blosc2.URLPath(node_info["path"], urlbase=urlbase), mode="r")
        # It is safer to copy data here, as the reference to the SChunk may disappear
        # Use from_cframe so we can deserialize either an NDArray or an SChunk
        return blosc2.from_cframe(serialized_data, copy=True)

    def get(
        self, key: str, default: Any = None
    ) -> blosc2.NDArray | SChunk | blosc2.ObjectArray | blosc2.BatchArray | Any:
        """Retrieve a node, or default if not found."""
        try:
            return self[key]
        except KeyError:
            return default

    def __delitem__(self, key: str) -> None:
        """Remove a node from the embed store."""
        if self.mode == "r":
            raise ValueError("Cannot delete items in read-only mode.")
        with self._write_bracket():
            if key not in self._embed_map:
                raise KeyError(f"Key '{key}' not found in the embed store.")
            del self._embed_map[key]
            self._save_metadata()

    def __contains__(self, key: str) -> bool:
        """Check if a key exists."""
        self._sync_metadata()
        return key in self._embed_map

    def __len__(self) -> int:
        """Return number of nodes."""
        self._sync_metadata()
        return len(self._embed_map)

    def __iter__(self) -> Iterator[str]:
        """Iterate over keys."""
        self._sync_metadata()
        return iter(self._embed_map)

    def keys(self) -> KeysView[str]:
        """Return all keys."""
        self._sync_metadata()
        return self._embed_map.keys()

    def values(self) -> Iterator[blosc2.NDArray | SChunk | blosc2.ObjectArray | blosc2.BatchArray]:
        """Iterate over all values."""
        self._sync_metadata()
        for key in self._embed_map:
            yield self[key]

    def items(
        self,
    ) -> Iterator[tuple[str, blosc2.NDArray | SChunk | blosc2.ObjectArray | blosc2.BatchArray]]:
        """Iterate over (key, value) pairs."""
        self._sync_metadata()
        for key in self._embed_map:
            yield key, self[key]

    def _save_metadata(self) -> None:
        """Save embed store map to vlmeta."""
        metadata = {"embed_map": self._embed_map, "current_offset": self._current_offset}
        self._store.vlmeta["estore_metadata"] = metadata
        self._meta_tick = self._backing_schunk.change_tick

    def _load_metadata(self) -> None:
        """Load embed store map from vlmeta."""
        if "estore_metadata" in self._store.vlmeta:
            metadata = self._store.vlmeta["estore_metadata"]
            self._embed_map = metadata["embed_map"]
            self._current_offset = metadata["current_offset"]
        else:
            self._embed_map = {}
            self._current_offset = 0
        self._meta_tick = self._backing_schunk.change_tick

    def to_cframe(self) -> bytes:
        """Serialize embed store to CFrame format."""
        return self._store.to_cframe()

    def __enter__(self):
        """Context manager enter."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit."""
        # No need to close anything as SChunk/NDArray handles persistence automatically
        return False


def estore_from_cframe(cframe: bytes, copy: bool = False) -> EmbedStore:
    """
    Deserialize a CFrame to an EmbedStore object.

    Parameters
    ----------
    cframe : bytes
        CFrame data to deserialize.
    copy : bool, optional
        If True, copy the data. Default is False.

    Returns
    -------
    estore : EmbedStore
        The deserialized EmbedStore object.
    """
    schunk = blosc2.schunk_from_cframe(cframe, copy=copy)
    return EmbedStore(_from_schunk=schunk)


if __name__ == "__main__":
    # Example usage
    persistent = False
    if persistent:
        estore = EmbedStore(urlpath="example_estore.b2e", mode="w")  # , cparams=blosc2.CParams(clevel=0))
    else:
        estore = EmbedStore()  # , cparams=blosc2.CParams(clevel=0))
    # import pdb;  pdb.set_trace()
    estore["/node1"] = np.array([1, 2, 3])
    estore["/node2"] = blosc2.ones(2)
    urlpath = blosc2.URLPath("@public/examples/ds-1d.b2nd", "https://cat2.cloud/demo")
    arr_remote = blosc2.open(urlpath, mode="r")
    estore["/dir1/node3"] = arr_remote

    print("EmbedStore keys:", list(estore.keys()))
    print("Node1 data:", estore["/node1"][:])
    print("Node2 data:", estore["/node2"][:])
    print("Node3 data (remote):", estore["/dir1/node3"][:3])

    del estore["/node1"]
    print("After deletion, keys:", list(estore.keys()))

    # Reading back the estore
    if persistent:
        estore_read = EmbedStore(urlpath="example_estore.b2e", mode="r")
    else:
        estore_read = blosc2.from_cframe(estore.to_cframe())

    print("Read keys:", list(estore_read.keys()))
    for key, value in estore_read.items():
        print(
            f"shape of {key}: {value.shape}, dtype: {value.dtype}, map: {estore_read._embed_map[key]}, "
            f"values: {value[:10] if len(value) > 3 else value[:]}"
        )
