import geopandas as gpd
import os
import requests

def create_directory(directory):
    """
    Creates a directory if it does not already exist.

    Parameters:
    - directory (str): The path of the directory to be created.

    Effects:
    - A new directory is created at the specified path if it does not exist. If the directory already exists, no action is taken.
    """
    # Check if the directory exists and create it if it doesn't
    if not os.path.exists(directory):
        os.makedirs(directory)

def show_available_datasets(geopackage_url):
    """
    Loads a GeoPackage and displays all unique countries available in the 'country' column of the dataset.

    Parameters:
    - geopackage_url (str): The URL to the GeoPackage file.

    Outputs:
    - Prints a list of unique countries found in the GeoPackage's 'country' column.
    """
    # Load the GeoPackage directly
    gdf = gpd.read_file(geopackage_url, layer='caribbean_watersheds')
    unique_countries = gdf['Country'].unique()
    print("Available countries in the dataset:")
    for Country in unique_countries:
        print(Country)

def download_country(country, output_dir=".//Watershed//", target_crs=None):
    """
    Downloads and saves watershed data for a specified country from a GeoPackage.

    Parameters:
    - country (str): The name of the country for which watershed data is to be downloaded.
    - output_dir (str): The directory where the GeoPackage will be saved. Defaults to ".//Watershed//".
    - target_crs (str, optional): The coordinate reference system to which the data should be converted. If None, the original CRS is used.

    Outputs:
    - A new GeoPackage file in the specified directory containing the watershed data for the specified country.
    - Console messages about the status of the download and data processing steps.

    Raises:
    - Prints an error message if the GeoPackage cannot be downloaded (e.g., due to network issues or an incorrect URL).
    """
    # URL of the GeoPackage hosted on GitHub
    geopackage_url = "https://github.com/kgdthomas/IslandSheds/raw/main/CARIBBEAN_WATERSHEDS.gpkg"
    
    # Download the GeoPackage
    response = requests.get(geopackage_url)
    
    if response.status_code == 200:
        # Ensure the output directory exists
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        
        # Save the GeoPackage locally
        local_path = os.path.join(output_dir, 'temp.gpkg')
        with open(local_path, 'wb') as f:
            f.write(response.content)
        
        # Load the GeoPackage from the local file
        gdf = gpd.read_file(local_path, layer='caribbean_watersheds')
        
        # Filter data for the specified country
        country_gdf = gdf[gdf['Country'] == country]
        
        # Check and inform user about CRS
        crs = gdf.crs
        print(f"The CRS of the data is {crs}")
        
        # Ask user if they want to change the CRS
        if target_crs:
            # Change CRS if target_crs is specified
            country_gdf = country_gdf.to_crs(target_crs)
            print(f"Data CRS changed to {target_crs}")
        else:
            print("Data will be downloaded with the original CRS.")
        
        # Save the filtered data to a new GeoPackage file
        output_file_path = os.path.join(output_dir, f"{country}_watersheds.gpkg")
        country_gdf.to_file(output_file_path, driver='GPKG')
        
        print(f"Data for {country} downloaded successfully.")
    else:
        print(f"Failed to download GeoPackage. Status code: {response.status_code}")

def download_watershed(geopackage_url, country, download_dir="//Watershed//", target_crs=None):
    """
    Downloads watershed data for a specified country and optionally a specific watershed.

    Parameters:
    - geopackage_url (str): The URL to the GeoPackage file.
    - country (str): The country for which watershed data is to be downloaded.
    - download_dir (str): The directory where the watershed data will be saved. Defaults to "//Watershed//".

    Outputs:
    - Downloads and saves the specified watershed data as a GPKG file in the specified directory.
    - Console messages about the status of the download.

    Raises:
    - Prints an error message if the data cannot be downloaded (e.g., due to network issues or an incorrect URL).
    """
    # Ensure the output directory exists
    if not os.path.exists(download_dir):
        os.makedirs(download_dir)

    # Load the GeoPackage
    gdf = gpd.read_file(geopackage_url, layer='caribbean_watersheds')
    
    # Filter data for the specified country
    country_gdf = gdf[gdf['Country'] == country]
    
    # Check and inform user about CRS
    crs = gdf.crs
    print(f"The CRS of the data is {crs}")
    
    # Display available watersheds in the country
    print(f"Available watersheds in {country}:")
    unique_watersheds = country_gdf['Watershed'].unique()
    for watershed in unique_watersheds:
        print(watershed)
    
    # Ask user which watershed to download
    watershed = input("Enter the watershed name to download (or press enter for all watersheds): ")
    
    # Filter for the specific watershed if provided
    if watershed:
        watershed_gdf = country_gdf[country_gdf['Watershed'] == watershed]
    else:
        watershed_gdf = country_gdf
        
    # Ask user if they want to change the CRS
    if target_crs:
        # Change CRS if target_crs is specified
        watershed_gdf = watershed_gdf.to_crs(target_crs)
        print(f"Data CRS changed to {target_crs}")
    else:
        print("Data will be downloaded with the original CRS.")
    
    # Save the filtered data to a GPKG file
    output_file_path = os.path.join(download_dir, f"{country}_{watershed or 'all_watersheds'}.gpkg")
    watershed_gdf.to_file(output_file_path, layer='data', driver='GPKG')
    
    print(f"Downloaded {watershed or 'all watersheds'} for {country} successfully.")
    
def main():
    """
    Main function to run the Watershed Downloader program.

    Interactively prompts the user to enter a country name, an optional watershed name, and a download directory. It then downloads the specified watershed data using the provided details.
    """
    # URL to the GeoPackage
    geopackage_url = "https://github.com/kgdthomas/IslandSheds/raw/main/CARIBBEAN_WATERSHEDS.gpkg"
    print("Welcome to the Watershed Downloader!")
    
    # Display available datasets
    show_available_datasets(geopackage_url)
    
    # Ask user if they want to download data for a country or a specific watershed
    choice = input("Would you like to download data for a (1) specific country or (2) specific watershed? Enter 1 or 2: ")
    
    if choice == '1':
        # Get user input for country and download directory
        country = input("Enter the country name: ")
        output_dir = input("Enter download directory (default is .//Watershed//): ")
        output_dir = output_dir if output_dir else ".//Watershed//"
        target_crs = input("Enter target CRS (optional): ")
        
        # Download the country data
        download_country(country, output_dir, target_crs)
    
    elif choice == '2':
        # Get user input for country, watershed, and download directory
        country = input("Enter the country name for the watershed: ")
        # watershed = input("Enter the watershed name (or press enter for all watersheds): ")
        download_dir = input("Enter download directory (default is .//Watershed//): ")
        download_dir = download_dir if download_dir else ".//Watershed//"
        target_crs = input("Enter target CRS (optional): ")
        
        # Download the watershed data
        download_watershed(geopackage_url, country, download_dir, target_crs)
    
    else:
        print("Invalid choice. Please restart the program and enter 1 or 2.")

if __name__ == "__main__":
    main()