import warnings
from collections.abc import Callable, Iterator
from datetime import datetime, timezone
from functools import lru_cache
from typing import (
    TYPE_CHECKING,
    Any,
    Optional,
)

from pystac import Collection, Extent

from pystac_client._utils import Modifiable, call_modifier
from pystac_client.conformance import ConformanceClasses
from pystac_client.free_text import sqlite_text_search
from pystac_client.item_search import (
    BaseSearch,
    BBox,
    BBoxLike,
    Datetime,
    DatetimeLike,
    FieldsLike,
    FilterLangLike,
    FilterLike,
    QueryLike,
    SortbyLike,
)
from pystac_client.stac_api_io import StacApiIO
from pystac_client.warnings import DoesNotConformTo, PystacClientWarning

if TYPE_CHECKING:
    from pystac_client import client as _client


TemporalInterval = tuple[Optional[datetime], Optional[datetime]]


def temporal_intervals_overlap(
    interval1: TemporalInterval,
    interval2: TemporalInterval,
) -> bool:
    start1, end1 = interval1
    start2, end2 = interval2
    dtmin = datetime.min.replace(tzinfo=timezone.utc)
    dtmax = datetime.max.replace(tzinfo=timezone.utc)

    return (start2 or dtmin) <= (end1 or dtmax) and (start1 or dtmin) <= (end2 or dtmax)


def bboxes_overlap(bbox1: BBox, bbox2: BBox) -> bool:
    xmin1, ymin1, xmax1, ymax1 = bbox1
    xmin2, ymin2, xmax2, ymax2 = bbox2

    return xmin1 <= xmax2 and xmin2 <= xmax1 and ymin1 <= ymax2 and ymin2 <= ymax1


def _extent_matches(
    extent: Extent,
    bbox: BBox | None = None,
    temporal_interval_str: Datetime | None = None,
) -> tuple[bool, bool]:
    bbox_overlaps = not bbox or (
        any(
            bboxes_overlap(bbox, tuple(collection_bbox))
            for collection_bbox in extent.spatial.bboxes
        )
    )

    # check for overlap between the provided temporal interval and the collection's
    # temporal extent
    collection_temporal_extent = extent.temporal

    # process the user-provided temporal interval
    search_temporal_interval = (
        temporal_interval_str.split("/") if temporal_interval_str else []
    )

    # replace .. in open intervals with actual strings
    if search_temporal_interval:
        if search_temporal_interval[0] == "..":
            search_temporal_interval[0] = datetime.min.replace(
                tzinfo=timezone.utc
            ).isoformat()
        if search_temporal_interval[1] == "..":
            search_temporal_interval[1] = datetime.max.replace(
                tzinfo=timezone.utc
            ).isoformat()

    datetime_overlaps = not temporal_interval_str or (
        any(
            temporal_intervals_overlap(
                (
                    datetime.fromisoformat(
                        search_temporal_interval[0].replace("Z", "+00:00")
                    ),
                    datetime.fromisoformat(
                        search_temporal_interval[1].replace("Z", "+00:00")
                    ),
                ),
                (
                    collection_temporal_interval[0],
                    collection_temporal_interval[1],
                ),
            )
            for collection_temporal_interval in collection_temporal_extent.intervals
        )
    )
    return bbox_overlaps, datetime_overlaps


def collection_matches(
    collection_dict: dict[str, Any],
    bbox: BBox | None = None,
    temporal_interval_str: Datetime | None = None,
    q: str | None = None,
) -> bool:
    # check for overlap between provided bbox and the collection's spatial extent
    try:
        extent = Extent.from_dict(collection_dict.get("extent", {}))
    except Exception:
        warnings.warn(
            f"Unable to parse extent from collection={collection_dict.get('id', None)}",
            PystacClientWarning,
        )
        bbox_overlaps = True
        datetime_overlaps = True
    else:
        bbox_overlaps, datetime_overlaps = _extent_matches(
            extent, bbox, temporal_interval_str
        )

    # check for overlap between the provided free-text search query (q) and the
    # collection's title, description, and keywords
    text_fields: dict[str, str] = {
        key: text
        for key, text in collection_dict.items()
        if text and key in ["title", "description", "keywords"]
    }

    if text_fields.get("keywords"):
        text_fields["keywords"] = ", ".join(text_fields["keywords"])

    text_overlaps = not q or sqlite_text_search(q, text_fields)

    return bbox_overlaps and datetime_overlaps and text_overlaps


