# 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 date, datetime, timezone
from typing import cast

import datacube
import odc.geo
import sqlalchemy.exc
from odc.geo.crs import CRS
from odc.geo.geom import Geometry
from psycopg.types.json import Json as Json3
from psycopg2.extras import Json as Json2
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),
    )

    conn = get_sqlconn(layer.dc)
    txn = conn.begin()
    if meta in cache:
        template = cache[meta][0]
        # A layer has same signature - reuse range"
        cache[meta].append(layer.name)
        try:
            conn.execute(
                text("""
            INSERT INTO ows.layer_ranges
                (layer, lat_min, lat_max, lon_min, lon_max, dates, bboxes, meta, last_updated)
            SELECT :layer_id, lat_min, lat_max, lon_min, lon_max, dates, bboxes, meta, last_updated
            FROM ows.layer_ranges lr2
            WHERE lr2.layer = :template_id"""),
                {"layer_id": layer.name, "template_id": template},
            )
        except sqlalchemy.exc.IntegrityError:
            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:
        Json = Json3 if get_driver_name(layer.dc.index) == "psycopg" else Json2
        # insert empty row if one does not already exist
        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=timezone.utc),
            },
        )

        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(timezone.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.index, 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},
        )

        cache[meta] = [layer.name]

    txn.commit()
    conn.close()


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:
    conn = get_sqlconn(layer.dc)
    results = conn.execute(
        text("""
        SELECT *
        FROM ows.layer_ranges
        WHERE layer=:pname"""),
        {"pname": layer.name},
    )

    for result in results:
        conn.close()
        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
