import streamlit as st
import geopandas as gpd
import tempfile
import os
import urllib.request
import json
from pathlib import Path
import datetime
from osgeo import gdal
import io
import zipfile
import base64

# 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',
}

# Helper functions
def wgs84_to_lv95(lat, lon):
    url = f'http://geodesy.geo.admin.ch/reframe/wgs84tolv95?easting={lat}&northing={lon}&format=json'
    with urllib.request.urlopen(url) as response:
        data = json.load(response)
    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'
    with urllib.request.urlopen(url) as response:
        data = json.load(response)
    return data['easting'], data['northing']

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, wgs84_max[0], wgs84_max[1], wgs84_min[1])
        return (bbox_wgs84, bbox)
    
    return None

def get_list_from_STAC_swisstopo(url, est, sud, ouest, nord, gdb=False):
    if gdb:
        lst_indesirables = []
    else:
        lst_indesirables = ['.xyz.zip', '.gdb.zip']

    sufixe_url = f"/items?bbox={est},{sud},{ouest},{nord}"
    url += sufixe_url
    res = []

    while url:
        with urllib.request.urlopen(url) as response:
            json_res = json.load(response)
        url = None
        links = json_res.get('links', None)
        if links:
            for link in links:
                if link['rel'] == 'next':
                    url = link['href']   
        
        for item in json_res['features']:
            for k, dic in item['assets'].items():
                href = dic['href']
                if gdb:
                    if href[-8:] == '.gdb.zip':
                        if len(dic['href'].split('/')[-1].split('_')) == 7:
                            res.append(dic['href'])
                else:
                    if href[-8:] not in lst_indesirables:
                        res.append(dic['href'])           
    return res

def suppr_doublons_bati3D_v2(lst_url):    
    dico = {}
    dxf_files = [url for url in lst_url if url[-8:] == '.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[-8:] == '.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

@st.cache_data
def get_urls(bbox_wgs84, data_types, resolutions):
    urls = []
    for data_type, enabled in data_types.items():
        if enabled:
            url = URL_STAC_SWISSTOPO_BASE + DIC_LAYERS[data_type]
            if data_type in ['mnt', 'ortho']:
                resolution = resolutions[data_type]
                tri = f'_{resolution}_'
                lst = [v for v in get_list_from_STAC_swisstopo(url, *bbox_wgs84) if tri in v]
                if data_type == 'mnt':
                    lst = suppr_doublons_list_mnt(lst)
                else:
                    lst = suppr_doublons_list_ortho(lst)
            elif data_type == 'mns':
                lst = [v for v in get_list_from_STAC_swisstopo(url, *bbox_wgs84) if 'raster' in v]
                lst = suppr_doublons_list_mnt(lst)
            elif data_type == 'bati3D_v2':
                lst = get_list_from_STAC_swisstopo(url, *bbox_wgs84)
                lst = suppr_doublons_bati3D_v2(lst)
            elif data_type == 'bati3D_v3':
                lst = get_list_from_STAC_swisstopo(url, *bbox_wgs84, gdb=True)
                lst = suppr_doublons_bati3D_v3(lst)
            urls.extend(lst)
    return urls

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:
        data = json.load(response)
        
    with open(fn_geojson, 'w') as f:
        json.dump(data, f)

def create_geojson_with_links(urls, bbox):
    features = []
    for url in urls:
        feature = {
            "type": "Feature",
            "geometry": {
                "type": "Polygon",
                "coordinates": [bbox]
            },
            "properties": {
                "url": url,
                "type": url.split('/')[-2].split('_')[0]
            }
        }
        features.append(feature)
    
    geojson = {
        "type": "FeatureCollection",
        "features": features
    }
    return json.dumps(geojson)

def merge_ortho_images(urls, output_format='GTiff'):
    with tempfile.TemporaryDirectory() as temp_dir:
        local_files = []
        for i, url in enumerate(urls):
            local_filename = os.path.join(temp_dir, f"ortho_{i}.tif")
            urllib.request.urlretrieve(url, local_filename)
            local_files.append(local_filename)
        
        vrt_options = gdal.BuildVRTOptions(resampleAlg='nearest', addAlpha=False)
        vrt_path = os.path.join(temp_dir, "merged.vrt")
        vrt = gdal.BuildVRT(vrt_path, local_files, options=vrt_options)
        vrt = None  # Close the dataset
        
        output_path = os.path.join(temp_dir, f"merged.{output_format.lower()}")
        if output_format == 'GTiff':
            translate_options = gdal.TranslateOptions(format="GTiff", creationOptions=["COMPRESS=LZW", "TILED=YES"])
        elif output_format == 'JPEG':
            translate_options = gdal.TranslateOptions(format="JPEG", creationOptions=["QUALITY=85"])
        elif output_format == 'PNG':
            translate_options = gdal.TranslateOptions(format="PNG", creationOptions=["COMPRESS=DEFLATE"])
        
        gdal.Translate(output_path, vrt_path, options=translate_options)
        
        with open(output_path, 'rb') as f:
            return f.read()

def prepare_download_package(urls, bbox, ortho_format):
    geojson_data = create_geojson_with_links(urls, bbox)
    ortho_urls = [url for url in urls if 'swissimage-dop10' in url]
    ortho_data = merge_ortho_images(ortho_urls, ortho_format) if ortho_urls else None
    
    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
        zip_file.writestr('download_links.geojson', geojson_data)
        if ortho_data:
            zip_file.writestr(f'merged_ortho.{ortho_format.lower()}', ortho_data)
    
    return zip_buffer.getvalue()

# 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")
data_types = {
    '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),
}

resolutions = {
    'mnt': st.sidebar.selectbox("MNT Resolution", [0.5, 2.0], index=0),
    'ortho': st.sidebar.selectbox("Orthophoto Resolution", [0.1, 2.0], index=0),
}

ortho_format = st.sidebar.selectbox("Ortho Output Format", ['GTiff', 'JPEG', 'PNG'], 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.0500, step=0.0001, format="%.4f")
with col2:
    ymin = st.number_input("Min Latitude", value=46.1800, step=0.0001, format="%.4f")
with col3:
    xmax = st.number_input("Max Longitude", value=6.2200, step=0.0001, format="%.4f")
with col4:
    ymax = st.number_input("Max Latitude", value=46.2500, step=0.0001, format="%.4f")

if 'bbox' not in st.session_state:
    st.session_state.bbox = None

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

if st.session_state.bbox:
    st.write(f"Selected bounding box (WGS84): {st.session_state.bbox}")
    
    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 Package"):
            with st.spinner("Preparing download package..."):
                urls = get_urls(bbox_wgs84, data_types, resolutions)
                if urls:
                    zip_data = prepare_download_package(urls, bbox_wgs84, ortho_format)
                    b64 = base64.b64encode(zip_data).decode()
                    href = f'<a href="data:application/zip;base64,{b64}" download="swiss_geospatial_data.zip">Download All Data</a>'
                    st.markdown(href, unsafe_allow_html=True)
                    st.success("Download package prepared. Click the link above to download.")
                else:
                    st.warning("No files found for the selected area and options.")
        
        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)
                    
                    # Provide download link for forest data
                    with open(tmp.name, 'r') as f:
                        forest_data = f.read()
                    b64 = base64.b64encode(forest_data.encode()).decode()
                    href = f'<a href="data:application/json;base64,{b64}" download="forest_data.geojson">Download Forest Data</a>'
                    st.markdown(href, unsafe_allow_html=True)
                
                os.unlink(tmp.name)
                st.success("Forest data prepared. Click the link above to download.")
    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 Package' to prepare all data for download.")