class CollectionSearch(BaseSearch):
    """Represents a deferred query to a STAC collection search endpoint as described in
    the `STAC API - Collection Search extension
    <https://github.com/stac-api-extensions/collection-search>`__.

    Due to potential conflicts between the collection-search and transactions extensions
    pystac_client will only send GET requests to the /collections endpoint.

    No request is sent to the API until a method is called to iterate
    through the resulting STAC Collections, either :meth:`CollectionSearch.collections`,
    or :meth:`CollectionSearch.collections_as_dicts`.

    All parameters except `url``, ``method``, ``max_collections``, and ``client``
    correspond to query parameters
    described in the `STAC API - Collection Extension: Query Parameters Table
    <https://github.com/stac-api-extensions/collection-search?tab=readme-ov-file#query-parameters-and-fields>`__
    docs. Please refer to those docs for details on how these parameters filter search
    results.

    Args:
        url: The URL to the search page of the STAC API.
        max_collections : The maximum number of collections to return from the
            search, even if there are more matching results. This client to limit
            the total number of Collections returned from the :meth:`collections`,
            :meth:`collection_list`, and :meth:`collections_as_dicts methods`. The
            client will continue to request pages of collections until the number of
            max collections is reached. Setting this to ``None`` will allow
            iteration over a possibly very large number of results.
        stac_io: An instance of StacIO for retrieving results. Normally comes
            from the Client that returns this CollectionSearch client: An instance of a
            root Client used to set the root on resulting Collections.
        client: An instance of Client for retrieving results. This is normally populated
            by the client that returns this CollectionSearch instance.
        limit: A recommendation to the service as to the number of collections to return
            *per page* of results. Defaults to 100.
        bbox: A list, tuple, or iterator representing a bounding box of 2D
            or 3D coordinates. Results will be filtered
            to only those intersecting the bounding box.
        datetime: Either a single datetime or datetime range used to filter results.
            You may express a single datetime using a :class:`datetime.datetime`
            instance, a `RFC 3339-compliant <https://tools.ietf.org/html/rfc3339>`__
            timestamp, or a simple date string (see below). Instances of
            :class:`datetime.datetime` may be either
            timezone aware or unaware. Timezone aware instances will be converted to
            a UTC timestamp before being passed
            to the endpoint. Timezone unaware instances are assumed to represent UTC
            timestamps. You may represent a
            datetime range using a ``"/"`` separated string as described in the spec,
            or a list, tuple, or iterator
            of 2 timestamps or datetime instances. For open-ended ranges, use either
            ``".."`` (``'2020-01-01:00:00:00Z/..'``,
            ``['2020-01-01:00:00:00Z', '..']``) or a value of ``None``
            (``['2020-01-01:00:00:00Z', None]``).

            If using a simple date string, the datetime can be specified in
            ``YYYY-mm-dd`` format, optionally truncating
            to ``YYYY-mm`` or just ``YYYY``. Simple date strings will be expanded to
            include the entire time period, for example:

            - ``2017`` expands to ``2017-01-01T00:00:00Z/2017-12-31T23:59:59Z``
            - ``2017-06`` expands to ``2017-06-01T00:00:00Z/2017-06-30T23:59:59Z``
            - ``2017-06-10`` expands to ``2017-06-10T00:00:00Z/2017-06-10T23:59:59Z``

            If used in a range, the end of the range expands to the end of that
            day/month/year, for example:

            - ``2017/2018`` expands to
              ``2017-01-01T00:00:00Z/2018-12-31T23:59:59Z``
            - ``2017-06/2017-07`` expands to
              ``2017-06-01T00:00:00Z/2017-07-31T23:59:59Z``
            - ``2017-06-10/2017-06-11`` expands to
              ``2017-06-10T00:00:00Z/2017-06-11T23:59:59Z``
        q: Free-text search query. See the `STAC API - Free Text Extension
            Spec <https://github.com/stac-api-extensions/freetext-search>`__ for
            syntax.
        query: List or JSON of query parameters as per the STAC API `query` extension
        filter: JSON of query parameters as per the STAC API `filter` extension
        filter_lang: Language variant used in the filter body. If `filter` is a
            dictionary or not provided, defaults
            to 'cql2-json'. If `filter` is a string, defaults to `cql2-text`.
        sortby: A single field or list of fields to sort the response by
        fields: A list of fields to include in the response. Note this may
            result in invalid STAC objects, as they may not have required fields.
            Use `collections_as_dicts` to avoid object unmarshalling errors.
        modifier : A callable that modifies the children collection and items
            returned by this Client. This can be useful for injecting
            authentication parameters into child assets to access data
            from non-public sources.

            The callable should expect a single argument, which will be one
            of the following types:

            * :class:`pystac.Collection`
            * :class:`pystac.Item`
            * :class:`pystac.ItemCollection`
            * A STAC item-like :class:`dict`
            * A STAC collection-like :class:`dict`

            The callable should mutate the argument in place and return ``None``.

            ``modifier`` propagates recursively to children of this Client.
            After getting a child collection with, e.g.
            :meth:`Client.get_collection`, the child items of that collection
            will still be signed with ``modifier``.
    """

    _stac_io: StacApiIO
    _collection_search_extension_enabled: bool
    _collection_search_free_text_enabled: bool

    def __init__(
        self,
        url: str,
        *,
        max_collections: int | None = None,
        stac_io: StacApiIO | None = None,
        client: Optional["_client.Client"] = None,
        limit: int | None = None,
        bbox: BBoxLike | None = None,
        datetime: DatetimeLike | None = None,
        query: QueryLike | None = None,
        filter: FilterLike | None = None,
        filter_lang: FilterLangLike | None = None,
        sortby: SortbyLike | None = None,
        fields: FieldsLike | None = None,
        q: str | None = None,
        modifier: Callable[[Modifiable], None] | None = None,
        collection_search_extension_enabled: bool = False,
        collection_search_free_text_enabled: bool = False,
    ):
        super().__init__(
            url=url,
            method="GET",
            max_items=max_collections,
            stac_io=stac_io,
            client=client,
            limit=limit,
            bbox=bbox,
            datetime=datetime,
            query=query,
            filter=filter,
            filter_lang=filter_lang,
            sortby=sortby,
            fields=fields,
            q=q,
            modifier=modifier,
        )

        if client and client._stac_io is not None and stac_io is None:
            self._stac_io = client._stac_io
            self._collection_search_extension_enabled = client.conforms_to(
                ConformanceClasses.COLLECTION_SEARCH
            )
            self._collection_search_free_text_enabled = client.conforms_to(
                ConformanceClasses.COLLECTION_SEARCH_FREE_TEXT
            )
            if any([bbox, datetime, q, query, filter, sortby, fields]):
                if not self._collection_search_extension_enabled:
                    warnings.warn(
                        str(DoesNotConformTo("COLLECTION_SEARCH"))
                        + ". Filtering will be performed client-side where only bbox, "
                        "datetime, and q arguments are supported"
                    )
                    self._validate_client_side_args()
                else:
                    if not self._collection_search_free_text_enabled:
                        warnings.warn(
                            str(DoesNotConformTo("COLLECTION_SEARCH#FREE_TEXT"))
                            + ". Free-text search is not enabled for collection search"
                            "Free-text filters will be applied client-side."
                        )

        else:
            self._stac_io = stac_io or StacApiIO()
            self._collection_search_extension_enabled = (
                collection_search_extension_enabled
            )
            self._collection_search_free_text_enabled = (
                collection_search_free_text_enabled
            )
            self._validate_client_side_args()

    def _validate_client_side_args(self) -> None:
        """Client-side filtering only supports the bbox, datetime, and q parameters."""
        args = {key for key, value in self._parameters.items() if value is not None}
        extras = args - {"bbox", "datetime", "q", "limit", "max_items"}

        if extras:
            raise ValueError(
                "Only the limit, max_collections, bbox, datetime, and q arguments are "
                "supported for client-side filtering but these extra arguments were "
                "provided: " + ",".join(extras)
            )

    @lru_cache(1)
    def matched(self) -> int | None:
        """Return number matched for search

        Returns the value from the `numberMatched` or `context.matched` field.
        Not all APIs will support counts in which case a warning will be issued

        Returns:
            int: Total count of matched collections. If counts are not supported `None`
            is returned.
        """
        found = None
        iter = self.pages_as_dicts()
        page = next(iter, None)

        if not page:
            return 0

        # if collection search and free-text are fully supported, try reading a value
        # from the search result context
        if (
            self._collection_search_extension_enabled
            and self._collection_search_free_text_enabled
        ):
            if "context" in page:
                found = page["context"].get("matched", None)
            elif "numberMatched" in page:
                found = page["numberMatched"]

        if not found:
            count = len(page["collections"])

            for page in iter:
                count += len(page["collections"])

            found = count

        return found

    # ------------------------------------------------------------------------
    # Result sets
    # ------------------------------------------------------------------------
    # By collection
    def collections(self) -> Iterator[Collection]:
        """Iterator that yields :class:`pystac.Collection` instances for each collection
        matching the given search parameters.

        Yields:
            Collection : each Collection matching the search criteria
        """
        for collection in self.collections_as_dicts():
            # already signed in collections_as_dicts
            yield Collection.from_dict(
                collection, root=self.client, preserve_dict=False
            )

    @lru_cache(1)
    def collections_as_dicts(self) -> Iterator[dict[str, Any]]:
        """Iterator that yields :class:`dict` instances for each collection matching
        the given search parameters.

        Yields:
            Collection : each Collection matching the search criteria
        """
        for page in self.pages_as_dicts():
            yield from page.get("collections", [])

    # ------------------------------------------------------------------------
    # By Page
    def pages(self) -> Iterator[list[Collection]]:
        """Iterator that yields lists of Collection objects.  Each list is
        a page of results from the search.

        Yields:
            List[Collection] : a group of Collections matching the search criteria
        """
        if isinstance(self._stac_io, StacApiIO):
            for page in self.pages_as_dicts():
                # already signed in pages_as_dicts
                yield [
                    Collection.from_dict(
                        collection, preserve_dict=False, root=self.client
                    )
                    for collection in page["collections"]
                ]

    def pages_as_dicts(self) -> Iterator[dict[str, Any]]:
        """Iterator that yields :class:`dict` instances for each page
        of results from the search.

        Yields:
            Dict : a group of collections matching the search
            criteria as a feature-collection-like dictionary.
        """
        if isinstance(self._stac_io, StacApiIO):
            num_collections = 0
            for page in self._stac_io.get_pages(
                self.url, self.method, self.get_parameters()
            ):
                call_modifier(self.modifier, page)
                collections = page.get("collections", [])
                page_has_collections = len(collections) > 0

                # apply client-side filter if the collection search extension
                # is not enabled in the API
                if not self._collection_search_extension_enabled:
                    args = {
                        "bbox": self._parameters.get("bbox"),
                        "temporal_interval_str": self._parameters.get("datetime"),
                        "q": self._parameters.get("q"),
                    }
                    collections = [
                        collection
                        for collection in filter(
                            lambda x: collection_matches(x, **args),
                            collections,
                        )
                    ]

                # apply client-side free-text filter if free-text extension is not
                # enabled in the API
                elif not self._collection_search_free_text_enabled:
                    if q := self._parameters.get("q"):
                        collections = [
                            collection
                            for collection in filter(
                                lambda x: collection_matches(x, q=q),
                                collections,
                            )
                        ]
                if collections:
                    num_collections += len(collections)
                    if self._max_items and num_collections > self._max_items:
                        # Slice the features down to make sure we hit max_collections
                        page["collections"] = collections[
                            0 : -(num_collections - self._max_items)
                        ]
                    else:
                        page["collections"] = collections

                    yield page
                    if self._max_items and num_collections >= self._max_items:
                        return
                # if there were collections on this page but they got filtered out keep
                # going
                elif page_has_collections:
                    continue
                else:
                    return

    # ------------------------------------------------------------------------
    # Everything

    @lru_cache(1)
    def collection_list(self) -> list[Collection]:
        """
        Get the matching collections as a list of :py:class:`pystac.Collection` objects.

        Return:
            List[Collection]: The list of collections
        """
        # Bypass the cache here, so that we can pass __preserve_dict__
        # without mutating what's in the cache.
        collection_list = self.collections_as_dicts.__wrapped__(self)
        # already signed in collections_as_dicts
        return [
            Collection.from_dict(collection, preserve_dict=False, root=self.client)
            for collection in collection_list
        ]

    @lru_cache(1)
    def collection_list_as_dict(self) -> dict[str, Any]:
        """
        Get the matching collections as a dict.

        The dictionary will have a single key:

        1. ``'collections'`` with the value being a list of dictionaries
            for the matching collections.

        Return:
            Dict : A dictionary with the list of matching collections
        """
        collections = []
        for page in self.pages_as_dicts():
            for collection in page["collections"]:
                collections.append(collection)

        return {"collections": collections}