# This file is part of datacube-ows, part of the Open Data Cube project.
# See https://opendatacube.org for more information.
#
# Copyright (c) 2017-2024 OWS Contributors
# SPDX-License-Identifier: Apache-2.0

import logging
import math
from collections.abc import Callable
from datetime import UTC, date, datetime
from typing import cast

import datacube
import odc.geo
from odc.geo.crs import CRS
from odc.geo.geom import Geometry
from sqlalchemy import text

from datacube_ows.index.api import CoordRange, LayerExtent, LayerSignature
from datacube_ows.index.postgres.mv_index import MVSelectOpts, mv_search
from datacube_ows.ows_configuration import OWSNamedLayer
from datacube_ows.utils import get_driver_name, get_sqlconn

_LOG: logging.Logger = logging.getLogger(__name__)


def jsonise_bbox(bbox: odc.geo.geom.BoundingBox) -> dict[str, float]:
    if isinstance(bbox, dict):
        return bbox
    return {
        "top": bbox.top,
        "bottom": bbox.bottom,
        "left": bbox.left,
        "right": bbox.right,
    }


def create_range_entry(
    layer: OWSNamedLayer, cache: dict[LayerSignature, list[str]]
) -> None:
    # Take layer signature
    meta = LayerSignature(
        time_res=layer.time_resolution.value,
        products=tuple(layer.product_names),
        env=layer.local_env._name,
        datasets=layer.dc.index.datasets.count(product=layer.product_names),
    )

    with get_sqlconn(layer.dc) as conn:
        txn = conn.begin()
        cached = cache.get(meta)
        if get_driver_name(layer.dc.index) == "psycopg":
            from psycopg.types.json import Json
        else:
            from psycopg2.extras import Json
        # insert initial empty row if one does not already exist, so we can safely update it later.
        conn.execute(
            text("""
                 INSERT INTO ows.layer_ranges
                 (layer, lat_min, lat_max, lon_min, lon_max, dates, bboxes, meta, last_updated)
                 VALUES (:p_layer, 0, 0, 0, 0, :empty, :empty, :meta, :now) ON CONFLICT (layer) DO NOTHING
                 """),
            {
                "p_layer": layer.name,
                "empty": Json(""),
                "meta": Json(meta.as_json()),
                "now": datetime.now(tz=UTC),
            },
        )
        if cached is not None:
            template = cached[0]
            # A layer has same signature - reuse range"
            conn.execute(
                text("""
            UPDATE ows.layer_ranges lr1
            SET lat_min = lr2.lat_min,
                lat_max = lr2.lat_max,
                lon_min = lr2.lon_min,
                lon_max = lr2.lon_max,
                dates = lr2.dates,
                bboxes = lr2.bboxes,
                meta = lr2.meta,
                last_updated = lr2.last_updated
            FROM ows.layer_ranges lr2
            WHERE lr1.layer = :layer_id
            AND   lr2.layer = :template_id"""),
                {"layer_id": layer.name, "template_id": template},
            )
        else:
            prodids = [p.id for p in layer.products]
            # Update min/max lat/longs
            conn.execute(
                text(
                    """
              UPDATE ows.layer_ranges lr
              SET lat_min = st_ymin(subq.bbox),
                  lat_max = st_ymax(subq.bbox),
                  lon_min = st_xmin(subq.bbox),
                  lon_max = st_xmax(subq.bbox)
              FROM (
                SELECT st_extent(stv.spatial_extent) as bbox
                FROM ows.space_time_view stv
                WHERE stv.dataset_type_ref = ANY(:prodids)
              ) as subq
              WHERE lr.layer = :layer_id
              """
                ),
                {"layer_id": layer.name, "prodids": prodids},
            )

            # Set default timezone
            conn.execute(text("""set timezone to 'Etc/UTC'"""))

            # Loop over dates
            dates = set()  # Should get to here!
            if layer.time_resolution.is_solar():
                results = conn.execute(
                    text(
                        """
                  select
                        lower(temporal_extent), upper(temporal_extent),
                        ST_X(ST_Centroid(spatial_extent))
                  from ows.space_time_view
                  WHERE dataset_type_ref = ANY(:prodids)
                  """
                    ),
                    {"prodids": prodids},
                )
                for result in results:
                    dt1, dt2, lon = result
                    dt = dt1 + (dt2 - dt1) / 2
                    dt = dt.astimezone(UTC)

                    solar_day = datacube.api.query._convert_to_solar_time(
                        dt, lon
                    ).date()
                    dates.add(solar_day)
            else:
                results = conn.execute(
                    text(
                        """
                  select
                        array_agg(temporal_extent)
                  from ows.space_time_view
                  WHERE dataset_type_ref = ANY(:prodids)
                  """
                    ),
                    {"prodids": prodids},
                )
                for result in results:
                    for dat_ran in result[0]:
                        dates.add(dat_ran.lower)

            if layer.time_resolution.is_subday():
                date_formatter = lambda d: d.isoformat()  # noqa: E731
            else:
                date_formatter = lambda d: d.strftime("%Y-%m-%d")  # noqa: E731

            dates = sorted(dates)
            conn.execute(
                text("""
               UPDATE ows.layer_ranges
               SET dates = :dates
               WHERE layer= :layer_id
            """),
                {
                    "dates": Json(list(map(date_formatter, dates))),
                    "layer_id": layer.name,
                },
            )
            # calculate bounding boxes
            # Get extent polygon from materialised views

            extent_4386 = cast(
                Geometry,
                mv_search(layer.dc, MVSelectOpts.EXTENT, products=layer.products),
            )

            all_bboxes = bbox_projections(extent_4386, layer.global_cfg.crses)

            conn.execute(
                text("""
            UPDATE ows.layer_ranges
            SET bboxes = :bbox
            WHERE layer = :layer_id
            """),
                {"bbox": Json(all_bboxes), "layer_id": layer.name},
            )

        txn.commit()
        # Only update cache after successful commit.
        if cached is None:
            cache[meta] = []
        cache[meta].append(layer.name)


