"""
Data loaders for static classification and reference data.

This module provides functions to load CSV data files bundled with the package.
All functions return pandas DataFrames.

Example usage:
    from atlas_common_data import load_countries, load_products

    countries = load_countries()
    products = load_products(classification='hs92')
"""

from pathlib import Path
from typing import List, Optional

import pandas as pd

# Base path for data files
DATA_DIR = Path(__file__).parent / "data"


def _validate_data_dir(data_dir: Path) -> None:
    """Validate that the package data directory exists.

    Called at import time so a corrupt or incomplete installation is caught
    immediately with a clear message rather than surfacing later as a
    confusing FileNotFoundError or misleading ValueError.

    Raises:
        ImportError: If the data directory does not exist
    """
    if not data_dir.exists():
        raise ImportError(
            f"atlas-common-data package data not found at: {data_dir}\n"
            "This indicates a corrupt or incomplete installation.\n"
            "Try: pip install --force-reinstall atlas-common-data"
        )


_validate_data_dir(DATA_DIR)


def _get_data_path(subpath: str) -> Path:
    """Get the full path to a data file.

    Args:
        subpath: Relative path from DATA_DIR

    Returns:
        Full Path object

    Raises:
        FileNotFoundError: If the file doesn't exist
    """
    full_path = DATA_DIR / subpath
    if not full_path.exists():
        raise FileNotFoundError(f"Data file not found: {full_path}")
    return full_path


# =============================================================================
# Location Classification Loaders
# =============================================================================

def load_countries() -> pd.DataFrame:
    """Load country classification data.

    Returns:
        DataFrame with columns:
            - country_id: Unique identifier for the country
            - location_level: Type of location (always 'country')
            - name_en: English name
            - name_short_en: Short English name
            - name_es: Spanish name
            - name_short_es: Short Spanish name
            - iso3_code: ISO 3166-1 alpha-3 code
            - iso2_code: ISO 3166-1 alpha-2 code
            - legacy_location_id: Legacy identifier
            - the_prefix: Whether to use 'the' prefix
            - former_country: Whether country no longer exists
            - country_start_year: Year country started existing
            - country_end_year: Year country ceased existing (3000 if current)
            - in_complexity: Country included in base calculation for complexity metrics
            - in_cp: Country included in Atlas of Economic Complexity Country Profiles

    Example:
        >>> countries = load_countries()
        >>> usa = countries[countries['iso3_code'] == 'USA']
        >>> print(usa['name_en'].values[0])
        United States
    """
    path = _get_data_path("classification/location/country.csv")
    return pd.read_csv(path)


def load_groups() -> pd.DataFrame:
    """Load geographic group classification data.

    Returns:
        DataFrame with group definitions including:
            - group_id: Unique identifier
            - name_en: English name
            - group_type: Type of grouping (e.g., 'region', 'income')

    Example:
        >>> groups = load_groups()
        >>> print(groups.head())
    """
    path = _get_data_path("classification/location/group.csv")
    return pd.read_csv(path)


def load_group_members() -> pd.DataFrame:
    """Load country-to-group membership data.

    Returns:
        DataFrame mapping countries to their groups with columns:
            - country_id: Country identifier
            - group_id: Group identifier
            - Additional membership metadata

    Example:
        >>> members = load_group_members()
        >>> # Get all countries in a specific group
        >>> group_countries = members[members['group_id'] == 1]
    """
    path = _get_data_path("classification/location/group_member.csv")
    return pd.read_csv(path)


# =============================================================================
# Product Classification Loaders
# =============================================================================

def list_available_product_classifications() -> List[str]:
    """List available product classification systems.

    Returns:
        List of classification codes (e.g., ['hs92', 'hs12', 'hs22', 'sitc'])

    Example:
        >>> classifications = list_available_product_classifications()
        >>> print(classifications)
        ['hs12', 'hs22', 'hs92', 'sitc', 'services_bilateral', services_unilateral']
    """
    product_dir = DATA_DIR / "classification" / "product"
    if not product_dir.exists():
        return []

    # Find main classification files (not edges, clusters, or wip files)
    classification_files = []
    for f in product_dir.glob("*.csv"):
        name = f.stem
        # Skip auxiliary files
        if any(skip in name for skip in ['_ps_', '_wip', '_with_', '_need_']):
            continue
        classification_files.append(name)

    return sorted(classification_files)


