# This file is part of the Open Data Cube, see https://opendatacube.org for more information
#
# Copyright (c) 2015-2025 ODC Contributors
# SPDX-License-Identifier: Apache-2.0
import datetime
from collections import namedtuple
from contextlib import contextmanager
from copy import deepcopy
from typing import cast
from uuid import UUID

from datacube.drivers.postgres import PostgresDb
from datacube.index.exceptions import DuplicateRecordError
from datacube.index.postgres._datasets import DatasetResource
from datacube.index.postgres.index import Index
from datacube.model import Dataset, MetadataType, Product

_nbar_uuid = UUID("f2f12372-8366-11e5-817e-1040f381a756")
_ortho_uuid = UUID("5cf41d98-eda9-11e4-8a8e-1040f381a756")
_telemetry_uuid = UUID("4ec8fe97-e8b9-11e4-87ff-1040f381a756")

# An NBAR with source datasets. Many fields have been removed to keep it semi-focused to our ingest test.
_EXAMPLE_NBAR = {
    "id": str(_nbar_uuid),
    "product_type": "nbar_brdf",
    "checksum_path": "package.sha1",
    "ga_label": "LS8_OLITIRS_NBAR_P54_GALPGS01-002_112_079_20140126",
    "ga_level": "P54",
    "size_bytes": 4550,
    "platform": {"code": "LANDSAT_8"},
    "creation_dt": datetime.datetime(2014, 1, 26, 2, 5, 23, 126373).isoformat(),
    "instrument": {"name": "OLI_TIRS"},
    "format": {"name": "GeoTIFF"},
    "extent": {
        "center_dt": datetime.datetime(2014, 1, 26, 2, 5, 23, 126373).isoformat(),
        "coord": {
            "ul": {"lat": -26.37259, "lon": 116.58914},
            "lr": {"lat": -28.48062, "lon": 118.96145},
            "ur": {"lat": -26.36025, "lon": 118.92432},
            "ll": {"lat": -28.49412, "lon": 116.58121},
        },
    },
    "lineage": {
        "machine": {},
        "source_datasets": {
            "ortho": {
                "product_level": "L1T",
                "product_type": "ortho",
                "id": str(_ortho_uuid),
                "usgs": {"scene_id": "LC81120792014026ASA00"},
                "extent": {
                    "center_dt": datetime.datetime(
                        2014, 1, 26, 2, 5, 23, 126373
                    ).isoformat(),
                    "coord": {
                        "ul": {"lat": -26.37259, "lon": 116.58914},
                        "lr": {"lat": -28.48062, "lon": 118.96145},
                        "ur": {"lat": -26.36025, "lon": 118.92432},
                        "ll": {"lat": -28.49412, "lon": 116.58121},
                    },
                },
                "size_bytes": 1854924494,
                "platform": {"code": "LANDSAT_8"},
                "creation_dt": datetime.datetime(2015, 4, 7, 0, 58, 8).isoformat(),
                "instrument": {"name": "OLI_TIRS"},
                "checksum_path": "package.sha1",
                "ga_label": "LS8_OLITIRS_OTH_P51_GALPGS01-002_112_079_20140126",
                "grid_spatial": {
                    "projection": {
                        "spatial_reference": "EPSG:28350",
                        "resampling_option": "CUBIC_CONVOLUTION",
                        "geo_ref_points": {
                            "ul": {"y": 7082987.5, "x": 459012.5},
                            "lr": {"y": 6847987.5, "x": 692012.5},
                            "ur": {"y": 7082987.5, "x": 692012.5},
                            "ll": {"y": 6847987.5, "x": 459012.5},
                        },
                        "orientation": "NORTH_UP",
                    }
                },
                "acquisition": {
                    "groundstation": {
                        "code": "ASA",
                        "eods_domain_code": "002",
                        "label": "Alice Springs",
                    }
                },
                "format": {"name": "GEOTIFF"},
                "lineage": {
                    "algorithm": {"name": "LPGS", "parameters": {}, "version": "2.4.0"},
                    "machine": {},
                    "source_datasets": {
                        "satellite_telemetry_data": {
                            "product_type": "satellite_telemetry_data",
                            "checksum_path": "package.sha1",
                            "id": str(_telemetry_uuid),
                            "ga_label": "LS8_OLITIRS_STD-MD_P00_LC81160740742015089ASA00_"
                            "116_074_20150330T022553Z20150330T022657",
                            "ga_level": "P00",
                            "size_bytes": 637660782,
                            "platform": {"code": "LANDSAT_8"},
                            "creation_dt": datetime.datetime(
                                2015, 4, 22, 6, 32, 4
                            ).isoformat(),
                            "instrument": {"name": "OLI_TIRS"},
                            "format": {"name": "MD"},
                            "lineage": {"source_datasets": {}},
                        }
                    },
                },
            }
        },
    },
}

_EXAMPLE_METADATA_TYPE = MetadataType(
    {
        "name": "eo",
        "dataset": {
            "id": ["id"],
            "label": ["ga_label"],
            "creation_time": ["creation_dt"],
            "measurements": ["image", "bands"],
            "sources": ["lineage", "source_datasets"],
        },
    }
)

