import geopandas as gpd
import matplotlib.pyplot as plt
import requests
import os
import pandas as pd
from smolagents import tool
from typing import Dict, Tuple, Optional
from matplotlib.figure import Figure
from matplotlib.axes import Axes
from shapely.geometry.base import BaseGeometry

def _download_geojson(url: str, file_name: str) -> str:
    """Downloads a GeoJSON file if it doesn't exist.

    Args:
        url (str): The URL of the GeoJSON file.
        file_name (str): The name of the file to save the data in.

    Returns:
        str: The path to the downloaded file.
    """
    data_dir = "france_data"
    if not os.path.exists(data_dir):
        os.makedirs(data_dir)
        
    file_path = os.path.join(data_dir, file_name)
    
    if not os.path.exists(file_path):
        print(f"Downloading {file_name} from {url}...")
        response = requests.get(url)
        response.raise_for_status()  # Raise an exception for bad status codes
        
        with open(file_path, 'w') as f:
            f.write(response.text)
        print("Download complete.")
        
    return file_path

def get_france_geodata(level: str = 'regions') -> gpd.GeoDataFrame:
    """Gets a GeoDataFrame for Metropolitan France with its regions or departments.
    
    Args:
        level (str): The administrative level to draw ('regions' or 'departments').
    
    Returns:
        gpd.GeoDataFrame: A GeoDataFrame with the requested administrative level.
    """
    if level == 'regions':
        url = "https://raw.githubusercontent.com/gregoiredavid/france-geojson/master/regions.geojson"
        file_name = "regions.geojson"
    elif level == 'departments':
        url = "https://raw.githubusercontent.com/gregoiredavid/france-geojson/master/departements.geojson"
        file_name = "departements.geojson"
    else:
        raise ValueError("level must be 'regions' or 'departments'")

    geojson_path = _download_geojson(url, file_name)
    gdf = gpd.read_file(geojson_path)
    
    # Although the geojson files are for metropolitan France, we can filter to be safe.
    if level == 'regions':
        # Metropolitan region codes are between 11 and 94.
        gdf['code'] = gdf['code'].astype(int)
        france_metropolitan = gdf[gdf['code'].between(11, 94)]
    else: # departments
        # Metropolitan department codes are 01-19, 21-95, 2A, 2B. Corsica (20) is split.
        metro_codes = [f'{i:02d}' for i in range(1, 20)] + [f'{i:02d}' for i in range(21, 96)] + ['2A', '2B']
        france_metropolitan = gdf[gdf['code'].isin(metro_codes)]

    france_metropolitan = france_metropolitan.to_crs(epsg=2154)
    return france_metropolitan

@tool
def draw_france_map(level: str = 'regions') -> Tuple[Figure, Axes]:
    """Draws a map of Metropolitan France with its regions or departments.
    
    Args:
        level (str): The administrative level to draw ('regions' or 'departments').
    
    Returns:
        Tuple[Figure, Axes]: A tuple containing the Matplotlib figure and axes objects.
    """
    france_metropolitan = get_france_geodata(level)
    
    fig, ax = plt.subplots(1, 1, figsize=(15, 12))
    
    # Plot with a single color
    france_metropolitan.plot(ax=ax, color='lightgray', edgecolor='black')
    
    minx, miny, maxx, maxy = france_metropolitan.total_bounds
    
    padding = 0.1
    ax.set_xlim(minx - padding * (maxx - minx), maxx + padding * (maxx - minx))
    ax.set_ylim(miny - padding * (maxy - miny), maxy + padding * (maxy - miny))
    
    ax.set_aspect('equal', adjustable='box')
    ax.set_axis_off()
    ax.set_title(f'Metropolitan France with {level.capitalize()}', fontsize=20)
    
    return fig, ax

