"""
Satellite data acquisition via STAC (SpatioTemporal Asset Catalog) API.
Supports Landsat 5/8/9 and Sentinel-2 imagery with user authentication
for data download where required.
"""

import warnings

import numpy as np

try:
    import pystac_client
    HAS_PYSTAC = True
except ImportError:
    HAS_PYSTAC = False

try:
    import rasterio
    from rasterio.windows import from_bounds
    HAS_RIO = True
except ImportError:
    HAS_RIO = False

SEARCH_URLS = {
    'planetary_computer': (
        'https://planetarycomputer.microsoft.com/api/stac/v1'
    ),
    'usgs_landsat': 'https://landsatlook.usgs.gov/stac-server',
    'sentinel2_aws': (
        'https://earth-search.aws.element84.com/v1'
    ),
}

BAND_ALIASES = {
    'landsat': {
        5: {
            'green': 'SR_B2',
            'nir': 'SR_B4',
            'swir1': 'SR_B5',
        },
        8: {
            'green': 'SR_B3',
            'nir': 'SR_B5',
            'swir1': 'SR_B6',
        },
        9: {
            'green': 'SR_B3',
            'nir': 'SR_B5',
            'swir1': 'SR_B6',
        },
    },
    'sentinel2': {
        'green': 'B03',
        'nir': 'B08',
        'swir1': 'B11',
    },
}


def check_dependencies():
    missing = []
    if not HAS_PYSTAC:
        missing.append('pystac-client')
    if not HAS_RIO:
        missing.append('rasterio')
    return missing


def validate_date_range(start_date, end_date):
    from datetime import datetime
    for d in (start_date, end_date):
        if isinstance(d, str):
            datetime.strptime(d, '%Y-%m-%d')
    return start_date, end_date


def search_scenes(aoi_geometry, start_date, end_date,
                  max_cloud=20, satellite='landsat',
                  catalog='planetary_computer'):
    """
    Search for satellite scenes intersecting the AOI.

    Parameters
    ----------
    aoi_geometry : dict
        GeoJSON-like geometry dict (Polygon).
    start_date, end_date : str
        Date range in 'YYYY-MM-DD' format.
    max_cloud : int
        Maximum cloud cover percentage.
    satellite : str
        'landsat' or 'sentinel2'.
    catalog : str
        STAC catalog URL key.

    Returns
    -------
    list of dicts with scene metadata.
    """
    if not HAS_PYSTAC:
        raise ImportError(
            'pystac-client is required. Install with: '
            'pip install pystac-client'
        )

    url = SEARCH_URLS.get(catalog, catalog)
    catalog_obj = pystac_client.Client.open(url)

    collections = {
        'landsat': [
            'landsat-c2-l2',
        ],
        'sentinel2': [
            'sentinel-2-l2a',
        ],
    }

    search = catalog_obj.search(
        collections=collections.get(satellite, collections['landsat']),
        intersects=aoi_geometry,
        datetime=f'{start_date}/{end_date}',
        query={
            'eo:cloud_cover': {'lt': max_cloud},
        },
    )

    items = list(search.items())
    results = []
    for item in items:
        dt = (
            item.datetime.isoformat()
            if item.datetime else None
        )
        results.append({
            'id': item.id,
            'datetime': dt,
            'cloud_cover': item.properties.get(
                'eo:cloud_cover', None
            ),
            'bbox': item.bbox,
            'hrefs': item.assets,
        })

    return results


def build_epoch_composite(item_list, aoi_bounds, resolution=30.0,
                          out_path=None):
    """
    Build a median composite from multiple scenes for the AOI.

    Parameters
    ----------
    item_list : list
        Scene metadata from search_scenes().
    aoi_bounds : tuple
        (minx, miny, maxx, maxy) in target CRS.
    resolution : float
        Target resolution in metres.
    out_path : str or None
        If provided, write composite GeoTIFF.

    Returns
    -------
    dict with keys: 'bands' (band_name: 2D array), 'transform', 'crs'.
    """
    if not HAS_RIO:
        raise ImportError(
            'rasterio is required. Install with: pip install rasterio'
        )

    band_stack = {'green': [], 'nir': [], 'swir1': []}
    out_crs = rasterio.CRS.from_string('EPSG:4326')

    for item in item_list:
        for band_key in band_stack:
            asset = None
            for href_key, href_val in item['hrefs'].items():
                if band_key.lower() in href_key.lower():
                    asset = href_val
                    break
            if asset is None:
                continue
            href = asset.href if hasattr(asset, 'href') else asset
            try:
                with rasterio.open(href) as src:
                    src_crs = src.crs
                    if src_crs and src_crs != out_crs:
                        from rasterio.warp import transform_bounds
                        win_bounds = transform_bounds(
                            src_crs, out_crs,
                            *aoi_bounds,
                        )
                        window = from_bounds(
                            *win_bounds, src.transform
                        )
                    else:
                        window = from_bounds(
                            *aoi_bounds, src.transform
                        )
                    data = src.read(
                        1, window=window,
                        out_shape=(1, max(1, int(window.height)),
                                   max(1, int(window.width))),
                    )
                    band_stack[band_key].append(data.squeeze())
            except Exception as exc:
                warnings.warn(
                    f'Failed to read {band_key} from '
                    f'{item["id"]}: {exc}'
                )
                continue

    composites = {}
    transform = None

    for band_key, arrays in band_stack.items():
        if not arrays:
            continue
        stacked = np.stack(arrays, axis=0)
        valid = np.all(stacked > 0, axis=0)
        median = np.full(stacked.shape[1:], np.nan)
        median[valid] = np.median(stacked[:, valid], axis=0)
        composites[band_key] = np.nan_to_num(median, nan=0.0)

    if out_path and composites:
        height = list(composites.values())[0].shape[0]
        width = list(composites.values())[0].shape[1]
        transform = rasterio.transform.from_bounds(
            aoi_bounds[0], aoi_bounds[1],
            aoi_bounds[2], aoi_bounds[3],
            width, height
        )
        with rasterio.open(
            out_path, 'w', driver='GTiff',
            height=height, width=width,
            count=len(composites), dtype='float32',
            crs=out_crs, transform=transform,
        ) as dst:
            for i, band_key in enumerate(composites, start=1):
                dst.write(
                    composites[band_key].astype('float32'), i
                )

    return {
        'bands': composites,
        'transform': transform,
        'crs': str(out_crs) if out_crs else None,
    }
