"""UN World Population Prospects data fetching.

Provides the UNPopulationData class for fetching population forecast data
from the UN Population Division API.

Requires the 'all' extra or just requests: pip install atlas-common-data[all]
"""

import io
from pathlib import Path

import pandas as pd
import requests


class UNPopulationData:
    """Fetch and process UN World Population Prospects data.

    Downloads population forecasts by variant (Medium, High, Low, etc.)
    from the UN Population Division.

    Example:
        un = UNPopulationData()
        df = un.get_population_forecast()
    """

    COLUMNS = {
        "ISO3 Alpha-code": "iso3_code",
        "Year": "year",
        "Total Population, as of 1 January (thousands)": "population",
    }

    MULTIPLIERS = {
        "population": 1_000,
    }

    def __init__(self, latest_year):
        """Initialize UNPopulationData.

        Args:
            latest_year: Default end year for queries. Methods will use
                this when end_year is not explicitly passed.
            variants: Default list of variant names to fetch
                (e.g. ['Medium', 'High']). Methods will use this when
                variants is not explicitly passed. Defaults to ['Medium'].
        """
        self._country_classification = None
        self.latest_year = latest_year
        

    def get_population_forecast(self, start_year=None, end_year=None):
        """Fetch UN population forecast data.

        Returns:
            DataFrame with columns: iso3_code, year, population
        """
        df_forecast = self._fetch_population_data("Medium variant")
        df_pop = self._fetch_population_data("Estimates")
        df = pd.concat([df_forecast, df_pop])

        if start_year is not None:
            df = df[df["year"] >= start_year]

        if end_year is not None:
            df = df[df["year"] <= end_year]

        df = df[df["iso3_code"].notna()]
        df["population"] = (df["population"] * self.MULTIPLIERS["population"]).astype(int)
        df = df.sort_values(by=['iso3_code', 'year'])
        return df

    @classmethod
    def list_available_columns(cls):
        """List output columns returned by get_population_forecast.

        Returns:
            dict mapping output column names to descriptions.
        """
        return {
            "iso3_code": "ISO 3166-1 alpha-3 country code",
            "year": "Year of observation",
            "population": "Total population (individuals)",
        }


    def _fetch_population_data(self, sheet_name):
        """Download population data from UN, trying multiple release years.

        Returns:
            Raw DataFrame with original UN column names.
        """
        import datetime

        current_year = datetime.datetime.now().year
        # range is to detect the latest file
        for year in range(current_year + 1, current_year - 3, -1):
            url = (
                f'https://population.un.org/wpp/assets/Excel%20Files/1_Indicator%20(Standard)/EXCEL_FILES/1_General/WPP{year}_GEN_F01_DEMOGRAPHIC_INDICATORS_COMPACT.xlsx'
            )
            try:
                r = requests.get(url, timeout=30)
                if r.status_code == 200:
                    df = pd.read_excel(url, sheet_name=sheet_name, skiprows=15, header=1)
                    df = df[df["Type"] == "Country/Area"]
                    return df[self.COLUMNS.keys()].rename(columns=self.COLUMNS)
            except Exception:
                continue

        raise ValueError(
            "Could not download UN Population Prospects data. "
            "The UN download URL may have changed."
        )



