import datetime
import enum
import logging
import typing
import uuid
from typing import Literal, TypeVar, Any, Union, Annotated

import pydantic_core
from pydantic import (
    AwareDatetime,
    BaseModel,
    Field,
    GetCoreSchemaHandler,
    model_validator,
)
from pydantic_core import CoreSchema, core_schema
from typing_extensions import Self

_logger = logging.getLogger("earthscale")


class PrettyBaseModel(BaseModel):
    """Base model with IPython pretty-printing support"""

    # Skipping type hints to avoid dependency on IPython
    def _repr_pretty_(self, p, cycle):  # type: ignore[no-untyped-def]
        if cycle:
            p.text(f"{self.__class__.__name__}(...)")
        else:
            p.text(pydantic_core.to_json(self, indent=4).decode())


class DatasetType(str, enum.Enum):
    RASTER = "raster"
    VECTOR = "vector"


class VersionCheckResponse(PrettyBaseModel):
    is_supported: bool
    newest_supported_version: int

    is_deprecated: bool = False
    will_be_removed_after: datetime.date | None = None
    deprecation_message: str | None = None


class DatasetLabel(BaseModel):
    """User-defined label for a dataset"""

    name: str
    """
    Name of the label. Can be any string.
    """

    value: str
    """
    Value of the label. Can be any string.
    """


class SimpleDatasetMetadata(BaseModel):
    """Simplified metadata with only essential fields"""

    description: str | None = None
    thumbnail_url: str | None = None
    attributes: list[DatasetLabel] = Field(default_factory=list)
    license: str | None = None


class FilenameBandPattern(BaseModel):
    """Pattern for mapping filename patterns to band names"""

    pattern: str
    band: str | None = None
    band_mapping: dict[int, str] | None = None

    @model_validator(mode="after")
    def _validate_band_or_mapping(self) -> Self:
        if self.band is None and self.band_mapping is None:
            raise ValueError("Either 'band' or 'band_mapping' must be provided")
        if self.band is not None and self.band_mapping is not None:
            raise ValueError("Cannot specify both 'band' and 'band_mapping'")
        if self.band_mapping is not None:
            invalid_indices = [k for k in self.band_mapping.keys() if k < 1]
            if invalid_indices:
                raise ValueError(
                    f"Band indices must be >= 1 (GDAL uses 1-based indexing), "
                    f"got invalid indices: {invalid_indices}"
                )
        return self


def validate_filename_band_patterns_not_mixed(
    patterns: list[FilenameBandPattern],
) -> None:
    if len(patterns) <= 1:
        return
    has_band = any(p.band is not None for p in patterns)
    has_band_mapping = any(p.band_mapping is not None for p in patterns)
    if has_band and has_band_mapping:
        raise ValueError(
            "Cannot mix 'band' and 'band_mapping' patterns in the same list. "
            "All patterns must use either 'band' (for single-band files) or "
            "'band_mapping' (for multi-band files), but not both."
        )


class BaseAddDatasetRequest(BaseModel):
    """Base class for all dataset add requests with shared fields"""

    name: str = Field(
        description="The name of the dataset",
        examples=["my_dataset"],
    )
    labels: list[DatasetLabel] | None = Field(
        description="DEPRECATED: Please use `tags` instead. User-defined labels to add "
        "to the dataset. If both `tags` and `labels` are provided, `tags` will be "
        "used.",
        examples=[{"name": "key", "value": "value"}],
        default=None,
        deprecated=True,
    )
    tags: dict[str, str] = Field(
        description="User-defined tags to add to the dataset. Tags are key-value pairs "
        "that can be used to categorize and filter datasets.",
        examples=[{"key": "value"}],
        default_factory=dict,
    )
    visualization_optimization: bool | Literal["auto"] = Field(
        description=(
            "Whether to optimize the dataset for visualization. If set to 'auto', "
            "the dataset will be optimized if it is small enough."
        ),
        examples=["auto"],
        default="auto",
    )
    pixel_info_optimizations: list[str] = Field(
        description=(
            "List of dimensions to optimize for the pixel info API. This is useful to "
            "quickly retrieve all values of a dimension for a given pixel, e.g. "
            "a time series."
        ),
        examples=[["time"]],
        default_factory=list,
    )

    @model_validator(mode="after")
    def check_labels_and_tags(self) -> "BaseAddDatasetRequest":
        if self.labels:
            _logger.warning(
                "`labels` are provided. `tags` will be joined with "
                "`labels` with `tags` taking precedence. `labels` is deprecated. "
                "Please use `tags` instead. The `labels` field will be removed in a "
                "future release. Please update your code to use `tags` instead."
            )
            for label in self.labels:
                if label.name not in self.tags:
                    self.tags[label.name] = label.value
        return self