@tool
def get_geodata_mapping(level: str = 'regions') -> Dict[str, BaseGeometry]:
    """Returns a mapping from region/department name to its polygon.
    
    Args:
        level (str): The administrative level to get the mapping for ('regions' or 'departments').
    
    Returns:
        Dict[str, BaseGeometry]: A dictionary mapping the name to the polygon.
    """
    france_metropolitan = get_france_geodata(level)
        
    mapping = {row['nom']: row['geometry'] for _, row in france_metropolitan.iterrows()}
    
    return mapping

@tool
def plot_geodata(geodata: gpd.GeoDataFrame, ax: Axes, color: str = None, edgecolor: str = 'black', alpha: float = 1.0, output_path: Optional[str] = None) -> Optional[str]:
    """Plots geodata on a given map axes and optionally saves the map as an image file.
    
    Args:
        geodata (gpd.GeoDataFrame): The geodata to plot.
        ax (Axes): The axes to plot on.
        color (str, optional): The color for the geometries. Defaults to None.
        edgecolor (str, optional): The color for the geometry edges. Defaults to 'black'.
        alpha (float, optional): The alpha blending value, between 0 and 1. Defaults to 1.0.
        output_path (Optional[str], optional): Path to save the map image file (e.g., 'map.png'). Defaults to None.
    
    Returns:
        Optional[str]: The path to the saved file if output_path is provided, otherwise None.
    """
    # Ensure the geodata is in the same CRS
    geodata = geodata.to_crs(epsg=2154)
    geodata.plot(ax=ax, color=color, edgecolor=edgecolor, alpha=alpha)

    if output_path:
        fig = ax.get_figure()
        fig.savefig(output_path, bbox_inches='tight')
    
    return output_path

@tool
def plot_departments_data(
    data: pd.DataFrame, 
    dep_col: str = 'dep', 
    value_col: str = 'value',
    map_title: str = 'French Departments Data',
    output_path: Optional[str] = 'france_data.png'
) -> Optional[str]:
    """
    Plots data for French departments on a map of France.
    
    Args:
        data (pd.DataFrame): DataFrame with department data. Must contain at least two columns:
                             one for department codes and one for the values to plot.
        dep_col (str): The name of the column in `data` that contains the department codes.
        value_col (str): The name of the column in `data` that contains the values to plot.
        map_title (str): The title of the map.
        output_path (Optional[str]): Path to save the map image file. If None, the plot is not saved.
                                     Defaults to 'france_data.png'.
                                     
    Returns:
        Optional[str]: The path to the saved file if output_path is provided, otherwise None.
    """
    # Get the geodata for departments
    departments_gdf = get_france_geodata('departments')
    
    # Ensure department codes in user data are strings for merging
    data[dep_col] = data[dep_col].astype(str).str.zfill(2)
    
    # Merge the geodata with the user's data
    merged_gdf = departments_gdf.merge(data, left_on='code', right_on=dep_col)
    
    # Create the plot
    fig, ax = plt.subplots(1, 1, figsize=(15, 12))
    ax.set_aspect('equal')
    ax.set_axis_off()
    
    # Plot the base map of all departments
    departments_gdf.plot(ax=ax, color='lightgray', edgecolor='black')
    
    # Plot the data on top
    if not merged_gdf.empty:
        merged_gdf.plot(column=value_col, ax=ax, legend=True, cmap='viridis')

    ax.set_title(map_title, fontsize=20)
    
    if output_path:
        fig.savefig(output_path, bbox_inches='tight')
        print(f"Map saved to {output_path}")
        return output_path
    
    return None

if __name__ == '__main__':
    # Create sample data for 5 departments as requested
    # The user provided 5, 92, 63, 45, 32
    sample_data = {
        'dep': [5, 92, 63, 45, 32],
        'value': [10, 50, 20, 30, 45] # Some arbitrary values
    }
    data_df = pd.DataFrame(sample_data)
    
    print("Generating map with department data...")
    plot_departments_data(data_df, output_path='france_departments_data.png')