import hashlib
import json
import logging
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import date, timedelta
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 / ".sentinel_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 make_add_index(ndvi: bool, ndwi: bool):
    def add_index(img):
        bands = []
        if ndvi:
            bands.append(img.normalizedDifference(['B8', 'B4']).rename('NDVI'))
        if ndwi:
            bands.append(img.normalizedDifference(['B3', 'B8']).rename('NDWI'))
        return img.addBands(bands)
    return add_index


def generate_point_key(row: pd.Series, ndvi: bool, ndwi: bool) -> str:
    normalized = {
        "lat": float(row["latitude"]),
        "lon": float(row["longitude"]),
        "date": str(row["eventDate"]),
        "ndvi": ndvi,
        "ndwi": ndwi,
    }
    return hashlib.md5(json.dumps(normalized, sort_keys=True).encode()).hexdigest()


def get_index_from_sentinel_batch(
    batch: gpd.GeoDataFrame,
    ndvi: bool,
    ndwi: bool,
) -> gpd.GeoDataFrame:
    batch = batch.copy()
    logger.info(f"Processing batch of {len(batch)} records")

    if pd.Timestamp(batch['eventDate'].max()) < pd.Timestamp('2019-01-01'):
        start_date = date.strftime(date.today() - timedelta(30), '%Y-%m-%d')
        end_date = date.strftime(date.today(), '%Y-%m-%d')
    else:
        start_date = batch['eventDate'].min()
        end_date = batch['eventDate'].max()

    index = [name for name, flag in [('NDVI', ndvi), ('NDWI', ndwi)] if flag]

    img = (
        ee.ImageCollection('COPERNICUS/S2_SR_HARMONIZED')
        .filterDate(start_date, end_date)
        .map(make_add_index(ndvi, ndwi))
        .select(index)
        .median()
    )

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

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

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

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

    return gdf_out


def get_sentinel_properties(
    gdf: gpd.GeoDataFrame,
    batch_size: int = 150,
    ndvi: bool = True,
    ndwi: bool = True,
) -> pd.DataFrame:
    if not ndvi and not ndwi:
        raise ValueError("At least one of NDVI or NDWI must be True.")

    gdf = gdf.sort_values("eventDate", ascending=True)

    keys = gdf.apply(lambda row: generate_point_key(row, ndvi, ndwi), axis=1)
    mask = keys.map(lambda k: k in cache)

    cached_rows = [cache[k] for k in keys[mask]]
    missing_gdf = gpd.GeoDataFrame(gdf[~mask], geometry=gdf.geometry.name, crs=gdf.crs)

    logger.info(f"Cache hit: {len(cached_rows)} | Missing: {len(missing_gdf)}")

    new_results: list[gpd.GeoDataFrame] = []

    if not missing_gdf.empty:
        batches = list(chunk_gdf(missing_gdf, batch_size))

        with ThreadPoolExecutor(max_workers=4) as executor:
            futures = {
                executor.submit(get_index_from_sentinel_batch, batch, ndvi, ndwi): i
                for i, batch in enumerate(batches)
            }
            for future in as_completed(futures):
                try:
                    new_results.append(future.result())
                except Exception as e:
                    logger.error(f"Batch {futures[future]} failed: {e}")

        if new_results:
            new_df = pd.concat(new_results, ignore_index=True)
            for _, row in new_df.iterrows():
                cache.set(generate_point_key(row, ndvi, ndwi), row)

    parts = []
    if cached_rows:
        parts.append(pd.DataFrame(cached_rows))
    if new_results:
        parts.append(pd.concat(new_results, ignore_index=True))

    return pd.concat(parts, ignore_index=True) if parts else pd.DataFrame()