"""
seafog.gfs provides methods to download GFS forecast results from `NOAA <https://nomads.ncep.noaa.gov/cgi-bin/filter_gfs_0p25.pl>`_.
"""

from copy import copy
from datetime import datetime, timedelta
from os import makedirs
from os.path import exists
from time import time
from typing import Dict, Union, Callable, Tuple
from urllib.parse import urlencode

from rich.progress import Progress

from .utils import logger, download_url

# The root URL to download GFS data
ROOT_URL = "https://nomads.ncep.noaa.gov/cgi-bin/filter_gfs_0p25.pl"

# GFS only keeps ten days of data, so we need to set an expiry time and check it
# Units: day
EXPIRY_TIME = 10

# Here is all the params which could be used to download GFS forecast data, which is parsed from the url generated by the website below.
# For more details, see this website: https://nomads.ncep.noaa.gov/gribfilter.php?ds=gfs_0p25
# Select params and click "Show Url", you will see the URL generated by the website.
# Default template dict is suitable to download GFS data for WRF running.
# There are four kinds of key:
# "dir": The store path of forecast data. {DATE} means the date of the forecast, {FORECAST_INDEX} is one of: 00, 06, 12, 18
# "file": From which file you want to collect data.
# "lev_*": Level settings.
# "var_*": Variables settings.
GFS_PARAMS_DICT_TEMPLATE = {
    "dir": "/gfs.{DATE}/{FORECAST_INDEX}/atmos",
    "file": "gfs.t{FORECAST_INDEX}z.pgrb2.0p25.f000",
    "lev_0-0.1_m_below_ground": "on",
    "lev_0.01_mb": "off",
    "lev_0.02_mb": "off",
    "lev_0.04_mb": "off",
    "lev_0.07_mb": "off",
    "lev_0.1-0.4_m_below_ground": "on",
    "lev_0.1_mb": "off",
    "lev_0.2_mb": "off",
    "lev_0.33-1_sigma_layer": "off",
    "lev_0.4-1_m_below_ground": "on",
    "lev_0.44-0.72_sigma_layer": "off",
    "lev_0.44-1_sigma_layer": "off",
    "lev_0.4_mb": "off",
    "lev_0.72-0.94_sigma_layer": "off",
    "lev_0.7_mb": "off",
    "lev_0.995_sigma_level": "off",
    "lev_0C_isotherm": "off",
    "lev_1-2_m_below_ground": "on",
    "lev_1000_m_above_ground": "off",
    "lev_1000_mb": "on",
    "lev_100_m_above_ground": "off",
    "lev_100_mb": "on",
    "lev_10_m_above_ground": "on",
    "lev_10_m_above_mean_sea_level": "off",
    "lev_10_mb": "off",
    "lev_150_mb": "on",
    "lev_15_mb": "off",
    "lev_180-0_mb_above_ground": "off",
    "lev_1829_m_above_mean_sea_level": "off",
    "lev_1_hybrid_level": "off",
    "lev_1_mb": "off",
    "lev_200_mb": "on",
    "lev_20_m_above_ground": "off",
    "lev_20_mb": "off",
    "lev_250_mb": "on",
    "lev_255-0_mb_above_ground": "off",
    "lev_2743_m_above_mean_sea_level": "off",
    "lev_2_hybrid_level": "off",
    "lev_2_m_above_ground": "on",
    "lev_2_mb": "off",
    "lev_30-0_mb_above_ground": "off",
    "lev_3000-0_m_above_ground": "off",
    "lev_300_mb": "on",
    "lev_30_m_above_ground": "off",
    "lev_30_mb": "off",
    "lev_350_mb": "on",
    "lev_3658_m_above_mean_sea_level": "off",
    "lev_3_mb": "off",
    "lev_4000_m_above_ground": "off",
    "lev_400_mb": "on",
    "lev_40_m_above_ground": "off",
    "lev_40_mb": "off",
    "lev_450_mb": "on",
    "lev_500_mb": "on",
    "lev_50_m_above_ground": "off",
    "lev_50_mb": "on",
    "lev_550_mb": "on",
    "lev_5_mb": "off",
    "lev_6000-0_m_above_ground": "off",
    "lev_600_mb": "on",
    "lev_650_mb": "on",
    "lev_700_mb": "on",
    "lev_70_mb": "on",
    "lev_750_mb": "on",
    "lev_7_mb": "off",
    "lev_800_mb": "on",
    "lev_80_m_above_ground": "off",
    "lev_850_mb": "on",
    "lev_90-0_mb_above_ground": "off",
    "lev_900_mb": "on",
    "lev_925_mb": "on",
    "lev_950_mb": "on",
    "lev_975_mb": "on",
    "lev_PV=-2e-06_(Km^2/kg/s)_surface": "off",
    "lev_PV=2e-06_(Km^2/kg/s)_surface": "off",
    "lev_boundary_layer_cloud_layer": "off",
    "lev_cloud_ceiling": "off",
    "lev_convective_cloud_bottom_level": "off",
    "lev_convective_cloud_layer": "off",
    "lev_convective_cloud_top_level": "off",
    "lev_entire_atmosphere": "off",
    "lev_entire_atmosphere_(considered_as_a_single_layer)": "off",
    "lev_high_cloud_bottom_level": "off",
    "lev_high_cloud_layer": "off",
    "lev_high_cloud_top_level": "off",
    "lev_highest_tropospheric_freezing_level": "off",
    "lev_low_cloud_bottom_level": "off",
    "lev_low_cloud_layer": "off",
    "lev_low_cloud_top_level": "off",
    "lev_max_wind": "off",
    "lev_mean_sea_level": "on",
    "lev_middle_cloud_bottom_level": "off",
    "lev_middle_cloud_layer": "off",
    "lev_middle_cloud_top_level": "off",
    "lev_planetary_boundary_layer": "off",
    "lev_surface": "on",
    "lev_top_of_atmosphere": "off",
    "lev_tropopause": "off",
    "var_4LFTX": "off",
    "var_ABSV": "off",
    "var_ACPCP": "off",
    "var_ALBDO": "off",
    "var_APCP": "off",
    "var_CAPE": "off",
    "var_CFRZR": "off",
    "var_CICEP": "off",
    "var_CIN": "off",
    "var_CLWMR": "off",
    "var_CNWAT": "off",
    "var_CPOFP": "off",
    "var_CPRAT": "off",
    "var_CRAIN": "off",
    "var_CSNOW": "off",
    "var_CWAT": "off",
    "var_CWORK": "off",
    "var_DLWRF": "off",
    "var_DPT": "off",
    "var_DSWRF": "off",
    "var_DZDT": "off",
    "var_FLDCP": "off",
    "var_FRICV": "off",
    "var_GFLUX": "off",
    "var_GRLE": "off",
    "var_GUST": "off",
    "var_HCDC": "off",
    "var_HGT": "on",
    "var_HINDEX": "off",
    "var_HLCY": "off",
    "var_HPBL": "off",
    "var_ICAHT": "off",
    "var_ICEC": "on",
    "var_ICEG": "off",
    "var_ICETK": "off",
    "var_ICETMP": "off",
    "var_ICMR": "off",
    "var_LAND": "on",
    "var_LCDC": "off",
    "var_LFTX": "off",
    "var_LHTFL": "off",
    "var_MCDC": "off",
    "var_MSLET": "off",
    "var_O3MR": "off",
    "var_PEVPR": "off",
    "var_PLPL": "off",
    "var_POT": "off",
    "var_PRATE": "off",
    "var_PRES": "on",
    "var_PRMSL": "on",
    "var_PWAT": "off",
    "var_REFC": "off",
    "var_REFD": "off",
    "var_RH": "on",
    "var_RWMR": "off",
    "var_SFCR": "off",
    "var_SHTFL": "off",
    "var_SNMR": "off",
    "var_SNOD": "off",
    "var_SOILL": "off",
    "var_SOILW": "on",
    "var_SOTYP": "off",
    "var_SPFH": "off",
    "var_SUNSD": "off",
    "var_TCDC": "off",
    "var_TMAX": "off",
    "var_TMIN": "off",
    "var_TMP": "on",
    "var_TOZNE": "off",
    "var_TSOIL": "on",
    "var_U-GWD": "off",
    "var_UFLX": "off",
    "var_UGRD": "on",
    "var_ULWRF": "off",
    "var_USTM": "off",
    "var_USWRF": "off",
    "var_V-GWD": "off",
    "var_VEG": "off",
    "var_VFLX": "off",
    "var_VGRD": "on",
    "var_VIS": "off",
    "var_VRATE": "off",
    "var_VSTM": "off",
    "var_VVEL": "off",
    "var_VWSH": "off",
    "var_WATR": "off",
    "var_WEASD": "on",
    "var_WILT": "off",
    "leftlon": "70",
    "rightlon": "165",
    "toplat": "65",
    "bottomlat": "8"
}