_EXAMPLE_PRODUCT = Product(
    _EXAMPLE_METADATA_TYPE,
    {"name": "eo", "description": "", "metadata_type": "eo", "metadata": {}},
)


def _build_dataset(doc):
    sources = {
        name: _build_dataset(src)
        for name, src in doc["lineage"]["source_datasets"].items()
    }
    return Dataset(_EXAMPLE_PRODUCT, doc, uri="file://test.zzz", sources=sources)


_EXAMPLE_NBAR_DATASET = _build_dataset(_EXAMPLE_NBAR)

DatasetRecord = namedtuple(
    "DatasetRecord",
    ["id", "metadata", "dataset_type_ref", "uris", "added", "added_by", "archived"],
)


class MockDb:
    def __init__(self) -> None:
        self.dataset: dict = {}
        self.dataset_source: set = set()

    @contextmanager
    def _connect(self):
        yield self

    def begin(self) -> None:
        pass

    def commit(self) -> None:
        pass

    def rollback(self) -> None:
        pass

    def get_dataset(self, id_):
        return self.dataset.get(id_, None)

    def get_locations(self, dataset):
        return ["file:xxx"]

    def datasets_intersection(self, ids):
        return [k for k in ids if k in self.dataset]

    def contains_dataset(self, dataset_id):
        return dataset_id in self.dataset

    def insert_dataset_location(self, *args, **kwargs) -> None:
        return

    def insert_dataset(self, metadata_doc, dataset_id, dataset_type_id) -> bool:
        # Will we pretend this one was already ingested?
        if dataset_id in self.dataset:
            raise DuplicateRecordError("already ingested")

        self.dataset[dataset_id] = DatasetRecord(
            dataset_id, deepcopy(metadata_doc), dataset_type_id, None, None, None, None
        )
        return True

    def insert_dataset_source(self, classifier, dataset_id, source_dataset_id) -> None:
        self.dataset_source.add((classifier, dataset_id, source_dataset_id))


class MockTypesResource:
    def __init__(self, type_) -> None:
        self.type = type_

    def get(self, *args, **kwargs):
        return self.type

    def get_by_name(self, *args, **kwargs):
        return self.type

    def get_by_name_unsafe(self, *args, **kwargs):
        return self.type

    @contextmanager
    def _db_connection(self, transaction=False):
        yield MockDb()


class MockIndex:
    def __init__(self, db, product) -> None:
        self._db = db
        self.products = MockTypesResource(product)

    def thread_transaction(self):
        return None

    @contextmanager
    def _active_connection(self, transaction=False):
        yield self._db


def test_index_dataset() -> None:
    mock_db = MockDb()
    mock_index = MockIndex(mock_db, _EXAMPLE_PRODUCT)
    datasets = DatasetResource(cast(PostgresDb, mock_db), cast(Index, mock_index))
    datasets.add(_EXAMPLE_NBAR_DATASET)

    ids = {d.id for d in mock_db.dataset.values()}
    assert ids == {_nbar_uuid, _ortho_uuid, _telemetry_uuid}

    # Three datasets (ours and the two embedded source datasets)
    assert len(mock_db.dataset) == 3

    # Our three datasets should be linked together
    # Nbar -> Ortho -> Telemetry
    assert len(mock_db.dataset_source) == 2
    assert mock_db.dataset_source == {
        ("ortho", _nbar_uuid, _ortho_uuid),
        ("satellite_telemetry_data", _ortho_uuid, _telemetry_uuid),
    }

    # Nothing ingested, because we reported the first as already ingested.
    datasets.add(_EXAMPLE_NBAR_DATASET)
    assert len(mock_db.dataset) == 3
    assert len(mock_db.dataset_source) == 2


def test_index_already_ingested_source_dataset() -> None:
    mock_db = MockDb()
    mock_index = MockIndex(mock_db, _EXAMPLE_PRODUCT)
    datasets = DatasetResource(cast(PostgresDb, mock_db), cast(Index, mock_index))
    datasets.add(_EXAMPLE_NBAR_DATASET.sources["ortho"])

    assert len(mock_db.dataset) == 2
    assert len(mock_db.dataset_source) == 1

    datasets.add(_EXAMPLE_NBAR_DATASET)
    assert len(mock_db.dataset) == 3
    assert len(mock_db.dataset_source) == 2


def test_index_two_levels_already_ingested() -> None:
    mock_db = MockDb()
    mock_index = MockIndex(mock_db, _EXAMPLE_PRODUCT)
    datasets = DatasetResource(cast(PostgresDb, mock_db), cast(Index, mock_index))
    datasets.add(
        _EXAMPLE_NBAR_DATASET.sources["ortho"].sources["satellite_telemetry_data"]
    )

    assert len(mock_db.dataset) == 1
    assert len(mock_db.dataset_source) == 0

    datasets.add(_EXAMPLE_NBAR_DATASET)
    assert len(mock_db.dataset) == 3
    assert len(mock_db.dataset_source) == 2
