import streamlit as st
import geopandas as gpd
import folium
from streamlit_folium import folium_static, st_folium
from folium.plugins import Draw
import requests
from shapely.geometry import box
import json

st.set_page_config(layout="wide", page_title="Extracteur de données géospatiales")

SWISSTOPO_LAYERS = {
    "SWISSIMAGE 10 cm": {"id": "ch.swisstopo.swissimage-dop10", "source": "swisstopo", "type": "image"},
    "Carte nationale 1:25'000": {"id": "ch.swisstopo.pixelkarte-farbe-pk25.noscale", "source": "swisstopo", "type": "image"},
    "MNT": {"id": "ch.swisstopo.swissalti3d", "source": "swisstopo", "type": "raster"},
    "Carte géologique": {"id": "ch.swisstopo.geologie-geologische_karte", "source": "swisstopo", "type": "image"},
    "Limites administratives": {"id": "ch.swisstopo.swissboundaries3d-gemeinde-flaeche.fill", "source": "swisstopo", "type": "vector"},
    "Réseau hydrographique": {"id": "ch.bafu.vec25-gewaessernetz", "source": "swisstopo", "type": "vector"},
    "swissBUILDINGS3D 3.0 Beta": {"id": "ch.swisstopo.swissbuildings3d_3_0", "source": "swisstopo", "type": "3d"},
}

ESRI_LAYERS = {
    "World Imagery": {"id": "World_Imagery", "source": "esri", "type": "image"},
    "World Elevation": {"id": "Elevation/World_Elevation", "source": "esri", "type": "raster"},
    "World Topographic": {"id": "World_Topo_Map", "source": "esri", "type": "image"},
    "World Street Map": {"id": "World_Street_Map", "source": "esri", "type": "image"},
    "World Terrain": {"id": "World_Terrain_Base", "source": "esri", "type": "image"},
}

def draw_box_on_map():
    m = folium.Map(location=[46.8, 8.2], zoom_start=8)
    Draw(draw_options={'polyline': False, 'polygon': False, 'circle': False, 'marker': False, 'circlemarker': False},
         edit_options={'edit': False}).add_to(m)
    return st_folium(m, width=700, height=500)

def create_geojson_from_box(bbox):
    gdf = gpd.GeoDataFrame({'geometry': [bbox]}, crs="EPSG:4326")
    return json.loads(gdf.to_json()), json.loads(gdf.to_crs("EPSG:2056").to_json())

@st.cache_data
def get_download_url(bbox, layer_info):
    if layer_info["source"] == "swisstopo":
        api_url = f"https://data.geo.admin.ch/api/stac/v0.9/collections/{layer_info['id']}/items"
        params = {"bbox": ",".join(map(str, bbox)), "limit": 1}
        response = requests.get(api_url, params=params)
        data = response.json()
        if data['features']:
            feature = data['features'][0]
            asset_keys = feature['assets'].keys()
            data_key = 'rgb' if 'rgb' in asset_keys else 'data' if 'data' in asset_keys else next(iter(asset_keys))
            return feature['assets'][data_key]['href']
    else:  # ESRI
        service_url = f"https://services.arcgisonline.com/arcgis/rest/services/{layer_info['id']}/MapServer/export"
        params = {
            "bbox": f"{bbox[0]},{bbox[1]},{bbox[2]},{bbox[3]}",
            "bboxSR": 4326,
            "size": "1000,1000",
            "format": "png",
            "f": "html",
            "imageSR": 4326,
            "transparent": "true"
        }
        return f"{service_url}?{'&'.join([f'{k}={v}' for k, v in params.items()])}"
    return None

def main():
    st.title("Extracteur de données Swisstopo et ESRI")

    col1, col2 = st.columns([1, 3])

    with col1:
        st.subheader("Sélection des couches")
        source = st.radio("Source des données", ["Swisstopo", "ESRI"])
        
        layers = SWISSTOPO_LAYERS if source == "Swisstopo" else ESRI_LAYERS
        
        selected_layers = st.multiselect("Couches à extraire", list(layers.keys()))
        
        st.markdown("---")
        st.markdown("## À propos\nExtracteur de données géospatiales Swisstopo et ESRI.\nDéveloppé par Vertdure")

    with col2:
        st.subheader("1. Définir la zone d'intérêt")
        method = st.radio("Méthode", ["Dessiner", "GeoJSON"], horizontal=True)
        if method == "Dessiner":
            if 'bbox' not in st.session_state:
                map_data = draw_box_on_map()
                if map_data['last_active_drawing']:
                    coords = map_data['last_active_drawing']['geometry']['coordinates'][0]
                    bbox = box(min(c[0] for c in coords), min(c[1] for c in coords),
                               max(c[0] for c in coords), max(c[1] for c in coords))
                    st.session_state['bbox'] = bbox
                    st.session_state['geojson_wgs84'], st.session_state['geojson_swiss'] = create_geojson_from_box(bbox)
                    st.success("Zone sélectionnée!")
            else:
                st.success("Zone déjà sélectionnée. Cliquez sur 'Effacer la zone' pour redessiner.")
                if st.button("Effacer la zone"):
                    del st.session_state['bbox']
                    del st.session_state['geojson_wgs84']
                    del st.session_state['geojson_swiss']
                    st.experimental_rerun()
        else:
            uploaded_file = st.file_uploader("GeoJSON", type="geojson")
            if uploaded_file:
                gdf = gpd.read_file(uploaded_file)
                st.session_state['geojson_wgs84'] = json.loads(gdf.to_crs(4326).to_json())
                st.session_state['geojson_swiss'] = json.loads(gdf.to_crs(2056).to_json())
                st.success("GeoJSON chargé!")

        if st.button("Obtenir les liens de téléchargement"):
            if 'geojson_wgs84' not in st.session_state or 'geojson_swiss' not in st.session_state:
                st.error("Définissez d'abord une zone d'intérêt.")
            elif not selected_layers:
                st.error("Sélectionnez au moins une couche à extraire.")
            else:
                for layer_name in selected_layers:
                    layer_info = layers[layer_name]
                    bbox = gpd.GeoDataFrame.from_features(st.session_state['geojson_swiss' if layer_info["source"] == "swisstopo" else 'geojson_wgs84']).total_bounds
                    st.write(f"Traitement de la couche : {layer_name}")
                    st.write(f"Bbox : {bbox}")
                    download_url = get_download_url(bbox, layer_info)
                    if download_url:
                        st.markdown(f"[Télécharger {layer_name}]({download_url})")
                    else:
                        st.error(f"Impossible d'obtenir le lien pour {layer_name}")

if __name__ == "__main__":
    main()