import re
import itertools
from lxml import html
from spatialist.vector import Vector, wkt2vector, bbox
from spatialist.auxil import utm_autodetect
from pyroSAR.drivers import ID
from cesard.ancillary import get_max_ext, buffer_min_overlap, get_kml, combine_polygons
from osgeo import ogr


def tile_from_aoi(
        vector: Vector | list[Vector],
        epsg: int | list[int] | None = None,
        strict: bool = True,
        return_geometries: bool = False,
        tilenames: list[str] | None = None
) -> list[str | Vector]:
    """
    Return a list of MGRS tile IDs or vector objects overlapping one or multiple areas of interest.
    
    Parameters
    -------
    vector:
        The vector object(s) to read. CRS must be EPSG:4236.
    epsg:
        Define which EPSG code(s) are allowed for the tile selection.
        If None, all tile IDs are returned regardless of projection.
    strict:
        Strictly only return the names/geometries of the overlapping tiles in the target projection
        or also allow reprojection of neighbouring tiles?
        In the latter case a tile name takes the form <tile ID>_<EPSG code>, e.g. `33TUN_32632`.
        Only applies if argument `epsg` is of type `int` or a list with one element.
    return_geometries:
        return a list of :class:`spatialist.vector.Vector` geometry objects (or just the tile names)?
    tilenames:
        an optional list of MGRS tile names to limit the selection
    
    Returns
    -------
        A list of unique MGRS tile IDs or :class:`spatialist.vector.Vector`
        objects with an attribute `mgrs` containing the tile ID.
    """
    if isinstance(epsg, int):
        epsg = [epsg]
    sortkey = None
    if return_geometries:
        sortkey = lambda x: x.mgrs
    kml = get_kml()
    with Vector(kml, driver='KML') as vec_kml:
        tiles = []
        with combine_polygons(vector, multipolygon=True) as vec_aoi:
            feature = vec_aoi.getFeatureByIndex(0)
            geom = feature.GetGeometryRef()
            vec_kml.layer.SetSpatialFilter(geom)
            feature = geom = None
            if tilenames is not None:
                values = ", ".join([f"'{x}'" for x in tilenames])
                sql_where = f"Name IN ({values})"
                vec_kml.layer.SetAttributeFilter(sql_where)
            for tile in vec_kml.layer:
                tilename = tile.GetField('Name')
                attrib = description2dict(tile.GetField('Description'))
                epsg_target = None
                if epsg is not None and attrib['EPSG'] not in epsg:
                    if len(epsg) == 1 and not strict:
                        epsg_target = int(epsg[0])
                        tilename += '_{}'.format(epsg_target)
                    else:
                        continue
                if return_geometries:
                    wkt = multipolygon2polygon(attrib['UTM_WKT'])
                    geom = wkt2vector_regrid(wkt=wkt,
                                             epsg_in=attrib['EPSG'],
                                             epsg_out=epsg_target)
                    geom.mgrs = tilename
                    field_defn = ogr.FieldDefn("mgrs", ogr.OFTString)
                    geom.layer.CreateField(field_defn)
                    feature = geom.getFeatureByIndex(0)
                    feature.SetField("mgrs", tilename)
                    geom.layer.SetFeature(feature)
                    tiles.append(geom)
                else:
                    tiles.append(tilename)
            tile = None
            geom = None
            layer = None
    return sorted(tiles, key=sortkey)


def aoi_from_tile(
        tile: str | list[str]
) -> Vector | list[Vector]:
    """
    Extract one or multiple MGRS tiles from the global Sentinel-2 tiling
    grid and return it as a :class:`~spatialist.vector.Vector` object.
    
    Parameters
    ----------
    tile:
        The MGRS tile ID(s) that should be extracted and returned as a vector object.
        Can also be expressed as <tile ID>_<EPSG code> (e.g. `33TUN_32632`). In this case the geometry
        of the tile is reprojected to the target EPSG code, its corner coordinates rounded to multiples
        of 10, and a new :class:`~spatialist.vector.Vector` object created.
    
    Returns
    -------
        either a single object or a list depending on `tile`
    """
    kml = get_kml()
    if isinstance(tile, str):
        tile = [tile]
    
    tilenames = []
    epsg_codes = []
    pattern = '([A-Z0-9]{5})_?([0-9]+)?'
    for i in tile:
        tilename, epsg = re.search(pattern, i).groups()
        tilenames.append(tilename)
        epsg_codes.append(epsg)
    
    values = ", ".join([f"'{x}'" for x in tilenames])
    sql_where = f"Name IN ({values})"
    out = []
    with Vector(kml, driver='KML') as vec:
        vec.layer.SetAttributeFilter(sql_where)
        for i, feat in enumerate(vec.layer):
            tilename = feat.GetField('Name')
            attrib = description2dict(feat.GetField('Description'))
            wkt = multipolygon2polygon(attrib['UTM_WKT'])
            epsg_target = epsg_codes[i]
            geom = wkt2vector_regrid(wkt=wkt,
                                     epsg_in=attrib['EPSG'],
                                     epsg_out=epsg_target)
            geom.mgrs = tilename
            field_defn = ogr.FieldDefn("mgrs", ogr.OFTString)
            geom.layer.CreateField(field_defn)
            feature = geom.getFeatureByIndex(0)
            feature.SetField("mgrs", tilename)
            geom.layer.SetFeature(feature)
            out.append(geom)
    if len(out) == 1:
        return out[0]
    else:
        return out


