from __future__ import annotations

import warnings
from typing import TYPE_CHECKING, Any, Literal, overload

import anndata as ad
import pandas as pd
from django.db import models
from django.db.models import CASCADE, PROTECT, Q
from lamin_utils import logger
from lamindb_setup.core.hashing import HASH_LENGTH, hash_set

from lamindb.base.fields import (
    CharField,
    ForeignKey,
    OneToOneField,
    TextField,
)

from ..base.ids import base62_20
from ..core._mapped_collection import MappedCollection
from ..core.storage._backed_access import _open_dataframe
from ..errors import FieldValidationError
from ..models._is_versioned import process_revises
from ._is_versioned import IsVersioned
from .artifact import (
    Artifact,
    describe_artifact_collection,
    get_run,
    populate_subsequent_run,
    save_schema_links,
    track_run_input,
)
from .has_parents import view_lineage
from .run import Run, TracksRun, TracksUpdates
from .sqlrecord import (
    BaseSQLRecord,
    IsLink,
    SQLRecord,
    _get_record_kwargs,
    init_self_from_db,
    update_attributes,
)

if TYPE_CHECKING:
    from collections.abc import Iterable, Iterator

    from polars import LazyFrame as PolarsLazyFrame
    from pyarrow.dataset import Dataset as PyArrowDataset

    from ..core.storage import UPath
    from .block import CollectionBlock
    from .project import Project, Reference
    from .query_set import QuerySet
    from .transform import Transform
    from .ulabel import ULabel


# below is a draft for the future, see also the tests in test_collection.py
#
# class CollectionFeatureManager:
#     """Query features of artifact in collection."""

#     def __init__(self, collection: Collection):
#         self._collection = collection

#     def _get_staged_feature_sets_union(self) -> dict[str, Schema]:
#         links_schema_artifact = Artifact.feature_sets.through.objects.filter(
#             artifact_id__in=self._collection.artifacts.values_list("id", flat=True)
#         )
#         feature_sets_by_slots = defaultdict(list)
#         for link in links_schema_artifact:
#             feature_sets_by_slots[link.slot].append(link.schema_id)
#         feature_sets_union = {}
#         for slot, schema_ids_slot in feature_sets_by_slots.items():
#             schema_1 = Schema.get(id=schema_ids_slot[0])
#             related_name = schema_1._get_related_name()
#             features_registry = getattr(Schema, related_name).field.model
#             # this way of writing the __in statement turned out to be the fastest
#             # evaluated on a link table with 16M entries connecting 500 feature sets with
#             # 60k genes
#             feature_ids = (
#                 features_registry.schemas.through.objects.filter(
#                     schema_id__in=schema_ids_slot
#                 )
#                 .values(f"{features_registry.__name__.lower()}_id")
#                 .distinct()
#             )
#             features = features_registry.filter(id__in=feature_ids)
#             feature_sets_union[slot] = Schema(features, dtype=schema_1.dtype)
#         return feature_sets_union


def _load_concat_artifacts(
    artifacts: list[Artifact], join: Literal["inner", "outer"] = "outer", **kwargs
) -> pd.DataFrame | ad.AnnData:
    suffixes = {artifact.suffix for artifact in artifacts}
    # Why is that? - Sergei
    if len(suffixes) != 1:
        raise ValueError(
            "Can only load collections where all artifacts have the same suffix"
        )

    # because we're tracking data flow on the collection-level, here, we don't
    # want to track it on the artifact-level
    first_object = artifacts[0].load(is_run_input=False)
    is_dataframe = isinstance(first_object, pd.DataFrame)
    is_anndata = isinstance(first_object, ad.AnnData)
    if not is_dataframe and not is_anndata:
        raise ValueError(f"Unable to concatenate {suffixes.pop()} objects.")

    objects = [first_object]
    artifact_uids = [artifacts[0].uid]
    for artifact in artifacts[1:]:
        objects.append(artifact.load(is_run_input=False))
        artifact_uids.append(artifact.uid)

    if is_dataframe:
        concat_object = pd.concat(objects, join=join, **kwargs)
    elif is_anndata:
        label = kwargs.pop("label", "artifact_uid")
        keys = kwargs.pop("keys", artifact_uids)
        concat_object = ad.concat(objects, join=join, label=label, keys=keys, **kwargs)
    return concat_object