def bbox_projections(
    starting_box: odc.geo.Geometry, crses: dict[str, CRS]
) -> dict[str, dict[str, float]]:
    result = {}
    for crsid, crs in crses.items():
        if crs.valid_region is not None:
            clipped_crs_region = starting_box & crs.valid_region
            if clipped_crs_region.wkt == "POLYGON EMPTY":
                continue
            clipped_crs_bbox = clipped_crs_region.to_crs(crs).boundingbox
        else:
            clipped_crs_bbox = None
        if clipped_crs_bbox is not None:
            result[crsid] = jsonise_bbox(clipped_crs_bbox)
        else:
            projbbox = starting_box.to_crs(crs).boundingbox
            result[crsid] = sanitise_bbox(projbbox)
    return result


def sanitise_bbox(bbox: odc.geo.geom.BoundingBox) -> dict[str, float]:
    def sanitise_coordinate(coord: float, fallback: float) -> float:
        return coord if math.isfinite(coord) else fallback

    return {
        "top": sanitise_coordinate(bbox.top, float("9.999999999e99")),
        "bottom": sanitise_coordinate(bbox.bottom, float("-9.999999999e99")),
        "left": sanitise_coordinate(bbox.left, float("-9.999999999e99")),
        "right": sanitise_coordinate(bbox.right, float("9.999999999e99")),
    }


def get_ranges(layer: OWSNamedLayer) -> LayerExtent | None:
    with get_sqlconn(layer.dc) as conn:
        results = conn.execute(
            text("""
            SELECT *
            FROM ows.layer_ranges
            WHERE layer=:pname"""),
            {"pname": layer.name},
        )

        for result in results:
            if layer.time_resolution.is_subday():
                dt_parser: Callable[[str], datetime | date] = (  # noqa: E731
                    lambda dts: datetime.fromisoformat(dts)
                )
            else:
                dt_parser = lambda dts: datetime.strptime(dts, "%Y-%m-%d").date()  # noqa: E731
            times = [dt_parser(d) for d in result.dates if d is not None]
            if not times:
                return None
            return LayerExtent(
                lat=CoordRange(min=float(result.lat_min), max=float(result.lat_max)),
                lon=CoordRange(min=float(result.lon_min), max=float(result.lon_max)),
                times=times,
                bboxes=result.bboxes,
            )
        return None
