import os
import glob
import yaml
import ee
import geemap
import geopandas as gpd
import rasterio
from shapely.geometry import box

def load_config():
    config_path = os.path.join(os.path.dirname(__file__), 'config.yaml')
    with open(config_path, 'r') as f:
        return yaml.safe_load(f)

def get_s2_image(roi, date_start, date_end):
    collection = ee.ImageCollection('COPERNICUS/S2_SR_HARMONIZED') \
        .filterBounds(roi) \
        .filterDate(date_start, date_end)
    
    def mask_s2_clouds(image):
        qa = image.select('QA60')
        cloudBitMask = 1 << 10
        cirrusBitMask = 1 << 11
        mask = qa.bitwiseAnd(cloudBitMask).eq(0).And(qa.bitwiseAnd(cirrusBitMask).eq(0))
        return image.updateMask(mask).divide(10000)
    
    s2_median = collection.map(mask_s2_clouds).median().clip(roi)
    
    ndvi = s2_median.normalizedDifference(['B8', 'B4']).rename('NDVI')
    ndwi = s2_median.normalizedDifference(['B3', 'B8']).rename('NDWI')
    ndsi = s2_median.normalizedDifference(['B3', 'B11']).rename('NDSI')
    
    return s2_median.select(['B2', 'B3', 'B4', 'B8', 'B11', 'B12']).addBands([ndvi, ndwi, ndsi])

def get_terrain_image(roi):
    def add_terrain(img):
        slope = ee.Terrain.slope(img).rename('slope')
        aspect = ee.Terrain.aspect(img).rename('aspect')
        return img.addBands([slope, aspect])
        
    terrain = ee.ImageCollection('COPERNICUS/DEM/GLO30').select('DEM').map(add_terrain).mosaic().clip(roi)
    elevation = terrain.select('DEM').rename('elevation')
    slope = terrain.select('slope')
    aspect = terrain.select('aspect')
    return elevation.addBands([slope, aspect])

def main():
    config = load_config()
    project_id = config['gee']['project_id']
    
    try:
        ee.Initialize(project=project_id)
    except Exception:
        print("\n--- Google Earth Engine Authentication Required ---")
        ee.Authenticate()
        ee.Initialize(project=project_id)
        print("Authentication successful!\n")

    mask_dir = os.path.join(os.path.dirname(__file__), config['paths']['output_dir'], 'masks')
    output_dir = os.path.join(os.path.dirname(__file__), config['paths']['output_dir'], 'features')
    os.makedirs(output_dir, exist_ok=True)
    
    target_crs = config['preprocessing']['target_crs']
    date_start = config['gee']['date_start']
    date_end = config['gee']['date_end']

    mask_files = sorted(glob.glob(os.path.join(mask_dir, 'mask_*.tif')))
    print(f"Found {len(mask_files)} masks. Extracting features for these locations...")
    
    for filepath in mask_files:
        filename = os.path.basename(filepath)
        idx_str = filename.split('_')[1].split('.')[0]
        idx = int(idx_str)
        
        out_tif = os.path.join(output_dir, f"feature_{idx:04d}.tif")
        if os.path.exists(out_tif):
            print(f"Skipping {idx}, feature stack already exists.")
            continue
            
        with rasterio.open(filepath) as src:
            bounds = src.bounds
            
        minx, miny, maxx, maxy = bounds
        padded_box_utm = gpd.GeoDataFrame(geometry=[box(minx, miny, maxx, maxy)], crs=target_crs)
        padded_box_4326 = padded_box_utm.to_crs('EPSG:4326')
        bounds_4326 = padded_box_4326.total_bounds
        
        roi = ee.Geometry.BBox(bounds_4326[0], bounds_4326[1], bounds_4326[2], bounds_4326[3])
        
        try:
            s2_img = get_s2_image(roi, date_start, date_end)
            terrain_img = get_terrain_image(roi)
            final_img = s2_img.addBands(terrain_img)
            
            print(f"Downloading feature stack {idx:04d}...")
            geemap.ee_export_image(
                final_img, 
                filename=out_tif, 
                scale=10, 
                region=roi, 
                crs=target_crs,
                file_per_band=False
            )
        except Exception as e:
            print(f"Failed to process {idx:04d}: {e}")

if __name__ == '__main__':
    main()