def load_products(classification: str = "hs92") -> pd.DataFrame:
    """Load product classification data.

    Args:
        classification: Classification system to load. Options:
            - 'hs92': HS 1992 (default, most comprehensive)
            - 'hs12': HS 2012
            - 'hs22': HS 2022
            - 'sitc': SITC Revision 2

    Returns:
        DataFrame with product hierarchy including:
            - product_id: Unique identifier
            - name_en: English name
            - name_short_en: Short English name
            - code: Product code in the classification
            - parent_id: Parent product ID (for hierarchy)
            - product_level: Level in hierarchy (1=section, 2=chapter, etc.)
            - top_parent_id: Top-level parent ID
            - is_shown: Whether product is displayed in visualizations

    Raises:
        ValueError: If classification is not available

    Example:
        >>> products = load_products('hs92')
        >>> chapters = products[products['product_level'] == 2]
        >>> print(f"Found {len(chapters)} chapters")
    """
    available = list_available_product_classifications()
    if classification not in available:
        raise ValueError(
            f"Classification '{classification}' not available. "
            f"Available options: {available}"
        )

    path = _get_data_path(f"classification/product/{classification}.csv")
    return pd.read_csv(path)


def list_available_product_space_classifications() -> List[str]:
    """List classifications that have product space network data available.

    Returns:
        List of classification codes with product space edges (e.g., ['hs92'])

    Example:
        >>> classifications = list_available_product_space_classifications()
        >>> print(classifications)
        ['hs92']
    """
    product_dir = DATA_DIR / "classification" / "product"
    if not product_dir.exists():
        return []
    return sorted(
        f.stem.replace("_ps_edges", "")
        for f in product_dir.glob("*_ps_edges.csv")
    )


def load_product_space_edges(classification: str = "hs92") -> pd.DataFrame:
    """Load product space network edges.

    The product space is a network representation of product relatedness,
    where edges connect products that tend to be co-exported.

    Args:
        classification: Classification system. Use
            list_available_product_space_classifications() to see options.

    Returns:
        DataFrame with network edges including:
            - source: Source product ID
            - target: Target product ID
            - weight: Edge weight (proximity/relatedness)

    Raises:
        ValueError: If product space data is not available for the classification

    Example:
        >>> edges = load_product_space_edges()
        >>> print(f"Product space has {len(edges)} edges")
    """
    available = list_available_product_space_classifications()
    if classification not in available:
        raise ValueError(
            f"Product space edges not available for '{classification}'. "
            f"Available classifications: {available}"
        )
    path = _get_data_path(f"classification/product/{classification}_ps_edges.csv")
    return pd.read_csv(path)


def load_product_space_clusters(classification: str = "hs92") -> pd.DataFrame:
    """Load product space cluster definitions.

    Clusters group related products in the product space visualization.

    Args:
        classification: Classification system. Use
            list_available_product_space_classifications() to see options.

    Returns:
        DataFrame with cluster definitions

    Raises:
        ValueError: If product space data is not available for the classification
    """
    available = list_available_product_space_classifications()
    if classification not in available:
        raise ValueError(
            f"Product space clusters not available for '{classification}'. "
            f"Available classifications: {available}"
        )
    path = _get_data_path(f"classification/product/{classification}_ps_clusters.csv")
    return pd.read_csv(path)