class Collection(SQLRecord, IsVersioned, TracksRun, TracksUpdates):
    """Versioned collections of artifacts.

    Args:
        artifacts: `Artifact | list[Artifact]` One or several artifacts.
        key: `str` A file-path like key, analogous to the `key` parameter of `Artifact` and `Transform`.
        description: `str | None = None` A description.
        revises: `Collection | None = None` An old version of the collection.
        run: `Run | None = None` The run that creates the collection.
        meta: `Artifact | None = None` An artifact that defines metadata for the collection.
        reference: `str | None = None` A simple reference, e.g. an external ID or a URL.
        reference_type: `str | None = None` A way to indicate to indicate the type of the simple reference `"url"`.

    See Also:
        :class:`~lamindb.Artifact`

    Examples:

        Create a collection from a list of :class:`~lamindb.Artifact` objects:

        >>> collection = ln.Collection([artifact1, artifact2], key="my_project/my_collection")

        Create a collection that groups a data & a metadata artifact (e.g., here :doc:`docs:rxrx`):

        >>> collection = ln.Collection(data_artifact, key="my_project/my_collection", meta=metadata_artifact)

    """

    class Meta(SQLRecord.Meta, IsVersioned.Meta, TracksRun.Meta, TracksUpdates.Meta):
        abstract = False
        app_label = "lamindb"
        constraints = [
            models.UniqueConstraint(
                fields=["key", "hash"],
                name="unique_collection_key_hash_not_null",
            )
        ]

    _len_full_uid: int = 20
    _len_stem_uid: int = 16
    _name_field: str = "key"

    id: int = models.AutoField(primary_key=True)
    """Internal id, valid only in one DB instance."""
    uid: str = CharField(
        editable=False,
        unique=True,
        db_index=True,
        max_length=_len_full_uid,
        default=base62_20,
    )
    """Universal id, valid across DB instances."""
    key: str = CharField(db_index=True)
    """Name or path-like key."""
    # below is the only case in which we use a TextField
    # for description; we do so because users had descriptions exceeding 255 chars
    # in their instances
    description: str | None = TextField(null=True)
    """A description or title."""
    hash: str | None = CharField(
        max_length=HASH_LENGTH,
        db_index=True,
        null=True,
    )
    """Hash of collection content."""
    reference: str | None = CharField(max_length=255, db_index=True, null=True)
    """A reference like URL or external ID."""
    # also for reference_type here, we allow an extra long max_length
    reference_type: str | None = CharField(max_length=25, db_index=True, null=True)
    """Type of reference, e.g., cellxgene Census collection_id."""
    ulabels: ULabel = models.ManyToManyField(
        "ULabel", through="CollectionULabel", related_name="collections"
    )
    """ULabels annotating the collection (see :class:`~lamindb.Feature`)."""
    run: Run | None = ForeignKey(
        Run, PROTECT, related_name="output_collections", null=True, default=None
    )
    """:class:`~lamindb.Run` that created the `collection`."""
    input_of_runs: Run = models.ManyToManyField(Run, related_name="input_collections")
    """Runs that use this collection as an input."""
    _subsequent_runs: Run = models.ManyToManyField(
        "Run",
        related_name="_recreated_collections",
        db_table="lamindb_collection__previous_runs",  # legacy name, change in lamindb v2
    )
    """Runs that re-created the record after initial creation."""
    artifacts: Artifact = models.ManyToManyField(
        "Artifact", related_name="collections", through="CollectionArtifact"
    )
    """Artifacts in collection."""
    meta_artifact: Artifact | None = OneToOneField(
        "Artifact",
        PROTECT,
        null=True,
        unique=True,
        related_name="_meta_of_collection",
    )
    """An artifact that stores metadata that indexes a collection.

    It has a 1:1 correspondence with an artifact. If needed, you can access the
    collection from the artifact via a private field:
    `artifact._meta_of_collection`.
    """
    _actions: Artifact = models.ManyToManyField(Artifact, related_name="+")
    """Actions to attach for the UI."""
    projects: Project
    """Linked projects."""
    references: Reference
    """Linked references."""
    blocks: CollectionBlock
    """Blocks that annotate this collection."""

    @overload
    def __init__(
        self,
        artifacts: Artifact | list[Artifact],
        key: str,
        description: str | None = None,
        meta: Any | None = None,
        reference: str | None = None,
        reference_type: str | None = None,
        run: Run | None = None,
        revises: Collection | None = None,
        skip_hash_lookup: bool = False,
    ): ...

    @overload
    def __init__(
        self,
        *db_args,
    ): ...

    def __init__(
        self,
        *args,
        **kwargs,
    ):
        if len(args) == len(self._meta.concrete_fields):
            super().__init__(*args, **kwargs)
            return None
        # now we proceed with the user-facing constructor
        if len(args) > 1:
            raise ValueError("Only one non-keyword arg allowed: artifacts")
        artifacts: Artifact | list[Artifact] = (
            kwargs.pop("artifacts") if len(args) == 0 else args[0]
        )
        meta_artifact: Artifact | None = kwargs.pop("meta_artifact", None)
        tmp_key: str | None = kwargs.pop("key", None)
        description: str | None = kwargs.pop("description", None)
        reference: str | None = kwargs.pop("reference", None)
        reference_type: str | None = kwargs.pop("reference_type", None)
        run: Run | None = kwargs.pop("run", None)
        revises: Collection | None = kwargs.pop("revises", None)
        version: str | None = kwargs.pop("version", None)
        skip_hash_lookup: bool = kwargs.pop("skip_hash_lookup", False)
        branch = kwargs.pop("branch", None)
        branch_id = kwargs.pop("branch_id", 1)
        space = kwargs.pop("space", None)
        space_id = kwargs.pop("space_id", 1)
        key: str
        if "name" in kwargs:
            key = kwargs.pop("name")
            warnings.warn(
                f"argument `name` will be removed, please pass {key} to `key` instead",
                FutureWarning,
                stacklevel=2,
            )
        else:
            key = tmp_key
        if not len(kwargs) == 0:
            valid_keywords = ", ".join(
                [val[0] for val in _get_record_kwargs(Collection)]
            )
            raise FieldValidationError(
                f"Only {valid_keywords} can be passed, you passed: {kwargs}"
            )
        if revises is None:
            revises = (
                Collection.filter(key=key, is_latest=True)
                .order_by("-created_at")
                .first()
            )
        provisional_uid, version, key, description, revises = process_revises(
            revises, version, key, description, Collection
        )
        run = get_run(run)
        if isinstance(artifacts, Artifact):
            artifacts = [artifacts]
        else:
            if not hasattr(artifacts, "__getitem__"):
                raise ValueError("Artifact or list[Artifact] is allowed.")
            assert isinstance(artifacts[0], Artifact)  # type: ignore  # noqa: S101
        hash = from_artifacts(artifacts)  # type: ignore
        if meta_artifact is not None:
            if not isinstance(meta_artifact, Artifact):
                raise ValueError("meta_artifact has to be an Artifact")
            if isinstance(meta_artifact, Artifact):
                if meta_artifact._state.adding:
                    raise ValueError(
                        "Save meta_artifact artifact before creating collection!"
                    )
        # we ignore collections in trash containing the same hash
        if hash is not None and not skip_hash_lookup:
            existing_collection = Collection.objects.filter(
                ~Q(branch_id=-1),
                hash=hash,
            ).one_or_none()
        else:
            existing_collection = None
        if existing_collection is not None:
            logger.warning(
                f"returning collection with same hash: {existing_collection}; if you intended to query to track this collection as an input, use: ln.Collection.get()"
            )
            init_self_from_db(self, existing_collection)
            update_attributes(self, {"description": description, "key": key})
            populate_subsequent_run(self, run)
        else:
            _skip_validation = revises is not None and key == revises.key
            super().__init__(  # type: ignore
                uid=provisional_uid,
                key=key,
                description=description,
                reference=reference,
                reference_type=reference_type,
                meta_artifact=meta_artifact,
                hash=hash,
                run=run,
                version=version,
                branch=branch,
                branch_id=branch_id,
                space=space,
                space_id=space_id,
                revises=revises,
                _skip_validation=_skip_validation,
            )
        self._artifacts = artifacts
        if revises is not None and revises.uid != self.uid:
            track_run_input(revises, run=run)
        track_run_input(artifacts, run=run)

    @classmethod
    def get(
        cls,
        idlike: int | str | None = None,
        *,
        is_run_input: bool | Run = False,
        **expressions,
    ) -> Artifact:
        """Get a single collection.

        Args:
            idlike: Either a uid stub, uid or an integer id.
            is_run_input: Whether to track this collection as run input.
            expressions: Fields and values passed as Django query expressions.

        Raises:
            :exc:`docs:lamindb.errors.DoesNotExist`: In case no matching record is found.

        See Also:
            - Method in `SQLRecord` base class: :meth:`~lamindb.models.SQLRecord.get`

        Examples:

            ::

                collection = ln.Collection.get("okxPW6GIKBfRBE3B0000")
                collection = ln.Collection.get(key="scrna/collection1")
        """
        from .query_set import QuerySet

        return QuerySet(model=cls).get(idlike, is_run_input=is_run_input, **expressions)

    def append(self, artifact: Artifact, run: Run | None = None) -> Collection:
        """Append an artifact to the collection.

        This does not modify the original collection in-place, but returns a new version
        of the original collection with the appended artifact.

        Args:
            artifact: An artifact to add to the collection.
            run: The run that creates the new version of the collection.

        Examples:

            ::

                collection_v1 = ln.Collection(artifact, key="My collection").save()
                collection_v2 = collection.append(another_artifact)  # returns a new version of the collection
                collection_v2.save()  # save the new version

        """
        return Collection(  # type: ignore
            self.artifacts.all().to_list() + [artifact],
            # key is automatically derived from revises.key
            description=self.description,
            revises=self,
            run=run,
        )

    def open(
        self,
        engine: Literal["pyarrow", "polars"] = "pyarrow",
        is_run_input: bool | None = None,
        **kwargs,
    ) -> PyArrowDataset | Iterator[PolarsLazyFrame]:
        """Open a dataset for streaming.

        Works for `pyarrow` and `polars` compatible formats
        (`.parquet`, `.csv`, `.ipc` etc. files or directories with such files).

        Args:
            engine: Which module to use for lazy loading of a dataframe
                from `pyarrow` or `polars` compatible formats.
            is_run_input: Whether to track this artifact as run input.
            **kwargs: Keyword arguments for `pyarrow.dataset.dataset` or `polars.scan_*` functions.

        Notes:
            For more info, see guide: :doc:`/arrays`.
        """
        if self._state.adding:
            artifacts = self._artifacts
            logger.warning("the collection isn't saved, consider calling `.save()`")
        else:
            artifacts = self.ordered_artifacts.all()
        paths = [artifact.path for artifact in artifacts]

        dataframe = _open_dataframe(paths, engine=engine, **kwargs)
        # track only if successful
        track_run_input(self, is_run_input)
        return dataframe

    def mapped(
        self,
        layers_keys: str | list[str] | None = None,
        obs_keys: str | list[str] | None = None,
        obsm_keys: str | list[str] | None = None,
        obs_filter: dict[str, str | list[str]] | None = None,
        join: Literal["inner", "outer"] | None = "inner",
        encode_labels: bool | list[str] = True,
        unknown_label: str | dict[str, str] | None = None,
        cache_categories: bool = True,
        parallel: bool = False,
        dtype: str | None = None,
        stream: bool = False,
        is_run_input: bool | None = None,
    ) -> MappedCollection:
        """Return a map-style dataset.

        Returns a `pytorch map-style dataset
        <https://pytorch.org/docs/stable/data.html#map-style-datasets>`__ by
        virtually concatenating `AnnData` arrays.

        By default (`stream=False`) `AnnData` arrays are moved into a local
        cache first.

        `__getitem__` of the `MappedCollection` object takes a single integer index
        and returns a dictionary with the observation data sample for this index from
        the `AnnData` objects in the collection. The dictionary has keys for `layers_keys`
        (`.X` is in `"X"`), `obs_keys`, `obsm_keys` (under `f"obsm_{key}"`) and also `"_store_idx"`
        for the index of the `AnnData` object containing this observation sample.

        .. note::

            For a guide, see :doc:`docs:scrna-mappedcollection`.

            This method currently only works for collections or query sets of `AnnData` artifacts.

        Args:
            layers_keys: Keys from the ``.layers`` slot. ``layers_keys=None`` or ``"X"`` in the list
                retrieves ``.X``.
            obs_keys: Keys from the ``.obs`` slots.
            obsm_keys: Keys from the ``.obsm`` slots.
            obs_filter: Select only observations with these values for the given obs columns.
                Should be a dictionary with obs column names as keys
                and filtering values (a string or a list of strings) as values.
            join: `"inner"` or `"outer"` virtual joins. If ``None`` is passed,
                does not join.
            encode_labels: Encode labels into integers.
                Can be a list with elements from ``obs_keys``.
            unknown_label: Encode this label to -1.
                Can be a dictionary with keys from ``obs_keys`` if ``encode_labels=True``
                or from ``encode_labels`` if it is a list.
            cache_categories: Enable caching categories of ``obs_keys`` for faster access.
            parallel: Enable sampling with multiple processes.
            dtype: Convert numpy arrays from ``.X``, ``.layers`` and ``.obsm``
            stream: Whether to stream data from the array backend.
            is_run_input: Whether to track this collection as run input.

        Examples:
            >>> import lamindb as ln
            >>> from torch.utils.data import DataLoader
            >>> ds = ln.Collection.get(description="my collection")
            >>> mapped = collection.mapped(obs_keys=["cell_type", "batch"])
            >>> dl = DataLoader(mapped, batch_size=128, shuffle=True)
            >>> # also works for query sets of artifacts, '...' represents some filtering condition
            >>> # additional filtering on artifacts of the collection
            >>> mapped = collection.artifacts.all().filter(...).order_by("-created_at").mapped()
            >>> # or directly from a query set of artifacts
            >>> mapped = ln.Artifact.filter(..., otype="AnnData").order_by("-created_at").mapped()
        """
        path_list = []
        if self._state.adding:
            artifacts = self._artifacts
            logger.warning("the collection isn't saved, consider calling `.save()`")
        else:
            artifacts = self.ordered_artifacts.all()
        for artifact in artifacts:
            if ".h5ad" not in artifact.suffix and ".zarr" not in artifact.suffix:
                logger.warning(f"ignoring artifact with suffix {artifact.suffix}")
                continue
            elif not stream:
                path_list.append(artifact.cache())
            else:
                path_list.append(artifact.path)
        ds = MappedCollection(
            path_list,
            layers_keys,
            obs_keys,
            obsm_keys,
            obs_filter,
            join,
            encode_labels,
            unknown_label,
            cache_categories,
            parallel,
            dtype,
        )
        # track only if successful
        track_run_input(self, is_run_input)
        return ds

    def cache(self, is_run_input: bool | None = None) -> list[UPath]:
        """Download cloud artifacts in collection to local cache.

        Follows syncing logic: only downloads outdated artifacts.

        Returns ordered paths to locally cached on-disk artifacts via `.ordered_artifacts.all()`:

        Args:
            is_run_input: Whether to track this collection as run input.
        """
        path_list = []
        for artifact in self.ordered_artifacts.all():
            # do not want to track data lineage on the artifact level
            path_list.append(artifact.cache(is_run_input=False))
        track_run_input(self, is_run_input)
        return path_list

    def load(
        self,
        join: Literal["inner", "outer"] = "outer",
        is_run_input: bool | None = None,
        **kwargs,
    ) -> pd.DataFrame | ad.AnnData:
        """Cache and load to memory.

        Returns an in-memory concatenated `DataFrame` or `AnnData` object.
        """
        # cannot call track_run_input here, see comment further down
        artifacts = self.ordered_artifacts.all()
        concat_object = _load_concat_artifacts(artifacts, join, **kwargs)
        # only call it here because there might be errors during load or concat
        track_run_input(self, is_run_input)
        return concat_object

    def save(self, using: str | None = None) -> Collection:
        """Save the collection and underlying artifacts to database & storage.

        Args:
            using: The database to which you want to save.

        Examples:
            >>> collection = ln.Collection("./myfile.csv", name="myfile")
        """
        if self.meta_artifact is not None:
            self.meta_artifact.save()
        super().save()
        # we don't allow updating the collection of artifacts
        # if users want to update the set of artifacts, they
        # have to create a new collection
        if hasattr(self, "_artifacts"):
            links = [
                CollectionArtifact(collection_id=self.id, artifact_id=artifact.id)  # type: ignore
                for artifact in self._artifacts
            ]
            # the below seems to preserve the order of the list in the
            # auto-incrementing integer primary
            # merely using .artifacts.set(*...) doesn't achieve this
            # we need ignore_conflicts=True so that this won't error if links already exist
            CollectionArtifact.objects.bulk_create(links, ignore_conflicts=True)
        save_schema_links(self)
        if using is not None:
            logger.warning("using argument is ignored")
        return self

    def restore(self) -> None:
        """Restore collection record from trash.

        Examples:

            For any `Collection` object `collection`, call:

            >>> collection.restore()
        """
        self.branch_id = 1
        self.save()

    @property
    def transform(self) -> Transform | None:
        """Transform whose run created the collection."""
        return self.run.transform if self.run is not None else None

    @property
    def name(self) -> str:
        """Name of the collection.

        Splits `key` on `/` and returns the last element.
        """
        return self.key.split("/")[-1]

    @property
    def ordered_artifacts(self) -> QuerySet:
        """Ordered `QuerySet` of `.artifacts`.

        Accessing the many-to-many field `collection.artifacts` directly gives
        you non-deterministic order.

        Using the property `.ordered_artifacts` allows to iterate through a set
        that's ordered by the order of the list that created the collection.
        """
        return self.artifacts.order_by("links_collection__id")

    @property
    def data_artifact(self) -> Artifact | None:
        """Access to a single data artifact.

        If the collection has a single data & metadata artifact, this allows access via::

           collection.data_artifact  # first & only element of collection.artifacts
           collection.meta_artifact  # metadata

        """
        return self.artifacts.first()

    def describe(self) -> None:
        """Describe relations of record.

        Examples:
            >>> artifact.describe()
        """
        return describe_artifact_collection(self)