def get_gfs_param_template(date: Union[str, datetime], forecast_index: int) -> Dict[str, str]:
    """
    Get param template dict to download GFS data.

    :param date: Date of the GFS data you download, either a string or datetime object.
                 For example, ``"2024-04-02"``.
    :param forecast_index: GFS runs 4 times a day, so there will be four indexes, [0, 6, 12, 18].
    :return: Parameters dict to download GFS data.
    """
    global GFS_PARAMS_DICT_TEMPLATE

    # check the date form
    if isinstance(date, str):
        date = datetime.strptime(date, "%Y-%m-%d")

    # get today's date
    today_date = datetime.fromtimestamp(time())

    # check if `date` is ten days before today
    if today_date - date > timedelta(days=10):
        logger.error(f"You can't download GFS data more than 10 days before, because GFS only keeps 10 days results")
        raise ValueError

    # check forecast index
    if forecast_index not in [0, 6, 12, 18]:
        logger.error(f"`forecast_index` should be one of [0, 6, 12, 18]")
        raise ValueError

    # copy template
    new_template = copy(GFS_PARAMS_DICT_TEMPLATE)

    # convert forecast_index and date to string
    forecast_index = str(forecast_index).rjust(2, '0')
    date = date.strftime("%Y%m%d")

    # fill value
    new_template["dir"] = new_template["dir"].format(DATE=date, FORECAST_INDEX=forecast_index)
    new_template["file"] = new_template["file"].format(FORECAST_INDEX=forecast_index)

    return new_template