def load_services(service_type: str = "bilateral") -> pd.DataFrame:
    """Load services trade classification data.

    Args:
        service_type: Type of services classification:
            - 'bilateral': Services traded between specific country pairs
            - 'unilateral': Services where partner is not specified

    Returns:
        DataFrame with service category definitions

    Raises:
        ValueError: If service_type is not 'bilateral' or 'unilateral'
    """
    if service_type not in ("bilateral", "unilateral"):
        raise ValueError(f"service_type must be 'bilateral' or 'unilateral', got '{service_type}'")

    path = _get_data_path(f"classification/product/services_{service_type}.csv")
    return pd.read_csv(path)


# =============================================================================
# Static Reference Data Loaders
# =============================================================================

def load_comparator_countries() -> pd.DataFrame:
    """Load comparator countries data.

    This dataset contains lists of comparator countries for each country,
    based on various similarity metrics.

    Returns:
        DataFrame with comparator country mappings including:
            - iso: ISO3 code of the reference country
            - idc: Comparator rank (1-5)
            - iso_score1, iso_score2: Comparators by different metrics
            - iso_distance: Comparator by geographic distance

    Example:
        >>> comparators = load_comparator_countries()
        >>> usa_comparators = comparators[comparators['iso'] == 'USA']
    """
    path = _get_data_path("static/comparator_countries.csv")
    return pd.read_csv(path)

def load_geo_distances() -> pd.DataFrame:
    """Load geographic distance matrix.

    Returns:
        DataFrame with pairwise distances between countries.

    Note:
        This is a large file (~4MB). Consider filtering after loading
        if you only need specific country pairs.
    """
    path = _get_data_path("static/geo_dist.csv")
    return pd.read_csv(path)


def load_sitc_natural_resource_products() -> pd.DataFrame:
    """Load SITC classification products classified as natural resources 

    Returns:
        DataFrame with SITC product codes with natural resource flag

    """
    path = _get_data_path("static/sitc_natural_resources.csv")
    return pd.read_csv(path, dtype={"sitc_product_code":str})



# =============================================================================
# Metadata Loaders
# =============================================================================

def load_column_metadata() -> pd.DataFrame:
    """Load column metadata/data dictionary.

    This dataset describes the columns used across various data exports.

    Returns:
        DataFrame with column definitions including:
            - column_id: Unique column identifier
            - column_name: Name of the column
            - data_type: Data type (integer, string, etc.)
            - description: Human-readable description
            - notes: Additional notes
            - complexity: Complexity indicator
            - decimal_places: Number of decimal places (if numeric)

    Example:
        >>> metadata = load_column_metadata()
        >>> gdp_cols = metadata[metadata['column_name'].str.contains('gdp', case=False)]
    """
    path = _get_data_path("downloads/column_metadata.csv")
    return pd.read_csv(path)


# =============================================================================
# Convenience Functions
# =============================================================================

def get_country_by_iso3(iso3_code: str) -> Optional[pd.Series]:
    """Get a single country's data by ISO3 code.

    Args:
        iso3_code: ISO 3166-1 alpha-3 code (e.g., 'USA', 'CHN')

    Returns:
        Series with country data, or None if not found

    Example:
        >>> usa = get_country_by_iso3('USA')
        >>> print(usa['name_en'])
        United States
    """
    if not isinstance(iso3_code, str):
        raise TypeError(f"iso3_code must be a string, got {type(iso3_code).__name__!r}")
    countries = load_countries()
    matches = countries[countries['iso3_code'] == iso3_code.upper()]
    if len(matches) == 0:
        return None
    return matches.iloc[0]


def get_country_id_mapping() -> dict:
    """Get a mapping from ISO3 codes to country IDs.

    Returns:
        Dictionary mapping ISO3 codes to country_id values

    Example:
        >>> mapping = get_country_id_mapping()
        >>> print(mapping['USA'])
        840
    """
    countries = load_countries()
    return dict(zip(countries['iso3_code'], countries['country_id']))


def get_iso3_to_name_mapping() -> dict:
    """Get a mapping from ISO3 codes to country names.

    Returns:
        Dictionary mapping ISO3 codes to English names

    Example:
        >>> names = get_iso3_to_name_mapping()
        >>> print(names['USA'])
        United States
    """
    countries = load_countries()
    return dict(zip(countries['iso3_code'], countries['name_en']))