class AddImageDatasetRequest(BaseAddDatasetRequest):
    """Request to add an image dataset"""

    type: Literal["image"] = "image"
    urls: list[str] = Field(
        description=(
            "List of URLs or wildcards to the image files. "
            "Wildcards are supported using the ``*`` character. "
            "e.g. ``gs://example-bucket/image_*.tif``"
        ),
        examples=["gs://example-bucket/image_*.tif"],
    )
    bands: list[str] | None = Field(
        description=(
            "List of bands to add to the dataset. If not provided, all bands will be "
            "added."
        ),
        examples=[["band1", "band2"]],
        default=None,
    )
    groupby: str | None = Field(
        description=(
            "[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."
        ),
        examples=["time"],
        default=None,
        deprecated=True,
    )
    filename_date_pattern: str | None = Field(
        description="Pattern to extract the date from the filename.",
        examples=["%Y-%m-%d"],
        default=None,
    )
    filename_band_pattern: list[FilenameBandPattern] | None = Field(
        description="Pattern to extract the band from the filename.",
        examples=[{"pattern": "image_*.tif", "band": "band1"}],
        default=None,
    )

    @model_validator(mode="after")
    def _check_pixel_info_optimizations_not_supported(self) -> "AddImageDatasetRequest":
        if self.pixel_info_optimizations:
            raise ValueError(
                "pixel_info_optimizations is only supported for Zarr datasets"
            )
        return self

    @model_validator(mode="after")
    def _check_filename_band_pattern_not_mixed(self) -> "AddImageDatasetRequest":
        if self.filename_band_pattern is not None:
            validate_filename_band_patterns_not_mixed(self.filename_band_pattern)
        return self

    @model_validator(mode="after")
    def _check_bands_for_filename_band_pattern(
        self,
    ) -> "AddImageDatasetRequest":
        if self.filename_band_pattern is None:
            return self
        if self.bands is None:
            derived_bands: list[str] = []
            for pattern in self.filename_band_pattern:
                if pattern.band is not None:
                    if pattern.band not in derived_bands:
                        derived_bands.append(pattern.band)
                elif pattern.band_mapping is not None:
                    for band_name in pattern.band_mapping.values():
                        if band_name not in derived_bands:
                            derived_bands.append(band_name)
            self.bands = derived_bands
            return self
        bands_set = set(self.bands)
        for pattern in self.filename_band_pattern:
            if pattern.band is not None:
                if pattern.band not in bands_set:
                    raise ValueError(
                        f"Band name '{pattern.band}' from filename_band_pattern "
                        f"is not in the bands list: {self.bands}"
                    )
            elif pattern.band_mapping is not None:
                for band_name in pattern.band_mapping.values():
                    if band_name not in bands_set:
                        raise ValueError(
                            f"Band name '{band_name}' from filename_band_pattern "
                            f"band_mapping is not in the bands list: {self.bands}"
                        )
        return self


class AddZarrDatasetRequest(BaseAddDatasetRequest):
    """Request to add a Zarr dataset"""

    type: Literal["zarr"] = "zarr"
    urls: list[str] = Field(
        description=(
            "List of URLs to the Zarr files. Currently only supports 1 URL. "
            "Can 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."
        ),
        examples=["gs://example-bucket/image.zarr", "gs://example-bucket/{time}.zarr"],
    )
    rename: dict[str, str] | None = None

    @model_validator(mode="after")
    def check_only_one_url_for_now(self) -> "AddZarrDatasetRequest":
        if len(self.urls) > 1:
            raise ValueError("Only one URL is supported for now")
        return self


