import gradio as gr
import geopandas as gpd
import os
import shutil
import zipfile
from shapely.geometry import mapping
import folium
import plotly.express as px
import tempfile
import logging
from pathlib import Path

# Set up logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class GeoPulseApp:
    def __init__(self):
        # Create a temporary directory for processing files
        self.temp_dir = tempfile.mkdtemp(prefix="geopulse_")
        logger.info(f"Created temporary directory at {self.temp_dir}")

    def cleanup(self):
        # Clean up temporary directory after use
        if self.temp_dir and os.path.exists(self.temp_dir):
            try:
                shutil.rmtree(self.temp_dir)
                logger.info(f"Cleaned up temporary directory: {self.temp_dir}")
            except Exception as e:
                logger.error(f"Error cleaning up: {e}")

    def process_files(self, file_obj, view_type: str) -> str:
        """Process the uploaded file and create visualization."""
        try:
            if file_obj is None:
                raise gr.Error("Please upload a ZIP file or Shapefile")

            # Save the uploaded file to the temporary directory
            file_path = os.path.join(self.temp_dir, file_obj.name)
            with open(file_path, "wb") as f:
                f.write(file_obj.read())  # Correct method to handle Gradio file object

            # Check file extension
            if file_path.lower().endswith('.zip'):
                # Extract the shapefile from the ZIP
                with zipfile.ZipFile(file_path, 'r') as zip_ref:
                    zip_ref.extractall(self.temp_dir)

                # Find the .shp file
                shp_files = list(Path(self.temp_dir).rglob("*.shp"))
                if not shp_files:
                    raise gr.Error("No shapefile (.shp) found in the ZIP file")
                shapefile_path = str(shp_files[0])
            elif file_path.lower().endswith('.shp'):
                shapefile_path = file_path  # Use directly if it's a .shp file
            else:
                raise gr.Error("Please upload a valid ZIP file or Shapefile (*.shp)")

            # Load the shapefile
            gdf = gpd.read_file(shapefile_path)

            if gdf.empty:
                raise gr.Error("The shapefile contains no data")

            # Convert to WGS84 if needed
            if gdf.crs is None:
                gdf.set_crs("EPSG:4326", inplace=True)
            elif gdf.crs != "EPSG:4326":
                gdf = gdf.to_crs("EPSG:4326")

            # Create visualization based on the selected view type
            if view_type == "2D":
                return self.create_2d_map(gdf)
            elif view_type == "3D":
                return self.create_3d_map(gdf)
            else:
                raise gr.Error("Please select either 2D or 3D view")

        except zipfile.BadZipFile:
            raise gr.Error("The file is not a valid ZIP file")
        except Exception as e:
            logger.error(f"Error processing file: {e}")
            raise gr.Error(f"Error processing file: {str(e)}")
        finally:
            self.cleanup()

    def create_2d_map(self, gdf: gpd.GeoDataFrame) -> str:
        """Create 2D map visualization."""
        try:
            # Re-project to a projected CRS for accurate centroid calculation
            gdf_projected = gdf.to_crs(epsg=3857)
            center = [gdf_projected.geometry.centroid.y.mean(), gdf_projected.geometry.centroid.x.mean()]

            m = folium.Map(location=center, zoom_start=4)

            folium.GeoJson(
                gdf,
                style_function=lambda x: {
                    'fillColor': '#ffff00',
                    'color': '#000000',
                    'weight': 1,
                    'fillOpacity': 0.5
                }
            ).add_to(m)

            # Save map to temporary file
            map_path = os.path.join(self.temp_dir, "map_2d.html")
            m.save(map_path)

            return map_path  # Return the path for download link
        except Exception as e:
            raise gr.Error(f"Error creating 2D map: {str(e)}")

    def create_3d_map(self, gdf: gpd.GeoDataFrame) -> str:
        """Create 3D map visualization."""
        try:
            # Sample if too large
            if len(gdf) > 1000:
                gdf = gdf.sample(n=1000)

            # Calculate centroids in a projected CRS for accuracy
            gdf_projected = gdf.to_crs(epsg=3857)
            gdf['lon'] = gdf_projected.geometry.centroid.x
            gdf['lat'] = gdf_projected.geometry.centroid.y
            gdf['area'] = gdf.geometry.area

            fig = px.scatter_3d(
                gdf,
                x='lon',
                y='lat',
                z='area',
                title="3D GIS Visualization"
            )

            # Save to temporary file
            map_path = os.path.join(self.temp_dir, "map_3d.html")
            fig.write_html(map_path)

            return map_path  # Return the path for download link
        except Exception as e:
            raise gr.Error(f"Error creating 3D map: {str(e)}")

def main():
    app = GeoPulseApp()

    with gr.Blocks(title="GeoPulse: GIS Data Visualizer") as interface:
        gr.Markdown("# GeoPulse: GIS Data Visualizer")

        with gr.Row():
            with gr.Column():
                # File input component for uploading ZIP files
                file_input = gr.File(
                    label="Upload Shapefile ZIP or Shapefile",
                    file_types=[".zip", ".shp"],  # Accept zip or shapefiles
                )
                view_type = gr.Radio(
                    choices=["2D", "3D"],
                    value="2D",
                    label="View Type"
                )
                process_btn = gr.Button("Process and Visualize")

            with gr.Column():
                # Output area for the visualization
                output = gr.HTML(label="Visualization")

        # Trigger processing when the button is clicked
        process_btn.click(
            fn=app.process_files,
            inputs=[file_input, view_type],
            outputs=output
        )

        # Instructions for the user
        gr.Markdown(""" 
        ### Instructions:
        1. Prepare a ZIP file containing your shapefile (.shp, .shx, and .dbf files)
        2. Click 'Upload Shapefile ZIP or Shapefile' and select your file
        3. Choose between 2D or 3D visualization
        4. Click 'Process and Visualize'
        5. After processing, a download link for the map will appear below
        """)

    # Launch the interface
    interface.launch(server_name="0.0.0.0", server_port=7860)

if __name__ == "__main__":
    main()