import streamlit as st
import urllib
import json
from pathlib import Path
import datetime
import geopandas as gpd
import tempfile
import os
import asyncio
import aiohttp
from functools import partial
from multiprocessing import Pool
import rasterio
from rasterio.merge import merge
from rasterio.warp import calculate_default_transform, reproject, Resampling
import numpy as np

# Constants
CATEGORIES = {
    'Gebueschwald': 'Forêt buissonnante',
    'Wald': 'Forêt',
    'Wald offen': 'Forêt claisemée',
    'Gehoelzflaeche': 'Zone boisée',
}
MERGE_CATEGORIES = True

URL_STAC_SWISSTOPO_BASE = 'https://data.geo.admin.ch/api/stac/v0.9/collections/'

DIC_LAYERS = {
    'ortho': 'ch.swisstopo.swissimage-dop10',
    'mnt': 'ch.swisstopo.swissalti3d',
    'mns': 'ch.swisstopo.swisssurface3d-raster',
    'bati3D_v2': 'ch.swisstopo.swissbuildings3d_2',
    'bati3D_v3': 'ch.swisstopo.swissbuildings3d_3_0',
}

# Functions
def wgs84_to_lv95(lat, lon):
    url = f'http://geodesy.geo.admin.ch/reframe/wgs84tolv95?easting={lat}&northing={lon}&format=json'
    site = urllib.request.urlopen(url)
    data = json.load(site)
    return data['easting'], data['northing']

def lv95_to_wgs84(x, y):
    url = f'http://geodesy.geo.admin.ch/reframe/lv95towgs84?easting={x}&northing={y}&format=json'
    f = urllib.request.urlopen(url)
    txt = f.read().decode('utf-8')
    json_res = json.loads(txt)
    return json_res

def detect_and_convert_bbox(bbox):
    xmin, ymin, xmax, ymax = bbox
    
    wgs84_margin = 0.9
    wgs84_bounds = {
        'xmin': 5.96 - wgs84_margin,
        'ymin': 45.82 - wgs84_margin,
        'xmax': 10.49 + wgs84_margin,
        'ymax': 47.81 + wgs84_margin
    }
    
    lv95_margin = 100000
    lv95_bounds = {
        'xmin': 2485000 - lv95_margin,
        'ymin': 1075000 - lv95_margin,
        'xmax': 2834000 + lv95_margin,
        'ymax': 1296000 + lv95_margin
    }
    
    if (wgs84_bounds['xmin'] <= xmin <= wgs84_bounds['xmax'] and
        wgs84_bounds['ymin'] <= ymin <= wgs84_bounds['ymax'] and
        wgs84_bounds['xmin'] <= xmax <= wgs84_bounds['xmax'] and
        wgs84_bounds['ymin'] <= ymax <= wgs84_bounds['ymax']):
        
        lv95_min = wgs84_to_lv95(xmin, ymin)
        lv95_max = wgs84_to_lv95(xmax, ymax)
        
        bbox_lv95 = (lv95_min[0], lv95_min[1], lv95_max[0], lv95_max[1])
        return (bbox, bbox_lv95)
    
    if (lv95_bounds['xmin'] <= xmin <= lv95_bounds['xmax'] and
        lv95_bounds['ymin'] <= ymin <= lv95_bounds['ymax'] and
        lv95_bounds['xmin'] <= xmax <= lv95_bounds['xmax'] and
        lv95_bounds['ymin'] <= ymax <= lv95_bounds['ymax']):
        
        wgs84_min = lv95_to_wgs84(xmin, ymin)
        wgs84_max = lv95_to_wgs84(xmax, ymax)
        
        bbox_wgs84 = (wgs84_min['easting'], wgs84_min['northing'], wgs84_max['easting'], wgs84_max['northing'])
        return (bbox_wgs84, bbox)
    
    return None

async def fetch(session, url):
    async with session.get(url) as response:
        return await response.json()

