import ee
from te_schemas.results import Band, DataType
from te_schemas.schemas import BandInfo

from te_algorithms.gee.util import teimage_v1_to_teimage_v2

from .util import GEEImage, TEImage, TEImageV2


def _download_default(
    asset,
    name,
    temporal_resolution,
    year_initial,
    year_final,
    band_number=None,
    band_name=None,
    band_metadata=None,
    band_add_to_map=None,
):
    """
    Default function used to download data if no other function is provided for an asset
    """
    in_img = ee.Image(asset)

    info = in_img.getInfo()
    image_properties = info.get("properties", {})
    image_bands = info.get("bands", [])

    if band_number is not None:
        selected_band_number = int(band_number)
        if selected_band_number < 1 or selected_band_number > len(image_bands):
            raise ValueError(
                f"Requested band_number {selected_band_number} is out of range "
                f"for asset {asset} ({len(image_bands)} bands)."
            )

        band_index = selected_band_number - 1
        out = in_img.select([band_index])
        selected_band_metadata = image_bands[band_index]
        metadata = image_properties.copy()
        if isinstance(band_metadata, dict):
            metadata.update(band_metadata)
        metadata.update(
            {
                "band_number": selected_band_number,
                "source_band_id": selected_band_metadata.get("id", ""),
            }
        )

        resolved_band_name = band_name or name

        return teimage_v1_to_teimage_v2(
            TEImage(
                out,
                [
                    BandInfo(
                        resolved_band_name,
                        add_to_map=True,
                        metadata=metadata,
                    )
                ],
            )
        )

    out = in_img
    n_bands = len(image_bands)

    if n_bands == 1:
        metadata = image_properties.copy()
        if isinstance(band_metadata, dict):
            metadata.update(band_metadata)
        resolved_band_name = band_name or name
        add_to_map = band_add_to_map if band_add_to_map is not None else True
        band_info = [
            BandInfo(resolved_band_name, add_to_map=add_to_map, metadata=metadata)
        ]
    else:
        band_info = [BandInfo(name, add_to_map=True, metadata=image_properties)]
        if n_bands > 1:
            band_info.extend(
                [BandInfo(name, add_to_map=True, metadata=image_properties)]
                * (n_bands - 1)
            )

    return teimage_v1_to_teimage_v2(TEImage(out, band_info))


def _download_worldpop(asset, name, temporal_resolution, year_initial, year_final):
    """Download WorldPop data"""
    if year_initial is None:
        year_initial = 2000
    if year_final is None:
        year_final = 2020
    out = TEImageV2(
        {
            DataType.FLOAT32: GEEImage(
                **_get_population(year_initial, asset, add_to_map=True)
            )
        }
    )
    for year in range(year_initial + 1, year_final + 1):
        # Be inclusive of final year (+1) above, and recognize that initial
        # year was already added
        out.add_image(**_get_population(year, asset, add_to_map=True))

    return out


def _get_population(year, asset, add_to_map=False):
    """Return WorldPop population data for a given year"""
    wp = ee.ImageCollection(asset).filterDate(f"{year}-01-01", f"{int(year) + 1}-01-01")

    wp = (
        wp.select("male")
        .toBands()
        .rename(f"Population_{year}_male")
        .addBands(wp.select("female").toBands().rename(f"Population_{year}_female"))
    )

    return {
        "image": wp,
        "bands": [
            Band(
                "Population (number of people)",
                metadata={
                    "year": year,
                    "type": "male",
                },
                add_to_map=add_to_map,
            ),
            Band(
                "Population (number of people)",
                metadata={
                    "year": year,
                    "type": "female",
                },
                add_to_map=add_to_map,
            ),
        ],
        "datatype": DataType.FLOAT32,
    }


def _download_gpw(asset, name, temporal_resolution, year_initial, year_final):
    """Download GPW (Gridded Population of the World) data.

    GPWv411 is an ImageCollection with images for years 2000, 2005, 2010, 2015, 2020.
    """
    # GPW has data for specific years only
    gpw_years = [2000, 2005, 2010, 2015, 2020]

    if year_initial is None:
        year_initial = gpw_years[0]
    if year_final is None:
        year_final = gpw_years[-1]

    # Filter to years within requested range
    years_to_download = [y for y in gpw_years if year_initial <= y <= year_final]

    if not years_to_download:
        # If no exact matches, use the closest available year
        years_to_download = [min(gpw_years, key=lambda y: abs(y - year_initial))]

    out = TEImageV2(
        {
            DataType.FLOAT32: GEEImage(
                **_get_gpw_population(years_to_download[0], asset, add_to_map=True)
            )
        }
    )

    for year in years_to_download[1:]:
        out.add_image(**_get_gpw_population(year, asset, add_to_map=True))

    return out


