import pandas as pd
import geopandas as gpd
from pathlib import Path
from .metadata import FAFMetadata
from .download import download_and_process
from .visualization import FlowMap

class FAFRegionData:
    """
    Class to load and interact with the FAF5 Regional Database.
    """
    def __init__(self, data_dir=None):
        if data_dir is None:
            # Default to user home directory
            self.data_dir = Path.home() / ".tidyfaf_data"
            if not self.data_dir.exists():
                # Try to create it if it doesn't exist, though download() should handle this
                # But here we are just reading.
                pass
        else:
            self.data_dir = Path(data_dir)
            
        self.filename = "FAF5.7.1.parquet"
        self.filepath = self.data_dir / self.filename
        
        if not self.filepath.exists():
            # Fallback to CSV if Parquet doesn't exist
            self.filename = "FAF5.7.1.csv"
            self.filepath = self.data_dir / self.filename
            
        if not self.filepath.exists():
            raise FileNotFoundError(f"FAF Regional Database not found at {self.filepath}. Please run download_and_process() first.")
            
        print(f"Loading FAF Regional Database from {self.filepath}...")
        if self.filepath.suffix == '.parquet':
             self.df = pd.read_parquet(self.filepath)
        else:
             self.df = pd.read_csv(self.filepath, low_memory=False)
             
        self.metadata = FAFMetadata(self.data_dir)
        print("Data loaded successfully.")

    def get_flows(self, origin=None, destination=None, commodity=None, mode=None):
        """
        Filter flows based on origin, destination, commodity, and mode.
        Accepts names (strings) or codes (integers).
        """
        filtered_df = self.df.copy()
        
        if origin:
            if isinstance(origin, str) and not origin.isdigit():
                origin = self.metadata.lookup_state(origin)
            filtered_df = filtered_df[filtered_df['dms_orig'] == int(origin)]
        
        if destination:
            if isinstance(destination, str) and not destination.isdigit():
                destination = self.metadata.lookup_state(destination)
            filtered_df = filtered_df[filtered_df['dms_dest'] == int(destination)]
            
        if commodity:
            if isinstance(commodity, str) and not commodity.isdigit():
                commodity = self.metadata.lookup_commodity(commodity)
            filtered_df = filtered_df[filtered_df['sctg2'] == int(commodity)]
            
        if mode:
            if isinstance(mode, str) and not mode.isdigit():
                mode = self.metadata.lookup_mode(mode)
            filtered_df = filtered_df[filtered_df['dms_mode'] == int(mode)]
            
        return filtered_df

    def get_zones_gdf(self):
        """
        Returns a GeoDataFrame of FAF Zones with their geometries.
        Processes the raw county-level shapefile, dissolves by CFS Area, 
        and joins with metadata to get FAF Zone IDs.
        Caches the result.
        """
        cache_path = self.data_dir / "FAF5_Zones_Processed.parquet"
        if cache_path.exists():
            print(f"Loading processed zones from {cache_path}...")
            return gpd.read_parquet(cache_path)
            
        raw_zones_path = self.data_dir / "FAF5_Zones.parquet"
        if not raw_zones_path.exists():
             raise FileNotFoundError("FAF5_Zones.parquet not found. Run download_and_process() first.")
             
        print("Processing FAF Zones (dissolving and joining metadata)...")
        gdf_raw = gpd.read_parquet(raw_zones_path)
        
        # Load Metadata
        meta_df = self.metadata.get_table("FAF Zone (Domestic)")
        
        # Prepare for merge
        # Shapefile has 'CFS17_NAME'. Metadata has 'Long Description'.
        # We need to normalize strings to ensure match.
        
        # Helper to clean names
        def clean_name(name):
            if not isinstance(name, str): return ""
            return name.lower().strip().replace(" cfs area", "").replace(" (ct part)", "").replace(" (de part)", "").replace(" (md part)", "").replace(" (nj part)", "").replace(" (ny part)", "").replace(" (pa part)", "")

        gdf_raw['clean_name'] = gdf_raw['CFS17_NAME'].apply(clean_name)
        meta_df['clean_name'] = meta_df['Long Description'].apply(clean_name)
        
        # Merge
        # Note: 'Numeric Label' is the FAF ID
        merged = gdf_raw.merge(meta_df[['Numeric Label', 'clean_name']], on='clean_name', how='left')
        
        # Check for unmerged
        unmerged = merged[merged['Numeric Label'].isna()]['CFS17_NAME'].unique()
        if len(unmerged) > 0:
            print(f"Warning: Could not match {len(unmerged)} areas to FAF IDs. Examples: {unmerged[:5]}")
            # Fallback: Try matching 'Short Description' or other heuristics if needed.
            # For 'Rest of XX', the shapefile might say 'Remainder of XX'.
            
        # Drop rows without ID
        merged = merged.dropna(subset=['Numeric Label'])
        merged['FAFZONE'] = merged['Numeric Label'].astype(int)
        
        # Dissolve by FAFZONE
        print("Dissolving geometries by FAFZONE...")
        gdf_dissolved = merged.dissolve(by='FAFZONE')
        
        # Keep only necessary columns
        gdf_final = gdf_dissolved[['geometry']].reset_index()
        
        # Save cache
        print(f"Saving processed zones to {cache_path}...")
        gdf_final.to_parquet(cache_path)
        
        return gdf_final

