import streamlit as st
import folium
from streamlit_folium import st_folium
from folium.plugins import Draw
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
import concurrent.futures
import requests
from functools import partial

# 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={lon}&northing={lat}&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['northing'], data['easting']

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(ymin, xmin)
        lv95_max = wgs84_to_lv95(ymax, xmax)
        
        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[1], wgs84_min[0], wgs84_max[1], wgs84_max[0])
        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 = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
        future_to_data_type = {
            executor.submit(
                get_urls_for_data_type, 
                data_type, 
                bbox_wgs84, 
                resolutions.get(data_type)
            ): data_type for data_type, enabled in data_types.items() if enabled
        }
        for future in concurrent.futures.as_completed(future_to_data_type):
            data_type = future_to_data_type[future]
            try:
                urls.extend(future.result())
            except Exception as exc:
                st.error(f"Error fetching URLs for {data_type}: {exc}")
    return urls

def get_urls_for_data_type(data_type, bbox_wgs84, resolution=None):
    url = URL_STAC_SWISSTOPO_BASE + DIC_LAYERS[data_type]
    if data_type in ['mnt', 'ortho']:
        tri = f'_{resolution}_'
        lst = [v for v in get_list_from_STAC_swisstopo(url, *bbox_wgs84) if tri in v]
        if data_type == 'mnt':
            return suppr_doublons_list_mnt(lst)
        else:
            return 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]
        return suppr_doublons_list_mnt(lst)
    elif data_type == 'bati3D_v2':
        lst = get_list_from_STAC_swisstopo(url, *bbox_wgs84)
        return suppr_doublons_bati3D_v2(lst)
    elif data_type == 'bati3D_v3':
        lst = get_list_from_STAC_swisstopo(url, *bbox_wgs84, gdb=True)
        return suppr_doublons_bati3D_v3(lst)
    return []

def fetch_url(url):
    response = requests.get(url)
    return response.content

def merge_ortho_images(urls, output_format='GTiff'):
    try:
        with tempfile.TemporaryDirectory() as temp_dir:
            local_files = []
            with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
                future_to_url = {executor.submit(fetch_url, url): url for url in urls}
                for i, future in enumerate(concurrent.futures.as_completed(future_to_url)):
                    url = future_to_url[future]
                    try:
                        data = future.result()
                        local_filename = os.path.join(temp_dir, f"ortho_{i}.tif")
                        with open(local_filename, 'wb') as f:
                            f.write(data)
                        local_files.append(local_filename)
                    except Exception as exc:
                        st.error(f"Error downloading {url}: {exc}")
            
            if not local_files:
                st.error("No ortho images were successfully downloaded.")
                return None
            
            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"])
            else:
                st.error(f"Unsupported output format: {output_format}")
                return None
            
            gdal.Translate(output_path, vrt_path, options=translate_options)
            
            if not os.path.exists(output_path):
                st.error(f"Failed to create merged image: {output_path}")
                return None
            
            with open(output_path, 'rb') as f:
                return f.read()
    except Exception as e:
        st.error(f"Error in merge_ortho_images: {e}")
        return None

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)

@st.cache_data
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)
        else:
            st.warning("Failed to merge ortho images. Only download links will be included in the package.")
    
    return zip_buffer.getvalue()

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)

# 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("Select Bounding Box")

# Create a map centered on Switzerland
m = folium.Map(location=[46.8182, 8.2275], zoom_start=8)

# Add rectangle draw control
draw = Draw(
    draw_options={
        'rectangle': True,
        'polygon': False,
        'polyline': False,
        'circle': False,
        'marker': False,
        'circlemarker': False
    },
    edit_options={'edit': False}
)
draw.add_to(m)

# Use st_folium to render the map and get the drawn bbox
output = st_folium(m, width=700, height=500)

# Initialize session state for bbox
if 'bbox' not in st.session_state:
    st.session_state.bbox = [6.0, 46.0, 10.0, 47.0]  # Default values for Switzerland

# Update bbox if a new one is drawn
if output['last_active_drawing']:
    coordinates = output['last_active_drawing']['geometry']['coordinates'][0]
    st.session_state.bbox = [
        min(coord[0] for coord in coordinates),
        min(coord[1] for coord in coordinates),
        max(coord[0] for coord in coordinates),
        max(coord[1] for coord in coordinates)
    ]

# Display and allow editing of bounding box coordinates
st.subheader("Enter Bounding Box Coordinates")
col1, col2, col3, col4 = st.columns(4)
with col1:
    xmin = st.number_input("Min Longitude", value=st.session_state.bbox[0], format="%.4f", key="xmin")
with col2:
    ymin = st.number_input("Min Latitude", value=st.session_state.bbox[1], format="%.4f", key="ymin")
with col3:
    xmax = st.number_input("Max Longitude", value=st.session_state.bbox[2], format="%.4f", key="xmax")
with col4:
    ymax = st.number_input("Max Latitude", value=st.session_state.bbox[3], format="%.4f", key="ymax")

# Update session state if coordinates are manually changed
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, draw a bounding box on the map, and click 'Get Download Package' to prepare all data for download.")