def generate_gfs_download_url(param_dict: Dict[str, str]) -> str:
    """
    Generate GFS download URL based on the parameters' dict.

    :param param_dict: Parameters dict.
    :return: URL to download GFS data.

    Generate default url from template:

    >>> data_date = datetime.fromtimestamp(time()).strftime("%Y-%m-%d")
    >>> default_params = get_gfs_param_template(data_date, forecast_index=0)
    >>> url = generate_gfs_download_url(default_params)

    """
    # we need to remove the lev_* and var_* key with value "off"
    new_param_dict = {}
    for key in param_dict:

        if key.startswith("lev_") or key.startswith("var_"):
            if param_dict[key] != "off":
                new_param_dict[key] = param_dict[key]
        else:
            new_param_dict[key] = param_dict[key]

    # generate URL
    param_string = urlencode(new_param_dict)

    return f"{ROOT_URL}?{param_string}"


def gfs_find_data(date: str, save_path: str, area: Tuple[int, int, int, int], forecast_index: int, param_dict: Dict[str, str] = None,
                  proxy_host: str = None, proxy_port: int = None, progress: Progress = None, show_progress=True, callback: Callable = None) -> str:
    """
    Download GFS data in the last 10 days from `NOAA <https://nomads.ncep.noaa.gov/cgi-bin/filter_gfs_0p25.pl>`_.
    NOAA only keeps the last 10-day results, any date earlier than 10 days will raise an ValueError.

    :param date: Date of GFS data you want to download.
                 For example, ``"2024-04-02"``.
    :param save_path: Save path of GFS data.
    :param area: Longitude and latitude range, for example, (lon1, lon2, lat1, lat2).
    :param forecast_index: GFS runs 4 times a day, so there will be four indexes, [0, 6, 12, 18].
    :param param_dict: You can provide custom parameter dict to download GFS data; otherwise, we will download with default selection of variables.
    :param proxy_host: Http and https proxy server address.
    :param proxy_port: Http and https proxy server port.
    :param progress:  Living progress object.
    :param show_progress: If True, display download progress in console.
    :param callback: Callable object that will be called every iteration during data download.
                     ``callback`` should accept two params, `total_size` and `step_size`
    :return: path of downloaded data.

    Download today's data in area (116E-128E, 30N-40N) with ``forecast_index=0`` and default params:

    >>> today_date = datetime.fromtimestamp(time()).strftime("%Y-%m-%d")
    >>> gfs_find_data(today_date, "data", area=(116, 128, 30, 40), forecast_index=0, show_progress=False)
    'data/gfs.t00z.pgrb2.0p25.f000'
    
    Download data with custom params:

    >>> today_date = datetime.fromtimestamp(time()).strftime("%Y-%m-%d")
    >>> params = get_gfs_param_template(today_date, forecast_index=0)
    >>> # change any params
    >>> data_path = gfs_find_data(today_date, "data", area=(116, 128, 30, 40), forecast_index=0, show_progress=False)

    """
    # check the save path
    if not exists(save_path):
        makedirs(save_path)

    # get template parameters dict
    if param_dict is None:
        param_dict = get_gfs_param_template(date, forecast_index)

    # update longitude and latitude range
    lon1, lon2, lat1, lat2 = area
    param_dict["leftlon"] = str(lon1)
    param_dict["rightlon"] = str(lon2)
    param_dict["bottomlat"] = str(lat1)
    param_dict["toplat"] = str(lat2)

    # generate save name.
    # data of the same date may have different variables and in different areas, so we just use value `param_dict['file']` as the save name.
    save_name = param_dict["file"]

    # generate URL
    url = generate_gfs_download_url(param_dict)

    # download
    return_code = download_url(url, save_path, save_name, proxy_host=proxy_host, proxy_port=proxy_port, show_progress=show_progress, progress=progress, callback=callback)

    # check
    if return_code != 200:
        logger.error(f"Fail to download file {save_name}, status code is {return_code}")
        raise ConnectionError

    else:
        return f"{save_path}/{save_name}"


__all__ = ["get_gfs_param_template", "generate_gfs_download_url", "gfs_find_data"]