def description2dict(
        description: str
) -> dict[str, str | int]:
    """
    Convert the HTML description field of the MGRS tile KML file to a dictionary.
    
    Parameters
    ----------
    description:
        The plain text of the `Description` field
    
    Returns
    -------
        A dictionary with keys 'TILE_ID', 'EPSG', 'MGRS_REF', 'UTM_WKT' and 'LL_WKT'.
        The value of field 'EPSG' is of type integer, all others are strings.
    """
    attrib = html.fromstring(description)
    attrib = [x for x in attrib.xpath('//tr/td//text()') if x != ' ']
    attrib = dict(zip(attrib[0::2], attrib[1::2]))
    attrib['EPSG'] = int(attrib['EPSG'])
    return attrib


def aoi_from_scene(
        scene: ID,
        multi: bool = True,
        percent: int | float = 1
) -> list[dict[str, dict[str, float] | int]]:
    """
    Get processing AOIs for a SAR scene. The MGRS grid requires a SAR
    scene to be geocoded to multiple UTM zones depending on the overlapping
    MGRS tiles and their projection. This function returns the following
    for each UTM zone group:
    
    - the extent in WGS84 coordinates (key `extent`)
    - the extent in UTM coordinates (key `extent_utm`)
    - the EPSG code of the UTM zone (key `epsg`)
    
    A minimum overlap of the AOIs with the SAR scene is ensured by buffering
    the AOIs if necessary. The minimum overlap can be controlled with
    parameter `percent`.
    
    Parameters
    ----------
    scene:
        the SAR scene object
    multi:
        split into multiple AOIs per overlapping UTM zone or just one AOI
        covering the whole scene. In the latter case the best matching UTM
        zone is auto-detected
        (using function :func:`spatialist.auxil.utm_autodetect`).
    percent:
        the minimum overlap in percent of each AOI with the SAR scene.
        See function :func:`cesard.ancillary.buffer_min_overlap`.

    Returns
    -------
        a list of dictionaries with keys `extent`, `extent_utm`, `epsg`
    """
    out = []
    if multi:
        # extract all overlapping tiles
        with scene.geometry() as geom:
            tiles = tile_from_aoi(vector=geom, return_geometries=True)
        
        # group tiles by UTM zone
        def fn(x):
            return x.getProjection(type='epsg')
        
        for zone, group in itertools.groupby(tiles, lambda x: fn(x)):
            geometries = list(group)
            # get UTM EPSG code
            epsg = geometries[0].getProjection(type='epsg')
            # get maximum extent of tile group
            ext_utm = get_max_ext(geometries=geometries)
            del geometries
            with bbox(ext_utm, epsg) as geom1:
                # ensure a minimum overlap between AOI and pre-processed scene
                with scene.geometry() as geom2:
                    geom2.reproject(epsg)
                    # 60 m to keep aligned to MGRS tile size and overlaps
                    # see ancillary.check_spacing
                    with buffer_min_overlap(geom1=geom1, geom2=geom2,
                                            percent=percent, step=60) as buffered:
                        ext_utm = buffered.extent
                # convert extent to EPSG:4326
                geom1.reproject(projection=4326)
                ext = geom1.extent
            out.append({'extent': ext, 'epsg': epsg,
                        'extent_utm': ext_utm})
    else:
        with scene.bbox() as geom:
            ext = geom.extent
            # auto-detect UTM zone
            epsg = utm_autodetect(geom, 'epsg')
            # get all tiles, reprojected to the target UTM zone if necessary
            tiles = tile_from_aoi(vector=geom, epsg=epsg,
                                  return_geometries=True, strict=False)
        ext_utm = get_max_ext(geometries=tiles)
        del tiles
        out.append({'extent': ext, 'epsg': epsg,
                    'extent_utm': ext_utm})
    return out


def multipolygon2polygon(wkt: str) -> str:
    """
    Convert a MultiPolygon WKT with one Polygon to a simple Polygon WKT.
    The Sentinel-2 KML grid file stores all geometries as MultiPolygons.
    This function simply converts the geometries to simple Polygons.
    Not all geometries in the KML file have been checked. In case there are
    ever multiple Polygons in one MultiPolygon, an RuntimeError is raised.
    All other geometries are returned as is.
    
    Parameters
    ----------
    wkt:
        A geometry WKT representation.

    Returns
    -------
        the output WKT geometry. Either the original geometry or
        a Polygon extracted from a single-Polygon MultiPolygon.
    """
    geom1 = ogr.CreateGeometryFromWkt(wkt)
    poly_type = geom1.GetGeometryType()
    poly_num = geom1.GetGeometryCount()
    if poly_type == ogr.wkbMultiPolygon:
        if poly_num == 1:
            geom2 = geom1.GetGeometryRef(0).Clone()
            wkt = geom2.ExportToWkt()
            geom2 = None
        else:
            raise RuntimeError('got a MultiPolygon with multiple Polygons')
    geom1 = None
    return wkt


def wkt2vector_regrid(
        wkt: str,
        epsg_in: int,
        epsg_out: int | None = None
) -> Vector:
    """
    Convert a WKT geometry to a :class:`spatialist.vector.Vector` object and
    optionally reproject and regrid it.

    Parameters
    ----------
    wkt:
        the WKT string
    epsg_in:
        the EPSG code for the CRS of `wkt`
    epsg_out:
        and optional target CRS to reproject the geometry

    Returns
    -------
        the geometry object
    
    See Also
    --------
    spatialist.vector.wkt2vector
    """
    if epsg_out is None:
        return wkt2vector(wkt, epsg_in)
    else:
        with wkt2vector(wkt, epsg_in) as tmp:
            tmp.reproject(epsg_out)
            ext = tmp.extent
            for k, v in ext.items():
                ext[k] = round(v / 10) * 10
        return bbox(ext, crs=epsg_out)
