# SPDX-License-Identifier: PROPRIETARY
# SPDX-FileCopyrightText: Copyright The Geneva Authors

"""Checkpoint Store for Geneva Pipeline"""

import abc
import logging
import tempfile
from collections.abc import Iterator
from enum import Enum
from typing import TYPE_CHECKING, Optional

import attrs
import lance
import pyarrow as pa
from lance.file import LanceFileSession
from packaging.version import Version

from geneva.config import ConfigBase
from geneva.utils import retry_lance

if TYPE_CHECKING:
    from lance_namespace import LanceNamespace

_LOG = logging.getLogger(__name__)


class CheckpointStore(abc.ABC):
    """Abstract class for checkpoint store, which is used to store intermediate results
      of Geneva pipelines.

    It is implemented as a key-value store of :class:`pyarrow.RecordBatch` objects.

    This is a lighter weight version of collections.abc.MutableMapping
    where we don't expose length or deletion operations

    TODO: implementations are not consistently handling keys with '/'.  Please avoid it.
    """

    @abc.abstractmethod
    def __contains__(self, item: str) -> bool:
        pass

    @abc.abstractmethod
    def __getitem__(self, item: str) -> pa.RecordBatch:
        pass

    @abc.abstractmethod
    def __setitem__(self, key: str, value: pa.RecordBatch) -> None:
        pass

    @abc.abstractmethod
    def list_keys(self, prefix: str = "") -> Iterator[str]:
        """List all the available keys for check point."""

    @abc.abstractmethod
    def uri(self) -> str:
        pass

    @classmethod
    def from_uri(cls, uri: str) -> "CheckpointStore":
        """Construct a CheckpointStore from a URI."""
        if uri == "memory":
            return InMemoryCheckpointStore()
        try:
            if Version(lance.__version__) < Version("0.35.0b3"):
                _LOG.warning(
                    f"pylance {lance.__version__} has issues at scale.  "
                    "Upgrade to 0.35.0b3 or higher to avoid this."
                )
            return LanceCheckpointStore(uri)
        except Exception as e:
            raise ValueError(f"Invalid checkpoint store uri: {uri}") from e


class InMemoryCheckpointStore(CheckpointStore):
    """In memory checkpoint store for testing purposes."""

    def __init__(self) -> None:
        self._store = {}

    def __repr__(self) -> str:
        return self._store.__repr__()

    def __contains__(self, item: str) -> bool:
        return item in self._store

    def __getitem__(self, item: str) -> pa.RecordBatch:
        return self._store[item]

    def __setitem__(self, key: str, value: pa.RecordBatch) -> None:
        self._store[key] = value

    def list_keys(self, prefix: str = "") -> Iterator[str]:
        for key in self._store:
            if key.startswith(prefix):
                yield key

    def uri(self) -> str:
        return "memory:///"


