import hashlib
import json
import logging
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Iterator

import ee
import geopandas as gpd
import pandas as pd
from diskcache import Cache

logger = logging.getLogger(__name__)
path_root = Path(__file__).resolve().parent
cache: Cache = Cache(path_root / ".worldclim_cache")


def chunk_gdf(gdf: gpd.GeoDataFrame, size: int = 500) -> Iterator[gpd.GeoDataFrame]:
    for i in range(0, len(gdf), size):
        yield gdf.iloc[i:i + size]


def generate_cache_key(gdf: gpd.GeoDataFrame, index: list[str], scale: int) -> str:
    if isinstance(gdf, pd.DataFrame) and not isinstance(gdf, gpd.GeoDataFrame):
        gdf = gpd.GeoDataFrame(
            gdf,
            geometry=gpd.points_from_xy(gdf.longitude, gdf.latitude),
            crs="EPSG:4326",
        )

    bounds = gdf.total_bounds
    normalized = {
        "lat_min": float(bounds[1]),
        "lat_max": float(bounds[3]),
        "lon_min": float(bounds[0]),
        "lon_max": float(bounds[2]),
        "date_min": str(gdf['eventDate'].min()),
        "date_max": str(gdf['eventDate'].max()),
        "index": index,
        "scale": scale,
        "size": len(gdf),
    }
    return hashlib.md5(json.dumps(normalized, sort_keys=True).encode()).hexdigest()


def add_worldclim_properties(
    gdf: gpd.GeoDataFrame,
    img: ee.Image,
    scale: int,
) -> gpd.GeoDataFrame:
    gdf = gdf.copy()
    logger.info(f"Processing {len(gdf)} records...")

    gdf['eventDate'] = gdf['eventDate'].astype(str)
    fc = ee.FeatureCollection(json.loads(gdf.to_json()))

    result = img.reduceRegions(
        collection=fc,
        reducer=ee.Reducer.mean(),
        scale=scale,
    ).getInfo()

    return gpd.GeoDataFrame.from_features(result["features"], crs="EPSG:4326")


def get_average_temperature_and_precipation(
    gdf: gpd.GeoDataFrame | pd.DataFrame,
    scale: int = 1000,
    batch_size: int = 5000,
    temperature: bool = True,
    precipitation: bool = True,
) -> gpd.GeoDataFrame:
    index = [name for name, flag in [("tavg", temperature), ("prec", precipitation)] if flag]

    if not index:
        raise ValueError("At least temperature or precipitation must be True.")

    cache_key = generate_cache_key(gdf, index, scale)
    if cache_key in cache:
        logger.info("Cache hit — returning stored result")
        return cache[cache_key]

    logger.info("Downloading dataset from GEE...")

    img = (
        ee.ImageCollection("WORLDCLIM/V1/MONTHLY")
        .select(index)
        .mean()
        .multiply(0.1)
    )

    if len(gdf) > batch_size:
        with ThreadPoolExecutor(max_workers=4) as executor:
            futures = {
                executor.submit(add_worldclim_properties, batch, img, scale): i
                for i, batch in enumerate(chunk_gdf(gdf, size=batch_size))
            }
            all_batches: list[gpd.GeoDataFrame] = []
            for future in as_completed(futures):
                try:
                    all_batches.append(future.result())
                except Exception as e:
                    logger.error(f"Batch {futures[future]} failed: {e}")

        final_gdf = gpd.GeoDataFrame(pd.concat(all_batches, ignore_index=True), crs="EPSG:4326")
    else:
        final_gdf = add_worldclim_properties(gdf, img, scale)

    if len(index) == 1:
        final_gdf = final_gdf.rename(columns={"mean": index[0]})

    cache.set(cache_key, final_gdf)
    return final_gdf