class FAFHiLoData:
    """
    Class to load and interact with the FAF5 HiLo Forecasts.
    """
    def __init__(self, data_dir=None):
        if data_dir is None:
            self.data_dir = Path(__file__).resolve().parent.parent.parent / "data"
        else:
            self.data_dir = Path(data_dir)
            
        self.filename = "FAF5.7.1_HiLoForecasts.parquet"
        self.filepath = self.data_dir / self.filename
        
        if not self.filepath.exists():
            self.filename = "FAF5.7.1_HiLoForecasts.csv"
            self.filepath = self.data_dir / self.filename

        if not self.filepath.exists():
            raise FileNotFoundError(f"FAF HiLo Forecasts not found at {self.filepath}. Please run download_and_process() first.")
            
        print(f"Loading FAF HiLo Forecasts from {self.filepath}...")
        if self.filepath.suffix == '.parquet':
             self.df = pd.read_parquet(self.filepath)
        else:
             self.df = pd.read_csv(self.filepath, low_memory=False)
             
        self.metadata = FAFMetadata(self.data_dir)
        print("Data loaded successfully.")
        
    def get_forecasts(self, origin=None, destination=None, commodity=None):
        """
        Filter forecasts based on origin, destination, and commodity.
        Accepts names (strings) or codes (integers).
        """
        filtered_df = self.df.copy()
        
        if origin:
            if isinstance(origin, str) and not origin.isdigit():
                origin = self.metadata.lookup_state(origin)
            filtered_df = filtered_df[filtered_df['dms_orig'] == int(origin)]
        
        if destination:
            if isinstance(destination, str) and not destination.isdigit():
                destination = self.metadata.lookup_state(destination)
            filtered_df = filtered_df[filtered_df['dms_dest'] == int(destination)]
            
        if commodity:
            if isinstance(commodity, str) and not commodity.isdigit():
                commodity = self.metadata.lookup_commodity(commodity)
            filtered_df = filtered_df[filtered_df['sctg2'] == int(commodity)]
            
        return filtered_df