VectorCoordinatePrecision = Literal[
    "auto",
    "1250m",
    "600m",
    "300m",
    "150m",
    "80m",
    "40m",
    "20m",
    "10m",
    "5m",
    "2m",
    "1m",
    "50cm",
    "25cm",
    "15cm",
    "8cm",
]


class AddVectorDatasetRequest(BaseAddDatasetRequest):
    """Request to add a vector dataset"""

    type: Literal["vector"] = "vector"
    url: str = Field(
        description="URL to the vector dataset.",
        examples=["https://example.com/vector.geojson"],
    )
    coordinate_precision: "VectorCoordinatePrecision" = Field(
        description=(
            "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."
        ),
        default="auto",
    )

    @model_validator(mode="after")
    def _check_pixel_info_optimizations_not_supported(
        self,
    ) -> "AddVectorDatasetRequest":
        if self.pixel_info_optimizations:
            raise ValueError(
                "pixel_info_optimizations is only supported for Zarr datasets"
            )
        return self


class AddTileServerDatasetRequest(BaseAddDatasetRequest):
    type: Literal["tileserver"] = "tileserver"
    url: str = Field(
        description="URL to an XYZ tile server.",
        examples=["https://example.com/tileserver/{z}/{x}/{y}.png"],
    )

    @model_validator(mode="after")
    def _check_pixel_info_optimizations_not_supported(
        self,
    ) -> "AddTileServerDatasetRequest":
        if self.pixel_info_optimizations:
            raise ValueError(
                "pixel_info_optimizations is only supported for Zarr datasets"
            )
        return self


AddDatasetRequest = (
    AddImageDatasetRequest
    | AddZarrDatasetRequest
    | AddVectorDatasetRequest
    | AddTileServerDatasetRequest
)


class AddDatasetResponse(PrettyBaseModel):
    dataset_id: uuid.UUID
    """Dataset ID of the newly created dataset version"""

    dataset_version_id: uuid.UUID
    """Dataset Version ID of the newly created dataset version"""


class DeleteDatasetResponse(PrettyBaseModel):
    dataset_id: uuid.UUID
    """Dataset ID of the deleted dataset"""

    dataset_version_id: uuid.UUID
    """Dataset Version ID of the deletion marker version"""


class BaseAddConnectionRequest(PrettyBaseModel):
    name: str
    """Name of the connection"""
    default_tags: dict[str, str | int | float | bool | datetime.datetime] = Field(
        default_factory=dict
    )
    """Default tags for the connection, applied to all
    files in the connection. Tags are key-value pairs that
    can be used to categorize and filter files."""
    type: Literal["bucket", "stac"]
    """Type of the connection, either `stac` or `bucket`"""
    reread_existing_raster_source_files: bool = False
    """Whether to reread existing raster source files. If false,
    does not read any files that have already been indexed for the
    organization."""


class BucketAddConnectionRequest(BaseAddConnectionRequest):
    type: Literal["bucket"] = "bucket"
    glob_pattern: str
    """Glob pattern to match files in the bucket"""


class STACAddConnectionRequest(BaseAddConnectionRequest):
    type: Literal["stac"] = "stac"
    catalog_url: str
    """URL to the STAC catalog"""
    collection_name: str
    """Name of the collection to add"""


# Union type for all connection request types
AddConnectionRequest = Annotated[
    Union[BucketAddConnectionRequest, STACAddConnectionRequest],
    Field(discriminator="type"),
]


class DatasetOverviewsStatus(BaseModel):
    status: str
    """Status of the dataset overview generation. One of:
    - "not_started": Overview generation has not started yet
    - "pending": Overview generation is pending
    - "running": Overview generation is running
    - "success": Overview generation has completed successfully
    - "error": Overview generation has failed
    """
    updated_at: datetime.datetime | None = None
    """Timestamp of the last update to the status"""


OptimizationStatus = Literal[
    "not_started",
    "pending",
    "running",
    "success",
    "error",
]


