"""
World Bank World Development Indicators (WDI) data fetching and processing.

This module provides the WDIData class for fetching and processing
data from the World Bank Open Data API.

Example usage:
    from atlas_common_data.sources.wdi import WDIData

    wdi = WDIData()

    # Get specific indicators
    df = wdi.get_indicators(['NY.GDP.MKTP.CD', 'SP.POP.TOTL'])

    # Get GDP data specifically
    df = wdi.get_gdp_data()

    # Get services trade data
    df = wdi.get_services_data()
"""

from pathlib import Path
from typing import Dict, List, Optional, Tuple, Union
import warnings

import pandas as pd


class WDIData:
    """Fetch and process World Bank Development Indicators data.

    This class provides methods to query the World Bank API and process
    the returned data into analysis-ready DataFrames. Unlike the legacy
    implementation, this class returns DataFrames without writing files to disk.

    Attributes:
        WDI_INDICATORS: Default economic indicators mapping
        WDI_SERVICE_INDICATORS: Services trade indicators mapping
        INCOME_LEVEL_MAPPING: World Bank income level to standard names

    Example:
        >>> wdi = WDIData()
        >>> df = wdi.get_indicators(['NY.GDP.MKTP.CD'])
        >>> print(df.columns.tolist())
        ['iso3_code', 'year', 'gdp', 'country_id']
    """

    # Default economic indicators: WDI code -> friendly name
    WDI_INDICATORS: Dict[str, str] = {
        # GDP Current $
        "NY.GDP.MKTP.CD": "gdp",
        "NY.GDP.PCAP.CD": "gdppc",
        "NY.GDP.MKTP.PP.CD": "gdp_ppp",
        "NY.GDP.PCAP.PP.CD": "gdppc_ppp",
        # GDP Constant $
        "NY.GDP.MKTP.KD": "gdp_const",
        "NY.GDP.PCAP.KD": "gdppc_const",
        "NY.GDP.MKTP.PP.KD": "gdp_ppp_const",
        "NY.GDP.PCAP.PP.KD": "gdppc_ppp_const",
        # Population
        "SP.POP.TOTL": "population",
        # Balance of Payments
        "BN.CAB.XOKA.CD": "current_account",
        "BX.GSR.MRCH.CD": "exports_goods_bop",
        "BM.GSR.MRCH.CD": "imports_goods_bop",
    }

    # Services trade indicators: WDI code -> friendly name
    WDI_SERVICE_INDICATORS: Dict[str, str] = {
        # Total Services
        "BX.GSR.NFSV.CD": "services_export_value",
        "BM.GSR.NFSV.CD": "services_import_value",
        # Travel & Tourism
        "BX.GSR.TRVL.ZS": "travel_export_share",
        "BM.GSR.TRVL.ZS": "travel_import_share",
        # Insurance & Financial
        "BX.GSR.INSF.ZS": "finance_export_share",
        "BM.GSR.INSF.ZS": "finance_import_share",
        # Transport
        "BX.GSR.TRAN.ZS": "transport_export_share",
        "BM.GSR.TRAN.ZS": "transport_import_share",
        # Communications & Computer
        "BM.GSR.CMCP.ZS": "comms_import_share",
        "BX.GSR.CMCP.ZS": "comms_export_share",
    }

    WDI_NUMERIC_COLS = [
    "gdp", "gdp_const", "gdp_ppp", "gdp_ppp_const",
    "gdppc", "gdppc_const", "gdppc_ppp", "gdppc_ppp_const",
    "current_account", "exports_goods_bop", "imports_goods_bop",
    "population",
    ]

    # Income level mapping from World Bank names to standard
    INCOME_LEVEL_MAPPING: Dict[str, str] = {
        "High income": "high",
        "Low income": "low",
        "Lower middle income": "lower middle",
        "Upper middle income": "upper middle",
    }

    # World Bank API base URL
    API_BASE_URL = "https://api.worldbank.org/v2"

    def __init__(
        self,
        latest_year: Optional[int] = None,
    ):
        """Initialize WDIData.

        Args:
            latest_year: Default end year for all queries. Methods will use
                this when end_year is not explicitly passed.
        """
        self.latest_year = latest_year
        self._classification_dir = (
                Path(__file__).parent.parent / "data" / "classification"
            )

    def _build_api_url(self, indicators: List[str]) -> str:
        """Build the World Bank API URL for given indicators.

        Args:
            indicators: List of WDI indicator codes

        Returns:
            Full API URL string
        """
        indicator_string = ";".join(indicators)
        return (
            f"{self.API_BASE_URL}/countries/all/indicator/{indicator_string}"
            "?source=2&downloadformat=excel&dataformat=list"
        )

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

        Returns:
            DataFrame with country_id and iso3_code columns
        """
        path = self._classification_dir / "location" / "country.csv"
        return pd.read_csv(path)[["country_id", "iso3_code"]]

    def _load_in_country_profiles(self) -> set:
        """Load the set of iso3 codes for countries that are in country profiles

        Returns:
            Set of iso3_code strings where in_cp is True.
        """
        path = self._classification_dir.parent / "classification" / "location" / "country.csv"
        df = pd.read_csv(path)
        return set(df.loc[df["in_cp"] == True, "iso3_code"])

    def _merge_country_id(self, df: pd.DataFrame) -> pd.DataFrame:
        """Merge Atlas country_id into DataFrame.

        Args:
            df: DataFrame with iso3_code column

        Returns:
            DataFrame with country_id column added
        """
        countries = self._load_country_classification()
        return df.merge(countries, on="iso3_code", how="inner")

    def _fetch_excel_data(self, url: str) -> Dict[str, pd.DataFrame]:
        """Fetch Excel data from World Bank API.

        Args:
            url: World Bank API URL

        Returns:
            Dictionary of sheet name -> DataFrame
        """
        try:
            return pd.read_excel(
                url,
                sheet_name=None,
                header=None,
                names=[
                    "country_name",
                    "code",
                    "indicator_name",
                    "indicator",
                    "year",
                    "value",
                ],
            )
        except Exception as e:
            raise RuntimeError(
                f"Failed to fetch World Bank data from API: {e}"
            ) from e

    def _process_data_tabs(self, tabs: Dict[str, pd.DataFrame]) -> pd.DataFrame:
        """Process data tabs from World Bank Excel download.

        Args:
            tabs: Dictionary of sheet name -> DataFrame

        Returns:
            Combined DataFrame from all data tabs
        """
        dfs = []

        for name, df in tabs.items():
            if name == "Data":
                # First four rows of first tab are info about download
                df = df.drop(index=[0, 1, 2, 3])
                dfs.append(df)
            elif name.startswith("Data"):
                # Additional data tabs (Data (2), Data (3), etc.)
                # Drop their header row in favor of ours passed in read_excel
                df = df.drop(index=0)
                dfs.append(df)

        if not dfs:
            return pd.DataFrame()

        return pd.concat(dfs, ignore_index=True)

    def _process_metadata_tab(
        self, tabs: Dict[str, pd.DataFrame]
    ) -> Optional[pd.DataFrame]:
        """Process metadata tab to extract country info.

        Args:
            tabs: Dictionary of sheet name -> DataFrame

        Returns:
            DataFrame with iso3_code, region, incomelevel_enum, country_id
            or None if metadata tab not found
        """
        if "Metadata - Countries" not in tabs:
            return None

        df = tabs["Metadata - Countries"].copy()

        # Drop first row (header) and unnecessary columns
        df = df.drop(index=0)
        df = df.drop(columns=["year", "indicator", "value"], errors="ignore")
        df.columns = ["iso3_code", "region", "incomelevel_enum"]

        # Map income levels to standard names
        df["incomelevel_enum"] = df["incomelevel_enum"].map(self.INCOME_LEVEL_MAPPING)

        # Merge country_id
        df = self._merge_country_id(df)

        return df

    def _pivot_and_rename(
        self,
        df: pd.DataFrame,
        indicator_mapping: Dict[str, str],
    ) -> pd.DataFrame:
        """Pivot data to wide format and rename columns.

        Args:
            df: Long-format DataFrame with indicator, year, value columns
            indicator_mapping: Mapping of WDI codes to friendly names

        Returns:
            Wide-format DataFrame with one column per indicator
        """
        df = df.rename(columns={"code": "iso3_code"})
        df = df[["iso3_code", "indicator", "year", "value"]]

        # Map indicator codes to friendly names
        df["indicator"] = df["indicator"].map(indicator_mapping)

        # Pivot to wide format
        df = df.set_index(["iso3_code", "year"]).pivot(columns="indicator")
        df.columns = df.columns.droplevel()
        df = df.reset_index()

        return df

    def get_indicators(
        self,
        indicators: Optional[Dict[str, str]] = None,
        start_year: Optional[int] = None,
        end_year: Optional[int] = None,
        include_country_id: bool = True,
    ) -> pd.DataFrame:
        """Fetch WDI indicators and return as DataFrame.

        This is the primary method for fetching WDI data. It queries the API,
        processes the Excel response, and returns a clean DataFrame.

        Args:
            indicators: Dictionary mapping WDI indicator codes to friendly names.
                If None, uses WDI_INDICATORS (default economic indicators).
                Example: {'NY.GDP.MKTP.CD': 'gdp', 'SP.POP.TOTL': 'population'}
            start_year: Optional minimum year to include
            end_year: Optional maximum year to include
            include_country_id: If True, merge Atlas country_id from classification

        Returns:
            DataFrame with columns:
                - iso3_code: ISO 3166-1 alpha-3 country code
                - year: Year of observation
                - country_id: Atlas country ID (if include_country_id=True)
                - Plus one column per indicator (using friendly names)

        Example:
            >>> wdi = WDIData()
            >>> df = wdi.get_indicators(
            ...     indicators={'NY.GDP.MKTP.CD': 'gdp'},
            ...     start_year=2015,
            ...     end_year=2023
            ... )
        """
        if indicators is None:
            indicators = self.WDI_INDICATORS.copy()

        if not indicators:
            return pd.DataFrame()

        # Apply latest_year default
        if end_year is None and self.latest_year is not None:
            end_year = self.latest_year

        # World Bank Excel API rejects requests with more than 8 indicators.
        # Split into batches and merge results.
        MAX_INDICATORS_PER_REQUEST = 8
        indicator_items = list(indicators.items())
        batches = [
            dict(indicator_items[i : i + MAX_INDICATORS_PER_REQUEST])
            for i in range(0, len(indicator_items), MAX_INDICATORS_PER_REQUEST)
        ]

        batch_dfs = []
        for batch in batches:
            url = self._build_api_url(list(batch.keys()))
            tabs = self._fetch_excel_data(url)
            batch_df = self._process_data_tabs(tabs)
            if not batch_df.empty:
                batch_df = self._pivot_and_rename(batch_df, batch)
                batch_dfs.append(batch_df)

        if not batch_dfs:
            return pd.DataFrame()

        # Merge all batches on iso3_code + year
        df = batch_dfs[0]
        for other in batch_dfs[1:]:
            df = df.merge(other, on=["iso3_code", "year"], how="outer")

        # Convert year to int
        df["year"] = pd.to_numeric(df["year"], errors="coerce")
        df = df.dropna(subset=["year"])
        df["year"] = df["year"].astype(int)

        # Filter by year range if specified
        if start_year is not None:
            df = df[df["year"] >= start_year]
        if end_year is not None:
            df = df[df["year"] <= end_year]

        # Sort by country and year
        df = df.sort_values(["iso3_code", "year"])

        # Optionally merge country_id
        if include_country_id:
            df = self._merge_country_id(df)

        return df

    def get_country_metadata(self) -> pd.DataFrame:
        """Fetch country metadata including income level and region.

        This method fetches a minimal set of indicators just to get access
        to the metadata tab, which contains income level and region info.

        Returns:
            DataFrame with columns:
                - iso3_code: Country code
                - region: World Bank region
                - incomelevel_enum: Income level classification
                - country_id: Atlas country ID
        """
        # Fetch minimal data just to get metadata
        url = self._build_api_url(["SP.POP.TOTL"])
        tabs = self._fetch_excel_data(url)

        metadata = self._process_metadata_tab(tabs)

        if metadata is None:
            return pd.DataFrame()

        return metadata

    def get_gdp_data(
        self,
        start_year: Optional[int] = None,
        end_year: Optional[int] = None,
    ) -> pd.DataFrame:
        """Convenience method to fetch GDP-related indicators.

        Args:
            start_year: Optional minimum year
            end_year: Optional maximum year

        Returns:
            DataFrame with GDP indicators:
                - gdp: GDP, current USD
                - gdppc: GDP per capita, current USD
                - gdp_ppp: GDP, PPP
                - gdppc_ppp: GDP per capita, PPP
                - gdp_const: GDP, constant USD
                - gdppc_const: GDP per capita, constant USD
                - gdp_ppp_const: GDP, PPP constant
                - gdppc_ppp_const: GDP per capita, PPP constant
        """
        gdp_indicators = {
            "NY.GDP.MKTP.CD": "gdp",
            "NY.GDP.PCAP.CD": "gdppc",
            "NY.GDP.MKTP.PP.CD": "gdp_ppp",
            "NY.GDP.PCAP.PP.CD": "gdppc_ppp",
            "NY.GDP.MKTP.KD": "gdp_const",
            "NY.GDP.PCAP.KD": "gdppc_const",
            "NY.GDP.MKTP.PP.KD": "gdp_ppp_const",
            "NY.GDP.PCAP.PP.KD": "gdppc_ppp_const",
        }
        return self.get_indicators(
            indicators=gdp_indicators,
            start_year=start_year,
            end_year=end_year,
        )

    def get_services_data(
        self,
        start_year: Optional[int] = None,
        end_year: Optional[int] = None,
    ) -> pd.DataFrame:
        """Convenience method to fetch services trade indicators.

        Args:
            start_year: Optional minimum year
            end_year: Optional maximum year

        Returns:
            DataFrame with services trade indicators:
                - services_export_value: Total services exports
                - services_import_value: Total services imports
                - travel_export_share: Travel exports as % of services
                - travel_import_share: Travel imports as % of services
                - finance_export_share: Financial exports as % of services
                - finance_import_share: Financial imports as % of services
                - transport_export_share: Transport exports as % of services
                - transport_import_share: Transport imports as % of services
                - comms_export_share: Communications exports as % of services
                - comms_import_share: Communications imports as % of services
        """
        return self.get_indicators(
            indicators=self.WDI_SERVICE_INDICATORS,
            start_year=start_year,
            end_year=end_year,
        )

    def get_bop_data(
        self,
        start_year: Optional[int] = None,
        end_year: Optional[int] = None,
    ) -> pd.DataFrame:
        """Convenience method to fetch Balance of Payments indicators.

        Args:
            start_year: Optional minimum year
            end_year: Optional maximum year

        Returns:
            DataFrame with BOP indicators:
                - current_account: Current account balance
                - exports_goods_bop: Goods exports (BOP)
                - imports_goods_bop: Goods imports (BOP)
        """
        bop_indicators = {
            "BN.CAB.XOKA.CD": "current_account",
            "BX.GSR.MRCH.CD": "exports_goods_bop",
            "BM.GSR.MRCH.CD": "imports_goods_bop",
        }
        return self.get_indicators(
            indicators=bop_indicators,
            start_year=start_year,
            end_year=end_year,
        )

    def get_wdi_data(self, end_year: Optional[int] = None) -> pd.DataFrame:
        """Fetch all WDI indicators — matches schema of wdi_data.csv.

        This is the primary method whose output matches the legacy CSV file.
        Equivalent to get_indicators() with default indicators.

        Args:
            end_year: Optional maximum year to include. Falls back to latest_year.

        Returns:
            DataFrame with columns matching wdi_data.csv:
                iso3_code, year, current_account, exports_goods_bop, gdp,
                gdp_const, gdp_ppp, gdp_ppp_const, gdppc, gdppc_const,
                gdppc_ppp, gdppc_ppp_const, imports_goods_bop, population,
                country_id
        """
        df = self.get_indicators(end_year=end_year)

        if df.empty:
            return df

        # Warn if any countries are missing data for the requested end_year
        resolved_end_year = end_year if end_year is not None else self.latest_year
        if resolved_end_year is not None:
            all_countries = set(df["iso3_code"].unique())
            countries_in_latest = set(
                df[df["year"] == resolved_end_year]["iso3_code"].unique()
            )
            missing = sorted(all_countries - countries_in_latest)
            if missing:
                in_country_profiles = self._load_in_country_profiles()
                missing_ranked = sorted(c for c in missing if c in in_country_profiles)
                if missing_ranked:
                    warnings.warn(
                        f"{len(missing_ranked)} ranked countries have no WDI data for {resolved_end_year} "
                        f"and may not have reported yet: {missing_ranked}",
                        UserWarning,
                        stacklevel=2,
                    )
        df[self.WDI_NUMERIC_COLS] = df[self.WDI_NUMERIC_COLS].apply(pd.to_numeric, errors="coerce")
        return df

    @classmethod
    def list_available_indicators(cls) -> Dict[str, Dict[str, str]]:
        """List available WDI indicator presets.

        Returns:
            Dictionary with indicator categories and their mappings
        """
        return {
            "economic": cls.WDI_INDICATORS.copy(),
            "services": cls.WDI_SERVICE_INDICATORS.copy(),
        }

    @classmethod
    def get_indicator_description(cls, indicator_code: str) -> Optional[str]:
        """Get the friendly name for a WDI indicator code.

        Args:
            indicator_code: WDI indicator code ('NY.GDP.MKTP.PP.CD')

        Returns:
            Friendly name if found, None otherwise (gdp_ppp)
        """
        all_indicators = {**cls.WDI_INDICATORS, **cls.WDI_SERVICE_INDICATORS}
        return all_indicators.get(indicator_code)
