from Common.geotools import get_query_bbox
import planetary_computer as PC
from pystac_client import Client
import tqdm
import rasterio
import rasterio.mask
from pyproj import Transformer as T
import os
from shapely.geometry import Polygon
import urllib.request

def fetch(shp, output_fld, product='cop-dem-glo-30'):
    bbox, i_bbox, shp_srs = get_query_bbox(shp, return_all=True)
    api = Client.open('https://planetarycomputer.microsoft.com/api/stac/v1', modifier=PC.sign_inplace)
    res = api.search(collections=product, bbox=bbox)
    prg = tqdm.tqdm(total=len(res.item_collection()), desc="Fetching from Planetary")
    key = 'data'
    if product == 'nasadem':
        key = 'elevation'
    for item in res.items():
        with rasterio.open(item.assets[key].href) as ds:
            img_srs = ds.crs.to_epsg()
            aff = ds.profile['transform']
        if shp_srs == img_srs:
            o_bbox = [aff[0] * round((x + aff[0]*o) / aff[0]) for x,o in zip(i_bbox, [-2,-2,1,1])]
        else:
            tr = T.from_crs(shp_srs, img_srs, always_xy=True)
            x_min, y_min = tr.transform(i_bbox[0], i_bbox[1])
            x_max, y_max = tr.transform(i_bbox[2], i_bbox[3])
            o_bbox = [aff[0] * round((x + aff[0]*o) / aff[0]) for x,o in zip([x_min, y_min, x_max, y_max], [-2,-2,1,1])]

        ofn = os.path.join(output_fld, item.assets[key].get_absolute_href().split('?')[0].split('/')[-1])
        os.makedirs(os.path.dirname(ofn),exist_ok=True)
        with rasterio.open(item.assets[key].href) as img:
            out_img, out_geot = rasterio.mask.mask(img,
                                                   [Polygon.from_bounds(o_bbox[0], o_bbox[1],
                                                                        o_bbox[2], o_bbox[3])],
                                                   crop=True,
                                                   all_touched=True)
            out_meta = img.meta
        out_meta.update({"driver": "GTiff",
                         "height": out_img.shape[1],
                         "width": out_img.shape[2],
                         "transform": out_geot})
        with rasterio.open(ofn, "w", **out_meta) as dest:
            dest.write(out_img)
        prg.update()
    prg.close()

    return

def retrieve_geoid(out_fld):
    os.makedirs(out_fld, exist_ok=True)
    urllib.request.urlretrieve(
        'https://gitlab.orfeo-toolbox.org/orfeotoolbox/otb/-/raw/develop/Data/Input/DEM/egm96.grd?inline=false',
        os.path.join(out_fld, 'egm96.grd'))
    urllib.request.urlretrieve(
        'https://gitlab.orfeo-toolbox.org/orfeotoolbox/otb/-/raw/develop/Data/Input/DEM/egm96.grd.hdr?inline=false',
        os.path.join(out_fld, 'egm96.grd.hdr'))
    return