# =============================================================================
# Growth Projections Loaders
# =============================================================================

def list_available_growth_projection_years() -> List[int]:
    """List years for which growth projection data is available.

    Dynamically scans the data directory at call time, so new files are
    automatically discovered without code changes.

    Returns:
        Sorted list of available years as integers (e.g., [2024, 2025])

    Example:
        >>> years = list_available_growth_projection_years()
        >>> print(years)
        [2024]
    """
    gp_dir = DATA_DIR / "growth_projections"
    if not gp_dir.exists():
        return []
    return sorted(
        int(f.stem.replace("_growth_projections", ""))
        for f in gp_dir.glob("*_growth_projections.csv")
    )


def load_growth_projections(year) -> pd.DataFrame:
    """Load growth projections data.

    Args:
        year: Year of projections (integer). Use
            list_available_growth_projection_years() to see available options.

    Returns:
        DataFrame with columns: iso3_code, growth_proj, year

    Raises:
        ValueError: If no data is available for the given year

    Example:
        >>> df = load_growth_projections(2024)
        >>> usa = df[df['iso3_code'] == 'USA']
    """
    available = list_available_growth_projection_years()
    try:
        year_int = int(year)
    except (ValueError, TypeError):
        raise ValueError(
            f"Growth projections for year '{year}' not available. "
            f"Available years: {available}"
        )
    if year_int not in available:
        raise ValueError(
            f"Growth projections for year '{year}' not available. "
            f"Available years: {available}"
        )
    path = _get_data_path(f"growth_projections/{year_int}_growth_projections.csv")
    return pd.read_csv(path)


# =============================================================================
# Classification Constants and Conversion Weights
# =============================================================================

def load_classifications() -> dict:
    """Load trade classification metadata.

    Returns:
        Dictionary mapping classification codes to their metadata:
            - name: Display name (e.g. 'HS92')
            - release_year: Year the classification was released
            - description: Human-readable description

    Example:
        >>> classifications = load_classifications()
        >>> print(classifications['H0'])
        {'name': 'HS92', 'release_year': 1995, 'description': 'HS92'}
    """
    from atlas_common_data.constants.classifications import CLASSIFICATIONS
    return CLASSIFICATIONS.copy()


def list_available_conversion_weights() -> List[str]:
    """List available classification conversion weight pairs.

    Returns:
        List of conversion pair strings (e.g. ['H0_to_H1', 'H1_to_H0', ...])

    Example:
        >>> pairs = list_available_conversion_weights()
        >>> print(pairs[:3])
        ['H0_to_H1', 'H0_to_S3', 'H1_to_H0']
    """
    cw_dir = DATA_DIR / "conversion_weights"
    if not cw_dir.exists():
        return []
    return sorted(
        p.stem.replace("conversion_weights_", "")
        for p in cw_dir.glob("conversion_weights_*.csv")
    )


def load_conversion_weights(source: str, target: str) -> pd.DataFrame:
    """Load product code conversion weights between two classifications.

    Args:
        source: Source classification code (e.g. 'H0', 'S2')
        target: Target classification code (e.g. 'H1', 'S3')

    Returns:
        DataFrame with columns: source_code, target_code, weight, group_id

    Raises:
        FileNotFoundError: If the conversion pair is not available.
            Use list_available_conversion_weights() to see available pairs.

    Example:
        >>> df = load_conversion_weights('H0', 'H1')
        >>> print(df.head())
    """
    filename = f"conversion_weights/conversion_weights_{source}_to_{target}.csv"
    try:
        path = _get_data_path(filename)
    except FileNotFoundError:
        available = list_available_conversion_weights()
        raise ValueError(
            f"Conversion weights for {source} -> {target} not found. "
            f"Available pairs: {available}"
        )
    return pd.read_csv(path, dtype={target: str, source: str}).rename(columns={"weight": f"weight_{source}_{target}"})
