"""Utility functions."""

import io
from pathlib import Path
from zipfile import ZipFile

import httpx
import pandas as pd
import rich.progress
import us
from rich import print

from .constants import ACS_PUMS_URL, data_dir


def _clean_year(year: int) -> int:
    if 0 <= year <= 24:
        year += 2000

    if not (2000 <= year <= 2024):
        raise ValueError("Year must be between 2000 and 2024.")
    return year


def _clean_survey(_survey: str, _year: int) -> str:
    """
    Fixes URL part for survey. Some years don't have 3-Year surveys.
    If year <= 2006, _survey == ''.
    From 2007-2008, _survey can be either 1 or 3 years.
    From 2009-2013, _survey can be either 1, 3, or 5 years.
    From 2014 onward, only 1 or 5 years (except 2020: 5-Year only).
    """
    _survey = _survey.title()
    if _year <= 2006:
        if _survey != "1-Year":
            print("Prior to 2007, only 1-Year ACS are available, defaulting to 1-Year")
        return ""
    elif 2007 <= _year <= 2008:
        if _survey == "5-Year":
            print(f"There is no 5-Year ACS for {_year}, defaulting to 3-Year")
            return "3-Year/"
    elif _year == 2020:
        if _survey != "5-Year":
            print(
                f"The standard 1-Year ACS PUMS for {_year} was not released"
                " due to COVID-19, defaulting to 5-Year"
            )
        return "5-Year/"
    elif _year >= 2014:
        if _survey == "3-Year":
            print(f"There is no 3-Year ACS for {_year}, defaulting to 5-Year")
            return "5-Year/"
    return f"{_survey}/"


def _check_data_dirs(data_directory: data_dir) -> Path:
    """Validate data directory exists, creating 'raw/' and 'interim/' if needed."""
    # set directory's values
    _raw_data_directory = data_directory.joinpath("raw/")
    _interim_data_directory = data_directory.joinpath("interim/")

    # make sure they exists
    if not data_directory.exists():
        data_directory.mkdir()
    if not _raw_data_directory.exists():
        _raw_data_directory.mkdir()
    if not _interim_data_directory.exists():
        _interim_data_directory.mkdir()

    return data_directory


def _extract_data(downloaded_file: Path, extract_dir: Path):
    """Extract survey data downloaded from Census server."""
    *_, survey_year, _ = downloaded_file.parts
    state = downloaded_file.stem[-2:]

    full_extract_dir_path = extract_dir.joinpath(survey_year)
    if not full_extract_dir_path.exists():
        full_extract_dir_path.mkdir()
    full_extract_path = full_extract_dir_path.joinpath(state)
    if not full_extract_path.exists():
        full_extract_path.mkdir()
    content_file = ZipFile(downloaded_file)

    for file in rich.progress.track(
        content_file.filelist,
        description="Extracting...",
    ):
        content_file.extract(file, str(full_extract_path))

    print(f"Files extracted successfully at [magenta]{full_extract_path}[/magenta]")


def _download_data(
    url: str,
    name: str,
    data_directory: data_dir,
    extract: bool = True,
) -> None:
    """
    Downloads a file from Census FTP server.
    """

    data_directory = _check_data_dirs(data_directory=data_directory)
    _download_path = data_directory.joinpath("raw/")
    _extract_path = data_directory.joinpath("interim/")

    *_, year, _, _filename = url.split("/")
    year = year[-2:]
    _survey_dir = _download_path.joinpath(f"{name}_{year}/")

    if not _survey_dir.exists():
        _survey_dir.mkdir()

    _full_download_path = _survey_dir.joinpath(_filename)

    with (
        open(_full_download_path, "wb") as file,
        httpx.stream("GET", url, timeout=None) as response,
    ):
        total = response.headers.get("Content-Length")

        with rich.progress.Progress(
            "[progress.percentage]{task.percentage:>3.0f}%",
            rich.progress.BarColumn(bar_width=None),
            rich.progress.DownloadColumn(),
            rich.progress.TransferSpeedColumn(),
        ) as progress:
            if total:
                download_task = progress.add_task("Download", total=int(total))
            else:
                download_task = progress.add_task("Download")
            for chunk in response.iter_bytes():
                file.write(chunk)
                progress.update(download_task, completed=response.num_bytes_downloaded)
        print(
            f"File downloaded successfully at [magenta]{_full_download_path}[/magenta]"
        )

    if extract:
        _extract_data(_full_download_path, _extract_path)


def _download_as_dataframe(URL: str) -> pd.DataFrame:
    """Download a zip file from URL and return the CSV inside as a DataFrame.

    Parameters
    ----------
    URL : str
        URL to retrieve zip file from

    Returns
    -------
    pd.DataFrame
        DataFrame containing data from CSV
    """
    response = httpx.get(URL, timeout=None)

    with ZipFile(io.BytesIO(response.content)) as thezip:
        csv_files = [
            file for file in thezip.infolist() if file.filename.endswith(".csv")
        ]
        if len(csv_files) != 1:
            raise ValueError(
                f"Expected exactly 1 CSV file in zip archive, found {len(csv_files)}"
            )
        with thezip.open(csv_files[0]) as thefile:
            data = pd.read_csv(thefile)
    return data


def build_acs_url(
    year: int = 2023,
    survey: str = "1-year",
    sample_unit: str = "person",
    state: str = "California",
) -> str:
    """Build Census FTP server URL for ACS 1-, 3-, or 5-year estimates.

    Parameters
    ----------
    year : int, optional
        Year of survey, by default 2023
    survey : str, optional
        One of '1-year', '3-year', or '5-year', by default '1-year'
    sample_unit : str, optional
        Person-level or Household-level estimates, by default 'person'
    state : str, optional
        Name or abbreviation of state, by default 'California'

    Returns
    -------
    str
        URL to retrieve data from.
    """

    unit = sample_unit[0].lower()
    state_abbr = us.states.lookup(state).abbr.lower()

    if "5" in survey:
        survey = "5-Year"
    elif "3" in survey:
        survey = "3-Year"
    else:
        survey = "1-Year"

    year = _clean_year(year)

    survey = f"{_clean_survey(survey, year)}"

    survey_url = f"{ACS_PUMS_URL}{year}/{survey}csv_{unit}{state_abbr}.zip"
    return survey_url