class Optimization(BaseModel):
    status: OptimizationStatus
    """Status of the optimization. One of:
    - "not_started": Optimization has not started yet
    - "pending": Optimization is pending
    - "running": Optimization is running
    - "success": Optimization has completed successfully
    - "error": Optimization has failed
    """
    updated_at: datetime.datetime | None = None
    """Timestamp of the last update to the status"""


class TileServer(BaseModel):
    tile_url: str
    """Tileserver URL to be used in an application"""
    pixel_url: str
    """Pixel info URL to be used in an application"""
    min_zoom: int
    """Minimum zoom level supported by the tile server"""
    max_zoom: int
    """Maximum zoom level supported by the tile server"""


class Variable(BaseModel):
    """Estimated statistics for a variable in the dataset, useful for visualization"""

    sampled_min: float
    """Estimated minimum value in the dataset"""
    sampled_max: float
    """Estimated maximum value in the dataset"""


class DatasetResponse(PrettyBaseModel):
    dataset_id: uuid.UUID
    """Dataset ID of the dataset"""
    dataset_version_id: uuid.UUID
    """ID for the current version of the dataset"""
    name: str
    """Name of the dataset"""
    type: DatasetType
    """Type of the dataset, either `raster` or `vector`"""
    labels: list[DatasetLabel] = Field(default_factory=list)
    """Deprecated: User-defined labels for the dataset. Superseded by `tags`."""
    tags: dict[str, str] = Field(default_factory=dict)
    """User-defined tags for the dataset"""
    variables: dict[str, Variable] = Field(default_factory=dict)
    """Estimated statistics for each variable in the dataset, useful for visualization"""  # noqa: E501
    created_at: AwareDatetime
    """Timestamp of when the dataset was created"""
    visualization_optimization: Optimization
    """Status of the visualization optimization for faster visualization"""
    pixel_info_optimizations: dict[str, Optimization] = Field(default_factory=dict)
    """Status of the pixel info optimizations for faster pixel info retrieval"""

    dynamic_tile_server: TileServer
    """Dynamic tile server for fast visualization"""
    optimized_tile_server: TileServer | None
    """Optimized tile server for fast visualization"""


class ListDatasetResponse(PrettyBaseModel):
    """Single entry in the list of datasets with a subset of information"""

    dataset_id: uuid.UUID
    """Dataset ID of the dataset"""
    dataset_version_id: uuid.UUID
    """ID for the current version of the dataset"""
    name: str
    """Name of the dataset"""
    type: DatasetType
    """Type of the dataset, either `raster` or `vector`"""
    labels: list[DatasetLabel] = Field(default_factory=list)
    """Deprecated: User-defined labels for the dataset. Superseded by `tags`."""
    tags: dict[str, str] = Field(default_factory=dict)
    """User-defined tags for the dataset"""
    created_at: AwareDatetime
    """Timestamp of when the dataset was created"""
    updated_at: AwareDatetime
    """Timestamp of when the dataset was last updated"""


class PaginatedListDatasetResponse(PrettyBaseModel):
    """Paginated list of datasets with cursor-based pagination"""

    data: list[ListDatasetResponse]
    """List of datasets"""
    next_cursor: str | None = None
    """Opaque cursor token for fetching the next page"""


class ErrorResponse(PrettyBaseModel):
    message: str
    error_class: str


BaseModelT = TypeVar("BaseModelT", bound=BaseModel)


class PydanticList(list[BaseModelT]):
    """
    Custom list type to enable pretty-printing of our responses in IPython.
    """

    # Skipping type hints to avoid dependency on IPython
    def _repr_pretty_(self, p, cycle):  # type: ignore[no-untyped-def]
        if cycle:
            p.text("PydanticList(...)")
        else:
            p.text(pydantic_core.to_json(self, indent=4).decode())

    @classmethod
    def __get_pydantic_core_schema__(
        cls, source_type: Any, handler: GetCoreSchemaHandler
    ) -> CoreSchema:
        element_type = typing.get_args(source_type)[0]
        return core_schema.list_schema(
            handler(element_type),
        )