async def get_list_from_STAC_swisstopo_async(url, est, sud, ouest, nord, gdb=False):
    suffix_url = f"/items?bbox={est},{sud},{ouest},{nord}"
    url += suffix_url
    res = []

    async with aiohttp.ClientSession() as session:
        while url:
            json_res = await fetch(session, url)
            url = next((link['href'] for link in json_res.get('links', []) if link['rel'] == 'next'), None)
            
            for item in json_res['features']:
                for k, dic in item['assets'].items():
                    href = dic['href']
                    if gdb:
                        if href.endswith('.gdb.zip') and len(href.split('/')[-1].split('_')) == 7:
                            res.append(href)
                    elif not href.endswith(('.xyz.zip', '.gdb.zip')):
                        res.append(href)
    return res

def process_urls(args):
    layer, url, est, sud, ouest, nord, gdb, tri = args
    lst = [v for v in asyncio.run(get_list_from_STAC_swisstopo_async(url, est, sud, ouest, nord, gdb)) if tri in v]
    if layer == 'mnt' or layer == 'mns':
        return suppr_doublons_list_mnt(lst)
    elif layer == 'bati3D_v2':
        return suppr_doublons_bati3D_v2(lst)
    elif layer == 'bati3D_v3':
        return suppr_doublons_bati3D_v3(lst)
    elif layer == 'ortho':
        return suppr_doublons_list_ortho(lst)

def get_urls_async(bbox_wgs84, mnt=True, mns=True, bati3D_v2=True, bati3D_v3=True, ortho=True, mnt_resol=0.5, ortho_resol=0.1):
    est, sud, ouest, nord = bbox_wgs84
    tasks = []

    if mnt:
        mnt_resol = 0.5 if mnt_resol < 2 else 2
        tasks.append(('mnt', URL_STAC_SWISSTOPO_BASE + DIC_LAYERS['mnt'], est, sud, ouest, nord, False, f'_{mnt_resol}_'))

    if mns:
        tasks.append(('mns', URL_STAC_SWISSTOPO_BASE + DIC_LAYERS['mns'], est, sud, ouest, nord, False, 'raster'))

    if bati3D_v2:
        tasks.append(('bati3D_v2', URL_STAC_SWISSTOPO_BASE + DIC_LAYERS['bati3D_v2'], est, sud, ouest, nord, False, ''))

    if bati3D_v3:
        tasks.append(('bati3D_v3', URL_STAC_SWISSTOPO_BASE + DIC_LAYERS['bati3D_v3'], est, sud, ouest, nord, True, ''))

    if ortho:
        ortho_resol = 0.1 if ortho_resol < 2 else 2
        tasks.append(('ortho', URL_STAC_SWISSTOPO_BASE + DIC_LAYERS['ortho'], est, sud, ouest, nord, False, f'_{ortho_resol}_'))

    with Pool() as pool:
        results = pool.map(process_urls, tasks)

    return [url for sublist in results for url in sublist]

def suppr_doublons_bati3D_v2(lst_url):    
    dico = {}
    dxf_files = [url for url in lst_url if url.endswith('.dxf.zip')]
    for dxf in dxf_files:
        *a, date, feuille = dxf.split('/')[-2].split('_')
        dico.setdefault(feuille, []).append((date, dxf))
    res = []
    for k, liste in dico.items():
        res.append(sorted(liste, reverse=True)[0][1])
    return res

def suppr_doublons_bati3D_v3(lst_url):    
    dico = {}
    gdb_files = [url for url in lst_url if url.endswith('.gdb.zip')]
    for gdb in gdb_files:
        *a, date, feuille = gdb.split('/')[-2].split('_')
        dico.setdefault(feuille, []).append((date, gdb))
    res = []
    for k, liste in dico.items():
        res.append(sorted(liste, reverse=True)[0][1])
    return res

def suppr_doublons_list_ortho(lst):
    dic = {}
    for url in lst:
        nom, an, noflle, taille_px, epsg = url.split('/')[-1][:-4].split('_')
        dic.setdefault((noflle, float(taille_px)), []).append((an, url))
    res = []
    for noflle, lst in dic.items():
        an, url = sorted(lst, reverse=True)[0]
        res.append(url)
    return res

