"""Accessing watershed boundary-level data through web services."""
from __future__ import annotations

import importlib.util
import io
from pathlib import Path
from typing import TYPE_CHECKING, Union

import async_retriever as ar
import geopandas as gpd
import pandas as pd
import pygeoogc as ogc
import pyproj
import xarray as xr
from pygeoogc import ServiceURL
from pynhd import AGRBase
from pynhd.core import ScienceBase

from pygeohydro.exceptions import InputValueError

if TYPE_CHECKING:
    CRSTYPE = Union[int, str, pyproj.CRS]

__all__ = [
    "WBD",
    "huc_wb_full",
    "irrigation_withdrawals",
]


class WBD(AGRBase):
    """Access Watershed Boundary Dataset (WBD).

    Notes
    -----
    This web service offers Hydrologic Unit (HU) polygon boundaries for
    the United States, Puerto Rico, and the U.S. Virgin Islands.
    For more info visit: https://hydro.nationalmap.gov/arcgis/rest/services/wbd/MapServer

    Parameters
    ----------
    layer : str, optional
        A valid service layer. Valid layers are:

        - ``wbdline``
        - ``huc2``
        - ``huc4``
        - ``huc6``
        - ``huc8``
        - ``huc10``
        - ``huc12``
        - ``huc14``
        - ``huc16``

    outfields : str or list, optional
        Target field name(s), default to "*" i.e., all the fields.
    crs : str, int, or pyproj.CRS, optional
        Target spatial reference, default to ``EPSG:4326``.
    """

    def __init__(self, layer: str, outfields: str | list[str] = "*", crs: CRSTYPE = 4326):
        self.valid_layers = {
            "wbdline": "wbdline",
            "huc2": "2-digit hu (region)",
            "huc4": "4-digit hu (subregion)",
            "huc6": "6-digit hu (basin)",
            "huc8": "8-digit hu  (subbasin)",
            "huc10": "10-digit hu (watershed)",
            "huc12": "12-digit hu (subwatershed)",
            "huc14": "14-digit hu",
            "huc16": "16-digit hu",
        }
        _layer = self.valid_layers.get(layer)
        if _layer is None:
            raise InputValueError("layer", list(self.valid_layers))
        super().__init__(ServiceURL().restful.wbd, _layer, outfields, crs)


def huc_wb_full(huc_lvl: int) -> gpd.GeoDataFrame:
    """Get the full watershed boundary for a given HUC level.

    Notes
    -----
    This function is designed for cases where the full watershed boundary is needed
    for a given HUC level. If only a subset of the HUCs is needed, then use
    the ``pygeohydro.WBD`` class. The full dataset is downloaded from the National Maps'
    `WBD staged products <https://prd-tnm.s3.amazonaws.com/index.html?prefix=StagedProducts/Hydrography/WBD/HU2/Shape/>`__.

    Parameters
    ----------
    huc_lvl : int
        HUC level, must be even numbers between 2 and 16.

    Returns
    -------
    geopandas.GeoDataFrame
        The full watershed boundary for the given HUC level.
    """
    valid_hucs = [2, 4, 6, 8, 10, 12, 14, 16]
    if huc_lvl not in valid_hucs:
        raise InputValueError("huc_lvl", list(map(str, valid_hucs)))

    base_url = "https://prd-tnm.s3.amazonaws.com/StagedProducts/Hydrography/WBD/HU2/Shape"

    urls = [f"{base_url}/WBD_{h2:02}_HU2_Shape.zip" for h2 in range(1, 23)]
    fnames = [Path("cache", Path(url).name) for url in urls]
    _ = ogc.streaming_download(urls, fnames=fnames)  # type: ignore
    keys = (p.stem.split("_")[1] for p in fnames)
    engine = "pyogrio" if importlib.util.find_spec("pyogrios") else "fiona"
    huc = gpd.GeoDataFrame(
        pd.concat(
            (gpd.read_file(f"{p}!Shape/WBDHU{huc_lvl}.shp", engine=engine) for p in fnames),
            keys=keys,
        )
    )
    huc = huc.reset_index().rename(columns={"level_0": "huc2"}).drop(columns="level_1")
    return huc


def irrigation_withdrawals() -> xr.Dataset:
    """Get monthly water use for irrigation at HUC12-level for CONUS.

    Notes
    -----
    Dataset is retrieved from https://doi.org/10.5066/P9FDLY8P.
    """
    sb = ScienceBase()
    item = sb.get_file_urls("5ff7acf4d34ea5387df03d73")
    urls = item.loc[item.index.str.contains(".csv"), "url"]
    resp = ar.retrieve_text(urls.tolist())
    irr = {}
    for name, r in zip(urls.index, resp):
        df = pd.read_csv(io.StringIO(r), usecols=lambda x: "m3" in x or "huc12t" in x)
        df["huc12t"] = df["huc12t"].str.strip("'")
        df = df.rename(columns={"huc12t": "huc12"}).set_index("huc12")
        df = df.rename(columns={c: c[:3].capitalize() for c in df})
        irr[name[-6:-4]] = df.copy()
    ds = xr.Dataset(irr)
    ds = ds.rename({"dim_1": "month"})
    long_names = {
        "GW": "groundwater_withdrawal",
        "SW": "surface_water_withdrawal",
        "TW": "total_withdrawal",
        "CU": "consumptive_use",
    }
    for v, n in long_names.items():
        ds[v].attrs["long_name"] = n
        ds[v].attrs["units"] = "m3"
    ds.attrs["description"] = " ".join(
        (
            "Estimated Monthly Water Use for Irrigation by",
            "12-Digit Hydrologic Unit in the Conterminous United States for 2015",
        )
    )
    ds.attrs["source"] = "https://doi.org/10.5066/P9FDLY8P"
    return ds
