import datetime
import logging
import uuid
from collections.abc import Iterator
from typing import Any, Literal

import requests
from pydantic import TypeAdapter

from earthscale.v1._http_client import _EarthscaleHttpClient
from earthscale.v1.exceptions import (
    VersionIncompatibleError,
)
from earthscale.v1.models import (
    AddDatasetRequest,
    AddDatasetResponse,
    AddImageDatasetRequest,
    AddTileServerDatasetRequest,
    AddVectorDatasetRequest,
    AddZarrDatasetRequest,
    DatasetLabel,
    DatasetResponse,
    DeleteDatasetResponse,
    FilenameBandPattern,
    ListDatasetResponse,
    PaginatedListDatasetResponse,
    PydanticList,
    VectorCoordinatePrecision,
    VersionCheckResponse,
)
from earthscale.v1.settings import EarthscaleSettings

_logger = logging.getLogger("earthscale")


class EarthscaleClient:
    """Client for the Earthscale API.

    Can be used to add, list, and retrieve datasets.

    Note:
        There are two ways to authenticate with Earthscale:

        1. **Email and password environment variables:**
           Set the ``EARTHSCALE_EMAIL`` and ``EARTHSCALE_PASSWORD`` environment
           variables

        2. **OAuth:**
           If the environment variables are not set, a browser window will be opened to
           authenticate using OAuth. This only works if running in a graphical environment
           where a browser window can be opened.

    Examples:
        Usage with environment variables:

        .. code-block:: python

            import os
            from earthscale import EarthscaleClient

            os.environ["EARTHSCALE_EMAIL"] = "service@example.com"
            os.environ["EARTHSCALE_PASSWORD"] = "service_password"

            with EarthscaleClient() as client:
                client.list_datasets()
    """

    API_VERSION = "v1"

    def __init__(
        self,
        api_url: str | None = None,
        auth_url: str | None = None,
        anon_key: str | None = None,
        skip_version_check: bool = False,
        session: requests.Session | None = None,
        use_proxy: bool = False,
    ):
        """Initialize the Earthscale client.

        Args:
            api_url: Custom URL of the Earthscale API.
                Defaults to https://api.earthscale.ai.
            auth_url: Custom URL for the authentication service.
                Defaults to https://supabase.earthscale.ai.
            anon_key: Custom anon key for the authentication service.
                If not set, the client will use the default anon key for
                https://supabase.earthscale.ai.
            skip_version_check: Whether to skip version compatibility check.
            session: Optional custom requests session to use.
            use_proxy: Whether to use the proxy server for authentication
        """
        self._skip_version_check = skip_version_check

        settings = EarthscaleSettings()
        api_url = api_url or settings.api_url
        if use_proxy:
            auth_url = settings.auth_proxy_url
        else:
            auth_url = auth_url or settings.auth_url

        anon_key = anon_key or settings.auth_anon_key

        # Create the HTTP client
        self._http_client = _EarthscaleHttpClient(
            api_url=api_url,
            auth_url=auth_url,
            anon_key=anon_key,
            email=settings.email,
            password=settings.password,
            credentials_file=settings.credentials_file,
            session=session,
        )

    def __enter__(self) -> "EarthscaleClient":
        if not self._skip_version_check:
            self.check_api_support()
        # this triggers a login if not already logged in, or refreshes the token
        self._http_client.refresh_token()
        if self._http_client.user_email:
            _logger.info(
                "Authenticated as '%s' on API server '%s' using auth server '%s'",
                self._http_client.user_email,
                self._http_client.api_url,
                self._http_client.auth_url,
            )

        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: Any,
    ) -> None:
        # Close the HTTP client session
        self._http_client.close_session()

    def login(self) -> None:
        """Login using service account credentials from environment variables,
        or OAuth authentication if no environment variables are set.

        Raises:
            AuthenticationError: If authentication fails.
        """
        self._http_client.login()

    def _add_dataset(self, request: AddDatasetRequest) -> AddDatasetResponse:
        """Internal method to add a dataset using a request model.

        Args:
            request: The dataset request model.

        Returns:
            The dataset response.
        """
        raw_response = self._http_client.request(
            method="POST",
            endpoint="datasets",
            data=request,
            api_version=self.API_VERSION,
        )
        return AddDatasetResponse.model_validate_json(raw_response)

    def add_image_dataset(
        self,
        name: str,
        url: str | list[str],
        labels: list[DatasetLabel] | None = None,
        tags: dict[str, str] | None = None,
        bands: list[str] | None = None,
        groupby: str | None = None,
        filename_date_pattern: str | None = None,
        filename_band_pattern: list[dict[str, str]] | None = None,
        visualization_optimization: bool | Literal["auto"] = "auto",
        pixel_info_optimizations: list[str] | None = None,
    ) -> AddDatasetResponse:
        """Add an image dataset. Images must be in a format that can be
        read by rasterio, e.g. GeoTIFFs.

        This function supports creating a time dimension through the
        filename_date_pattern argument. This pattern uses strftime-style format
        codes to extract date information from the filenames.

        Examples:
            For filenames like `brasil_coverage_2011.tif`:

            .. code-block:: python

                from earthscale import EarthscaleClient

                with EarthscaleClient() as client:
                    client.add_image_dataset(
                        name="my_dataset",
                        url="gs://mybucket/my_dataset/brasil_coverage_*.tif",
                        filename_date_pattern="%Y",
                    )

        Args:
            name: The name of the dataset. Creates a new version of an existing dataset
                  if the latest version of a dataset has the same name.
            url: The URL or list of URLs of the dataset.
            labels: Optional. User-defined labels as key-value pairs. Deprecated, use
                `tags` instead. If only labels are provided, a warning will be logged
                and the labels will be converted to tags. If both `labels` and `tags`
                are provided, `tags` will take precedence.
            tags: Optional. User-defined tags as key-value pairs.
            bands: Optional list of bands to include.
            groupby: [DEPRECATED] Will be ignored. If `filename_date_pattern` or
                `filename_band_pattern` are provided, those will be used as group keys)
                for a time dimension and variables respectively. Otherwise, this
                defaults to putting all images onto the same plane.
            filename_date_pattern: Optional date pattern for filenames.
            filename_band_pattern: Optional band patterns for filenames.
                E.g. ``[{"pattern": "*_B[0-9]", "band": "band_1"}]`` would map all
                files matching the pattern ``*_B[0-9]`` to the band name ``band_1``.
                Uses Unix filename pattern rules (fnmatch).

                If specified, all files must match to some pattern, otherwise an error will be raised.
            visualization_optimization: Whether to optimize for visualization.
                Use `auto` for automatic optimization based on size, True to force,
                or False to disable. Defaults to `auto`.
            pixel_info_optimizations: List of dimensions to optimize for pixel info API.
                Defaults to None.

        Returns:
            The dataset response.
        """
        tags = tags or {}
        # Convert filename_band_pattern to the expected format
        formatted_band_patterns = None
        if filename_band_pattern:
            formatted_band_patterns = [
                FilenameBandPattern(pattern=p["pattern"], band=p["band"])
                for p in filename_band_pattern
            ]

        urls = url if isinstance(url, list) else [url]
        request = AddImageDatasetRequest(
            name=name,
            urls=urls,
            labels=labels,
            tags=tags,
            bands=bands,
            groupby=groupby,
            filename_date_pattern=filename_date_pattern,
            filename_band_pattern=formatted_band_patterns,
            visualization_optimization=visualization_optimization,
            pixel_info_optimizations=pixel_info_optimizations or [],
        )
        return self._add_dataset(request)

    def add_zarr_dataset(
        self,
        name: str,
        url: str,
        labels: list[DatasetLabel] | None = None,
        tags: dict[str, str] | None = None,
        rename: dict[str, str] | None = None,
        visualization_optimization: bool | Literal["auto"] = "auto",
        pixel_info_optimizations: list[str] | None = None,
    ) -> AddDatasetResponse:
        """Add a Zarr dataset.

        When loading into xarray, this dataset type will automatically standardize the
        dimensions of the dataset to `y`, `x` and `time` if present. It will infer
        spatial dimensions, so if `lon` or `longitude` is present, it will be renamed
        to `x`.

        This supports arbitrary multi-dimensional datasets, for example a dataset with
        `time` or `level` dimensions in addition to `y`, `x`.

        Args:
            name: The name of the dataset. Creates a new version of an existing dataset
                  if the latest version of a dataset has the same name.
            url: The URL of the dataset. Can optionally contain a placeholder for the
                dimension name. If specified, this concatenates multiple Zarrs along
                either an existing or new dimension as named in the pattern. Example:
                ``gs://mybucket/my_dataset/{time}.zarr``
            labels: Optional. User-defined labels as key-value pairs. Deprecated, use
                `tags` instead. If only labels are provided, a warning will be logged
                and the labels will be converted to tags. If both `labels` and `tags`
                are provided, `tags` will take precedence.
            tags: Optional. User-defined tags as key-value pairs.
            rename: Optional. Dictionary to rename dimensions.
            visualization_optimization: Whether to optimize for visualization.
                Use `auto` for automatic optimization based on size, True to force,
                or False to disable. Defaults to `auto`.
            pixel_info_optimizations: List of dimensions to optimize for pixel info API.
                Defaults to None.

        Returns:
            The dataset response.
        """
        tags = tags or {}
        request = AddZarrDatasetRequest(
            name=name,
            urls=[url] if isinstance(url, str) else url,
            labels=labels,
            tags=tags,
            rename=rename,
            visualization_optimization=visualization_optimization,
            pixel_info_optimizations=pixel_info_optimizations or [],
        )
        return self._add_dataset(request)

    def add_vector_dataset(
        self,
        name: str,
        url: str,
        labels: list[DatasetLabel] | None = None,
        tags: dict[str, str] | None = None,
        coordinate_precision: VectorCoordinatePrecision = "auto",
    ) -> AddDatasetResponse:
        """Add a vector dataset.

        This function supports adding vector datasets from a variety of sources,
        including GeoJSON, GeoParquet, FlatGeobuf, and more.

        Args:
            name: The name of the dataset. Creates a new version of an existing dataset
                  if the latest version of a dataset has the same name.
            url: The URL of the dataset.
            labels: Optional. User-defined labels as key-value pairs. Deprecated, use
                `tags` instead. If only labels are provided, a warning will be logged
                and the labels will be converted to tags. If both `labels` and `tags`
                are provided, `tags` will take precedence.
            tags: Optional. User-defined tags as key-value pairs.
            coordinate_precision: Controls the approximate precision of
                coordinates in the optimized dataset. 'auto' tries to detect
                an appropriate precision. Other values specify the desired
                coordinate precision explicitly. Defaults to 'auto'.

        Returns:
            The dataset response.
        """
        tags = tags or {}
        request = AddVectorDatasetRequest(
            name=name,
            url=url,
            labels=labels,
            tags=tags,
            coordinate_precision=coordinate_precision,
        )
        return self._add_dataset(request)

    def add_tile_server_dataset(
        self,
        name: str,
        url: str,
        labels: list[DatasetLabel] | None = None,
        tags: dict[str, str] | None = None,
    ) -> AddDatasetResponse:
        """Add a tile server dataset.

        The URL must be a template string with placeholders for the x, y, and z
        coordinates, e.g. ``https://server.com/tiles/{z}/{x}/{y}.png``.

        Args:
            name: The name of the dataset.
            url: The URL of the dataset.
            labels: Optional. User-defined labels as key-value pairs. Deprecated, use
                `tags` instead. If only labels are provided, a warning will be logged
                and the labels will be converted to tags. If both `labels` and `tags`
                are provided, `tags` will take precedence.
            tags: Optional. User-defined tags as key-value pairs.

        Returns:
            The dataset response.
        """
        tags = tags or {}
        request = AddTileServerDatasetRequest(
            name=name,
            url=url,
            labels=labels,
            tags=tags,
        )
        return self._add_dataset(request)

    def list_datasets(
        self,
        *,
        name: str | None = None,
        tags: dict[str, str] | None = None,
        bbox: tuple[float, float, float, float] | None = None,
        created_after: datetime.datetime | None = None,
        created_before: datetime.datetime | None = None,
        limit: int | None = None,
    ) -> list[ListDatasetResponse]:
        """List datasets with optional filtering.

        Returns up to ``limit`` datasets in a single request. For iterating
        over all matching datasets use :meth:`iter_datasets` instead.

        All filters are combined with AND logic. When multiple tags are
        given, only datasets matching **all** of them are returned.

        Args:
            name: Case-insensitive substring search on dataset name.
                For example, ``name="Zarr"`` matches "Zarr 1 Band" and
                "HLS Zarr".
            tags: Filter by tags (AND logic). Only datasets that have all
                specified key-value pairs are returned.
                Example: ``{"source": "satellite", "resolution": "10m"}``.
            bbox: Bounding box filter as
                ``(min_lon, min_lat, max_lon, max_lat)`` in EPSG:4326.
                Only datasets whose extent intersects the box are returned.
                For a point query, use the same value for min and max.
            created_after: Only return datasets created after this timestamp.
            created_before: Only return datasets created before this timestamp.
            limit: Maximum number of results (default 100, max 1000).

        Returns:
            A list of datasets matching the filters.

        Example::

            datasets = client.list_datasets(
                name="Zarr",
                tags={"source": "satellite"},
                bbox=(-60, -6, -59, -4),
            )
        """
        if limit is not None and (limit < 1 or limit > 1000):
            raise ValueError("limit must be between 1 and 1000")
        params = self._build_list_params(
            name=name,
            tags=tags,
            bbox=bbox,
            created_after=created_after,
            created_before=created_before,
            limit=limit,
        )
        raw_response = self._http_client.request(
            method="GET",
            endpoint="datasets",
            api_version=self.API_VERSION,
            params=params,
        )
        validated_data = TypeAdapter(list[ListDatasetResponse]).validate_json(
            raw_response
        )
        return PydanticList(validated_data)

    def iter_datasets(
        self,
        *,
        name: str | None = None,
        tags: dict[str, str] | None = None,
        bbox: tuple[float, float, float, float] | None = None,
        created_after: datetime.datetime | None = None,
        created_before: datetime.datetime | None = None,
        page_size: int = 100,
    ) -> Iterator[ListDatasetResponse]:
        """Iterate over all matching datasets with automatic pagination.

        Yields datasets one at a time, transparently fetching successive
        pages using cursor-based pagination. This is the recommended way
        to retrieve large or unbounded result sets.

        All filters are combined with AND logic. When multiple tags are
        given, only datasets matching **all** of them are returned.

        Args:
            name: Case-insensitive substring search on dataset name.
                For example, ``name="Zarr"`` matches "Zarr 1 Band" and
                "HLS Zarr".
            tags: Filter by tags (AND logic). Only datasets that have all
                specified key-value pairs are returned.
                Example: ``{"source": "satellite", "resolution": "10m"}``.
            bbox: Bounding box filter as
                ``(min_lon, min_lat, max_lon, max_lat)`` in EPSG:4326.
                Only datasets whose extent intersects the box are returned.
                For a point query, use the same value for min and max.
            created_after: Only return datasets created after this timestamp.
            created_before: Only return datasets created before this timestamp.
            page_size: Number of datasets to fetch per request (max 1000).

        Yields:
            :class:`ListDatasetResponse` for each matching dataset.

        Example::

            for ds in client.iter_datasets(
                tags={"source": "satellite"},
                bbox=(-60, -6, -59, -4),
            ):
                print(ds.name, ds.tags)
        """
        if page_size < 1 or page_size > 1000:
            raise ValueError("page_size must be between 1 and 1000")
        cursor: str | None = None
        while True:
            params = self._build_list_params(
                name=name,
                tags=tags,
                bbox=bbox,
                created_after=created_after,
                created_before=created_before,
                limit=page_size,
                cursor=cursor,
            )
            raw_response = self._http_client.request(
                method="GET",
                endpoint="datasets/cursor",
                api_version=self.API_VERSION,
                params=params,
            )
            page = PaginatedListDatasetResponse.model_validate_json(raw_response)
            yield from page.data
            cursor = page.next_cursor
            if cursor is None:
                break

    @staticmethod
    def _escape_tag_part(s: str) -> str:
        return s.replace("\\", "\\\\").replace(":", "\\:")

    @staticmethod
    def _build_list_params(
        *,
        name: str | None = None,
        tags: dict[str, str] | None = None,
        bbox: tuple[float, float, float, float] | None = None,
        created_after: datetime.datetime | None = None,
        created_before: datetime.datetime | None = None,
        limit: int | None = None,
        cursor: str | None = None,
    ) -> dict[str, Any]:
        params: dict[str, Any] = {}
        if name is not None:
            params["name"] = name
        if tags is not None:
            esc = EarthscaleClient._escape_tag_part
            params["tag"] = [f"{esc(k)}:{esc(v)}" for k, v in tags.items()]
        if bbox is not None:
            params["bbox"] = ",".join(f"{v:.6f}" for v in bbox)
        if created_after is not None:
            params["created_after"] = created_after.isoformat()
        if created_before is not None:
            params["created_before"] = created_before.isoformat()
        if limit is not None:
            params["limit"] = limit
        if cursor is not None:
            params["cursor"] = cursor
        return params

    def get_dataset(self, dataset_id: str | uuid.UUID) -> DatasetResponse:
        """Get the latest version of a dataset by dataset ID.

        Args:
            dataset_id: The dataset ID.

        Returns:
            The dataset response.
        """
        raw_response = self._http_client.request(
            method="GET",
            endpoint=f"datasets/{dataset_id}/latest",
            api_version=self.API_VERSION,
        )
        return DatasetResponse.model_validate_json(raw_response)

    def get_dataset_version_by_id(self, version_id: str | uuid.UUID) -> DatasetResponse:
        """Get a dataset version by version ID.

        Args:
            version_id: The version ID.

        Returns:
            The dataset response.
        """
        if not isinstance(version_id, uuid.UUID):
            version_id = uuid.UUID(version_id)

        raw_response = self._http_client.request(
            method="GET",
            endpoint=f"dataset-versions/{version_id}",
            api_version=self.API_VERSION,
        )
        return DatasetResponse.model_validate_json(raw_response)

    def delete_dataset(
        self,
        dataset_id: str | uuid.UUID,
        remove_from_maps: bool = True,
    ) -> DeleteDatasetResponse:
        """Soft-delete a dataset.

        Deletion works by creating a new dataset version that marks the
        dataset as deleted. This hides the dataset from the catalog and
        listing APIs, but tiling requests for existing versions are still
        accepted. Optionally removes the dataset from all maps.

        Both dataset IDs and dataset version IDs are accepted — it does
        not matter which is passed, the entire dataset is deleted either
        way.

        Args:
            dataset_id: The dataset ID or dataset version ID. Either
                can be used to identify the dataset — the entire
                dataset is always deleted, not just a single version.
            remove_from_maps: Whether to remove the dataset from all
                maps. Defaults to True.

        Returns:
            Response containing the dataset ID and the ID of the newly
            created deletion marker version.
        """
        raw_response = self._http_client.request(
            method="DELETE",
            endpoint=f"datasets/{dataset_id}",
            api_version=self.API_VERSION,
            params={"remove_from_maps": str(remove_from_maps).lower()},
        )
        return DeleteDatasetResponse.model_validate_json(raw_response)

    def check_api_support(self) -> VersionCheckResponse:
        """Check if the client's API version is compatible with the server.

        This method contacts the server to verify that the API version used by the
        client is supported by the server. It also provides information about the
        supported API versions and whether the current version is deprecated.

        Returns:
            The version check response from the server.

        Raises:
            VersionIncompatibleError: If the API version is not supported by the server.
        """
        raw_response = self._http_client.request(
            method="GET",
            endpoint=f"api-versions/{self.API_VERSION}",
            api_version="",
            use_auth=False,
        )
        response = VersionCheckResponse.model_validate_json(raw_response)

        if not response.is_supported:
            raise VersionIncompatibleError(
                f"API version {self.API_VERSION} is not supported. "
                f"Please upgrade to version {response.newest_supported_version}."
            )

        # Log a warning if the API version is deprecated
        if response.is_deprecated:
            removal_info = ""
            if response.will_be_removed_after:
                removal_info = (
                    f" It will be removed after {response.will_be_removed_after}."
                )

            _logger.warning(
                "API version '%s' is deprecated.%s Please upgrade to version '%s'.",
                self.API_VERSION,
                removal_info,
                response.newest_supported_version,
            )

        return response