class LanceCheckpointStore(CheckpointStore):
    """
    Stores checkpoint data as Lance formatted files

    The API mimics a dictionary.

    NOTE: The dict keys are actual paths in a file system and can be vulnerable to
    filesystem traversal attacks.
    """

    # Checkpoint subdirectory name (relative to dataset root)
    _CKP_SUBDIR = "_ckp"

    def __init__(
        self,
        root: str,
        namespace_client: Optional["LanceNamespace"] = None,
        table_id: Optional[list[str]] = None,
    ) -> None:
        self.root = root
        self.namespace_client = namespace_client
        self.table_id = table_id

        # Lazy-initialized runtime state (avoid getting this pickled)
        self._session: Optional[LanceFileSession] = None
        # Track if using namespace-based session (rooted at dataset dir)
        self._use_namespace_session = False

    def __getstate__(self) -> dict:
        """Exclude unpicklable session from pickle."""
        return {
            "root": self.root,
            "namespace_client": self.namespace_client,
            "table_id": self.table_id,
        }

    def __setstate__(self, state: dict) -> None:
        """Restore state from pickle, leaving session uninitialized."""
        self.root = state["root"]
        self.namespace_client = state.get("namespace_client")
        self.table_id = state.get("table_id")
        self._session = None
        self._use_namespace_session = False

    @property
    def session(self) -> "LanceFileSession":
        """Lazily create LanceFileSession on first access."""
        if self._session is None:
            if self.namespace_client and self.table_id:
                # Use dataset's file session for proper credential handling
                ds = lance.dataset(
                    namespace=self.namespace_client, table_id=self.table_id
                )
                self._session = ds.new_file_session()
                self._use_namespace_session = True
            else:
                # Fallback for non-namespace tables
                self._session = LanceFileSession(self.root)
        return self._session

    def _make_path(self, key: str) -> str:
        """Make the full path for a checkpoint key."""
        if self._use_namespace_session:
            # Session is rooted at dataset dir, include _ckp subdir
            return f"{self._CKP_SUBDIR}/{key}.lance"
        return f"{key}.lance"

    @retry_lance
    def __contains__(self, key: str) -> bool:
        _LOG.debug("contains: %s", key)
        return self.session.contains(self._make_path(key))

    @retry_lance
    def __getitem__(self, key: str) -> pa.RecordBatch:
        _LOG.debug("get: %s", key)
        reader = self.session.open_reader(self._make_path(key))
        return reader.read_all().to_table().combine_chunks().to_batches()[0]

    @retry_lance
    def __setitem__(self, key: str, value: pa.RecordBatch) -> None:
        _LOG.debug("set: %s", key)
        path = self._make_path(key)
        with self.session.open_writer(path, schema=value.schema) as writer:
            writer.write_batch(value)

    @retry_lance
    def list_keys(self, prefix: str = "") -> Iterator[str]:
        _LOG.debug("list_keys: %s", prefix)
        # LanceFileSession.list() lists by path prefix, not by filename prefix.
        # Since checkpoint keys are stored as flat files (no '/' separators by
        # convention), we list all keys and filter by string prefix here.
        list_prefix = self._CKP_SUBDIR if self._use_namespace_session else None
        files = self.session.list(list_prefix)
        for file_path in files:
            if not file_path.endswith(".lance"):
                continue
            # Remove _ckp/ prefix if using namespace session
            if self._use_namespace_session and file_path.startswith(
                f"{self._CKP_SUBDIR}/"
            ):
                file_path = file_path[len(self._CKP_SUBDIR) + 1 :]
            key = file_path.removesuffix(".lance")
            if prefix and not key.startswith(prefix):
                continue
            yield key

    def uri(self) -> str:
        return self.root


class CheckpointMode(Enum):
    OBJECT_STORE = "object_store"

    # Store checkpoints in temporary files, for local development
    # and testing.
    # It can be shared between process, i.e., local ray actors.
    TEMPFILE = "tempfile"

    # for testing only
    IN_MEMORY = "in_memory"

    @staticmethod
    def from_str(s: str) -> "CheckpointMode":
        if isinstance(s, CheckpointMode):
            return s
        return CheckpointMode(s)


@attrs.define
class ObjectStoreCheckpointConfig(ConfigBase):
    path: str

    @classmethod
    def name(cls) -> str:
        return "object_store"

    def make(self) -> CheckpointStore:
        return LanceCheckpointStore(self.path)


@attrs.define
class CheckpointConfig(ConfigBase):
    mode: CheckpointMode = attrs.field(
        default=CheckpointMode.OBJECT_STORE, converter=CheckpointMode.from_str
    )

    object_store: ObjectStoreCheckpointConfig | None = attrs.field(default=None)

    @classmethod
    def name(cls) -> str:
        return "checkpoint"

    def make(self) -> CheckpointStore:
        match self.mode:
            case CheckpointMode.TEMPFILE:
                temp_dir = tempfile.mkdtemp()
                _LOG.info("Create checkpoint store on %s", temp_dir)
                return LanceCheckpointStore(temp_dir)
            case CheckpointMode.OBJECT_STORE:
                if self.object_store is None:
                    raise ValueError("CheckpointConfig::object_store is required")
                return LanceCheckpointStore(self.object_store.path)
            case CheckpointMode.IN_MEMORY:
                return InMemoryCheckpointStore()
            case _:
                raise ValueError(f"Unknown checkpoint mode {self.mode}")