def suppr_doublons_list_mnt(lst):
    dic = {}
    for url in lst:
        nom, an, noflle, taille_px, epsg, inconnu = url.split('/')[-1][:-4].split('_')
        dic.setdefault((noflle, float(taille_px)), []).append((an, url))
    res = []
    for noflle, lst in dic.items():
        an, url = sorted(lst, reverse=True)[0]
        res.append(url)
    return res

def geojson_forest(bbox, fn_geojson):
    xmin, ymin, xmax, ymax = bbox
    url_base = 'https://hepiadata.hesge.ch/arcgis/rest/services/suisse/TLM_C4D_couverture_sol/FeatureServer/1/query?'

    sql = ' OR '.join([f"OBJEKTART='{cat}'" for cat in CATEGORIES.keys()])

    params = {
        "geometry": f"{xmin},{ymin},{xmax},{ymax}",
        "geometryType": "esriGeometryEnvelope",
        "returnGeometry": "true",
        "outFields": "OBJEKTART",
        "orderByFields": "OBJEKTART",
        "where": sql,
        "returnZ": "true",
        "outSR": '2056',
        "spatialRel": "esriSpatialRelIntersects",
        "f": "geojson"
    }
    query_string = urllib.parse.urlencode(params)
    url = url_base + query_string
    
    with urllib.request.urlopen(url) as response:
        response_data = response.read()
        data = json.loads(response_data)
        
        with open(fn_geojson, 'w') as f:
            json.dump(data, f)

def merge_forest_categories(gdf):
    if gdf.crs is None or gdf.crs.to_epsg() != 2056:
        gdf = gdf.to_crs(epsg=2056)

    gdf['category_fr'] = gdf['OBJEKTART'].map(CATEGORIES)
    dissolved = gdf.dissolve(by='category_fr', aggfunc='first')
    all_forests = dissolved.geometry.unary_union
    merged_gdf = gpd.GeoDataFrame(geometry=[all_forests], crs=gdf.crs)
    merged_gdf['category'] = 'All Forests'

    return merged_gdf

def download_and_merge_tiles(urls, output_path, target_crs='EPSG:2056'):
    with tempfile.TemporaryDirectory() as tmpdir:
        local_files = []
        for i, url in enumerate(urls):
            local_filename = os.path.join(tmpdir, f"tile_{i}.tif")
            urllib.request.urlretrieve(url, local_filename)
            local_files.append(local_filename)
        
        src_files_to_mosaic = []
        for file in local_files:
            src = rasterio.open(file)
            src_files_to_mosaic.append(src)
        
        mosaic, out_trans = merge(src_files_to_mosaic)
        
        out_meta = src_files_to_mosaic[0].meta.copy()
        out_meta.update({
            "driver": "GTiff",
            "height": mosaic.shape[1],
            "width": mosaic.shape[2],
            "transform": out_trans,
            "crs": src_files_to_mosaic[0].crs
        })
        
        with rasterio.open(output_path, "w", **out_meta) as dest:
            dest.write(mosaic)
        
        for src in src_files_to_mosaic:
            src.close()
        
        if out_meta['crs'] != target_crs:
            with rasterio.open(output_path) as src:
                transform, width, height = calculate_default_transform(
                    src.crs, target_crs, src.width, src.height, *src.bounds)
                kwargs = src.meta.copy()
                kwargs.update({
                    'crs': target_crs,
                    'transform': transform,
                    'width': width,
                    'height': height
                })
                
                reprojected_path = output_path.replace('.tif', '_reprojected.tif')
                with rasterio.open(reprojected_path, 'w', **kwargs) as dst:
                    for i in range(1, src.count + 1):
                        reproject(
                            source=rasterio.band(src, i),
                            destination=rasterio.band(dst, i),
                            src_transform=src.transform,
                            src_crs=src.crs,
                            dst_transform=transform,
                            dst_crs=target_crs,
                            resampling=Resampling.nearest)
                
                os.replace(reprojected_path, output_path)

    print(f"Merged image saved to {output_path}")

# Streamlit app
st.set_page_config(page_title="Swiss Geospatial Data Downloader", layout="wide")

st.title("Swiss Geospatial Data Downloader")