def _get_gpw_population(year, asset, add_to_map=False):
    """Return GPW population data for a given year."""
    gpw = ee.ImageCollection(asset).filterDate(
        f"{year}-01-01", f"{int(year) + 1}-01-01"
    )

    # GPW has a 'population_count' band - mosaic in case of multiple images
    gpw_img = gpw.select("population_count").mosaic().rename(f"Population_{year}")

    return {
        "image": gpw_img,
        "bands": [
            Band(
                "Population (number of people)",
                metadata={
                    "year": year,
                },
                add_to_map=add_to_map,
            ),
        ],
        "datatype": DataType.FLOAT32,
    }


def _download_spi_collection(
    asset, name, temporal_resolution, year_initial, year_final
):
    """Download SPI/SPEI data from a monthly ImageCollection (CHIRPS, GPCC, UKCEH).

    These datasets are stored as ImageCollections of monthly images, each
    containing bands for different accumulation lag periods (e.g. lag_9, lag_12).
    For each year in the requested range the December image is downloaded
    with all available lag bands included in the output.
    """
    if year_initial is None:
        year_initial = 1981
    if year_final is None:
        year_final = 2024

    collection = ee.ImageCollection(asset)
    proj = collection.first().projection()

    is_spei = "spei" in asset.lower()
    is_ukceh = "ukceh" in asset.lower()
    band_label = (
        "Standardized Precipitation Evapotranspiration Index (SPEI)"
        if is_spei
        else "Standardized Precipitation Index (SPI)"
    )

    # Discover available lag bands from the first image in the collection.
    # This makes one synchronous API call upfront (comparable to the getInfo()
    # call already made in _download_default).
    first_image_info = collection.first().getInfo()
    all_band_ids = [b["id"] for b in first_image_info.get("bands", [])]
    lag_bands = [b for b in all_band_ids if b.startswith("lag_")]

    def _get_december_image(year):
        if is_ukceh:
            # UKCEH images are indexed by system:index rather than date.
            prefix = "spei" if is_spei else "spi"
            lag = int(lag_bands[0].split("_")[1])
            return collection.filter(
                ee.Filter.equals("system:index", f"{prefix}{lag}_dec_{year}")
            ).first()
        else:
            # CHIRPS and GPCC are filtered by acquisition date.
            return collection.filterDate(f"{year}-12-01", f"{year + 1}-01-01").first()

    def _get_spi_year(year, add_to_map=False):
        year_img = (
            _get_december_image(year)
            .select(lag_bands)
            .unmask(-32768)
            .int16()
            .reproject(crs=proj)
        )
        bands = [
            Band(
                band_label,
                metadata={"year": year, "lag": int(b.split("_")[1])},
                add_to_map=add_to_map,
            )
            for b in lag_bands
        ]
        return {"image": year_img, "bands": bands, "datatype": DataType.INT16}

    first_output = _get_spi_year(year_initial, add_to_map=True)
    out = TEImageV2({DataType.INT16: GEEImage(**first_output)})

    for year in range(year_initial + 1, year_final + 1):
        add_to_map = bool(((year - year_initial) % 4) == 0)
        out.add_image(**_get_spi_year(year, add_to_map=add_to_map))

    return out


download_functions = {
    "Gridded Population Count (gender breakdown)": _download_worldpop,
    "Gridded Population of the World v411": _download_gpw,
    "CHIRPS v2025": _download_spi_collection,
    "GPCC V6 (Global Precipitation Climatology Centre)": _download_spi_collection,
    "UKCEH SPI": _download_spi_collection,
    "UKCEH SPEI": _download_spi_collection,
}


def download(
    asset,
    name,
    temporal_resolution,
    year_initial,
    year_final,
    logger,
    band_number=None,
    band_name=None,
    band_metadata=None,
    band_add_to_map=None,
):
    """
    Download dataset from GEE assets.
    """
    logger.debug("Entering download function.")

    # if temporal_resolution != "one time":
    #     assert (year_initial and year_final), "start year or end year not defined"
    #     out = in_img.select('y{}'.format(year_initial))
    #     band_info = [BandInfo(name, add_to_map=True, metadata={'year': year_initial})]
    #     for y in range(year_initial + 1, year_final + 1):
    #         out.addBands(in_img.select('y{}'.format(year_initial)))
    #         band_info.append(BandInfo(name, metadata={'year': year_initial}))
    # else:
    #     out = in_img
    #     band_info = [BandInfo(name, add_to_map=True)]
    if name in download_functions:
        return download_functions[name](
            asset, name, temporal_resolution, year_initial, year_final
        )
    else:
        return _download_default(
            asset,
            name,
            temporal_resolution,
            year_initial,
            year_final,
            band_number=band_number,
            band_name=band_name,
            band_metadata=band_metadata,
            band_add_to_map=band_add_to_map,
        )