# internal function, not exposed to user
def from_artifacts(artifacts: Iterable[Artifact]) -> tuple[str, dict[str, str]]:
    # assert all artifacts are already saved
    saved = not any(artifact._state.adding for artifact in artifacts)
    if not saved:
        raise ValueError("Not all artifacts are yet saved, please save them")
    # validate consistency of hashes - we do not allow duplicate hashes
    hashes = [artifact.hash for artifact in artifacts if artifact.hash is not None]
    hashes_set = set(hashes)
    if len(hashes) != len(hashes_set):
        seen = set()
        non_unique = [x for x in hashes if x in seen or seen.add(x)]  # type: ignore
        raise ValueError(
            "Please pass artifacts with distinct hashes: these ones are non-unique"
            f" {non_unique}"
        )
    hash = hash_set(hashes_set)
    return hash


class CollectionArtifact(BaseSQLRecord, IsLink, TracksRun):
    id: int = models.BigAutoField(primary_key=True)
    collection: Collection = ForeignKey(
        Collection, CASCADE, related_name="links_artifact"
    )
    artifact: Artifact = ForeignKey(Artifact, PROTECT, related_name="links_collection")

    class Meta:
        app_label = "lamindb"
        unique_together = ("collection", "artifact")


# mypy: ignore-errors
Collection.view_lineage = view_lineage