class FAFStateData:
    """
    Class to load and interact with the FAF5 State Database.
    """
    def __init__(self, data_dir=None):
        if data_dir is None:
            self.data_dir = Path(__file__).resolve().parent.parent.parent / "data"
        else:
            self.data_dir = Path(data_dir)
            
        self.filename = "FAF5.7.1_State.parquet"
        self.filepath = self.data_dir / self.filename
        
        if not self.filepath.exists():
            self.filename = "FAF5.7.1_State.csv"
            self.filepath = self.data_dir / self.filename

        if not self.filepath.exists():
            # Fallback check if the filename inside zip is different, e.g. just "FAF5.7.1_State.csv"
            # Let's assume it is standard for now.
            raise FileNotFoundError(f"FAF State Database not found at {self.filepath}. Please run download_and_process() first.")
            
        print(f"Loading FAF State Database from {self.filepath}...")
        if self.filepath.suffix == '.parquet':
             self.df = pd.read_parquet(self.filepath)
        else:
             self.df = pd.read_csv(self.filepath, low_memory=False)
             
        self.metadata = FAFMetadata(self.data_dir)
        print("Data loaded successfully.")
        
    def get_flows(self, origin=None, destination=None, commodity=None, mode=None):
        """
        Filter flows based on origin, destination, commodity, and mode.
        Note: Origin and destination for state data are 'dms_origst' and 'dms_destst'.
        Accepts names (strings) or codes (integers).
        """
        filtered_df = self.df.copy()
        
        if origin:
            if isinstance(origin, str) and not origin.isdigit():
                origin = self.metadata.lookup_state(origin)
            filtered_df = filtered_df[filtered_df['dms_origst'] == int(origin)]
        
        if destination:
            if isinstance(destination, str) and not destination.isdigit():
                destination = self.metadata.lookup_state(destination)
            filtered_df = filtered_df[filtered_df['dms_destst'] == int(destination)]
            
        if commodity:
            if isinstance(commodity, str) and not commodity.isdigit():
                commodity = self.metadata.lookup_commodity(commodity)
            filtered_df = filtered_df[filtered_df['sctg2'] == int(commodity)]

        if mode:
            if isinstance(mode, str) and not mode.isdigit():
                mode = self.metadata.lookup_mode(mode)
            filtered_df = filtered_df[filtered_df['dms_mode'] == int(mode)]
            
        return filtered_df

class FAFNetworkData:
    """
    Class to interact with the FAF5 Network Database.
    Currently provides paths to the network files.
    """
    def __init__(self, data_dir=None):
        if data_dir is None:
            self.data_dir = Path(__file__).resolve().parent.parent.parent / "data"
        else:
            self.data_dir = Path(data_dir)
            
        self.network_dir = self.data_dir / "Networks"
        
        if not self.network_dir.exists():
             raise FileNotFoundError(f"FAF Network directory not found at {self.network_dir}. Please run download_and_process() first.")

    def get_network_gdf(self, layer='links'):
        """
        Returns the network as a GeoDataFrame.
        Prioritizes the GeoParquet file if it exists.
        
        Args:
            layer (str): 'links' or 'nodes'. Defaults to 'links'.
        """
        layer_map = {
            'links': 'Links',
            'nodes': 'Nodes'
        }
        
        if layer.lower() not in layer_map:
             raise ValueError(f"Invalid layer '{layer}'. Options are: {list(layer_map.keys())}")
             
        suffix = layer_map[layer.lower()]
        parquet_path = self.data_dir / f"FAF5_Network_{suffix}.parquet"
        
        # Fallback for legacy single file if it exists and user wants links
        legacy_path = self.data_dir / "FAF5_Network.parquet"
        
        if parquet_path.exists():
            print(f"Loading network {layer} from {parquet_path}...")
            return gpd.read_parquet(parquet_path)
        elif layer.lower() == 'links' and legacy_path.exists():
             print(f"Loading network links from legacy {legacy_path}...")
             return gpd.read_parquet(legacy_path)
        
        gdb_path = self.get_gdb_path()
        if gdb_path:
            print(f"Loading network {layer} from {gdb_path} (this may be slow)...")
            # Map simple layer name to GDB layer name
            gdb_layer_map = {
                'links': 'FAF5_Links',
                'nodes': 'FAF5_Nodes'
            }
            return gpd.read_file(gdb_path, layer=gdb_layer_map[layer.lower()])
            
        raise FileNotFoundError(f"No network data found for layer '{layer}' (checked Parquet and GDB).")



    def get_transcad_path(self):
        """Returns the path to the TransCAD files."""
        transcad_path = self.network_dir / "TransCAD Format"
        if transcad_path.exists():
            return transcad_path
        return None