# Sidebar for data selection
st.sidebar.header("Data Selection")
mnt = st.sidebar.checkbox("Digital Terrain Model (MNT)", value=True)
mns = st.sidebar.checkbox("Digital Surface Model (MNS)", value=True)
bati3D_v2 = st.sidebar.checkbox("3D Buildings v2", value=True)
bati3D_v3 = st.sidebar.checkbox("3D Buildings v3", value=True)
ortho = st.sidebar.checkbox("Orthophotos", value=True)

mnt_resol = st.sidebar.selectbox("MNT Resolution", [0.5, 2.0], index=0)
ortho_resol = st.sidebar.selectbox("Orthophoto Resolution", [0.1, 2.0], index=0)

# Main content area
st.subheader("Enter Bounding Box Coordinates")
col1, col2, col3, col4 = st.columns(4)
with col1:
    xmin = st.number_input("Min Longitude", value=6.0, step=0.1)
with col2:
    ymin = st.number_input("Min Latitude", value=46.0, step=0.1)
with col3:
    xmax = st.number_input("Max Longitude", value=10.0, step=0.1)
with col4:
    ymax = st.number_input("Max Latitude", value=47.0, step=0.1)

if st.button("Set Bounding Box"):
    st.session_state.bbox = [xmin, ymin, xmax, ymax]

if 'bbox' in st.session_state:
    st.write(f"Selected bounding box (WGS84): {st.session_state.bbox}")
    
    # Convert bbox to LV95
    bbox_results = detect_and_convert_bbox(st.session_state.bbox)
    
    if bbox_results:
        bbox_wgs84, bbox_lv95 = bbox_results
        st.write(f"Converted bounding box (LV95): {bbox_lv95}")
        
        if st.button("Get Download Links"):
            with st.spinner("Fetching download links..."):
                urls = get_urls_async(bbox_wgs84, mnt, mns, bati3D_v2, bati3D_v3, ortho, mnt_resol, ortho_resol)
                
                if urls:
                    st.success(f"Found {len(urls)} files to download:")
                    for url in urls:
                        st.write(url)
                    
                    # Option to download files
                    if st.button("Download Files"):
                        with st.spinner("Downloading files..."):
                            download_path = Path("downloads") / f'swisstopo_extraction_{datetime.datetime.now().strftime("%Y%m%d_%H%M")}'
                            download_path.mkdir(parents=True, exist_ok=True)
                            for url in urls:
                                filename = url.split('/')[-1]
                                urllib.request.urlretrieve(url, download_path / filename)
                            st.success(f"Files downloaded to: {download_path}")
                else:
                    st.warning("No files found for the selected area and options.")
        
        # Option to download forest data
        if st.button("Download Forest Data"):
            with st.spinner("Downloading forest data..."):
                with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.geojson') as tmp:
                    geojson_forest(bbox_lv95, tmp.name)
                    gdf = gpd.read_file(tmp.name)
                    st.write(gdf)
                    
                    if MERGE_CATEGORIES:
                        merged_gdf = merge_forest_categories(gdf)
                        st.write("Merged Forest Data:")
                        st.write(merged_gdf)
                        
                        # Save merged data
                        merged_file = tmp.name.replace('.geojson', '_merged.geojson')
                        merged_gdf.to_file(merged_file, driver='GeoJSON')
                        st.success(f"Merged forest data saved to: {merged_file}")
                    
                os.unlink(tmp.name)
                st.success("Forest data downloaded and displayed.")
        
        # Option to merge image tiles
        if st.button("Merge Image Tiles"):
            with st.spinner("Merging image tiles..."):
                output_path = "merged_image.tif"
                download_and_merge_tiles(urls, output_path)
                st.success(f"Merged image saved to: {output_path}")
                
                # Optionally display the merged image
                with rasterio.open(output_path) as src:
                    image = src.read()
                    st.image(image.transpose(1, 2, 0), caption="Merged Image", use_column_width=True)
    else:
        st.error("Selected area is outside Switzerland. Please select an area within Switzerland.")

st.sidebar.info("This application allows you to download various types of geospatial data for Switzerland. Select the data types you want, enter the bounding box coordinates, and click 'Get Download Links' to see available files.")