"""Ordinance jurisdiction info"""

import logging
from warnings import warn
from pathlib import Path

import numpy as np
import pandas as pd

from compass.exceptions import COMPASSValueError
from compass.warn import COMPASSWarning


logger = logging.getLogger(__name__)
_COUNTY_DATA_FP = (
    Path(__file__).parent.parent / "data" / "conus_jurisdictions.csv"
)


def load_all_jurisdiction_info():
    """Load canonical jurisdiction metadata for the continental US

    Returns
    -------
    pandas.DataFrame
        Table containing jurisdiction names, FIPS codes, official
        websites, and related attributes.

    Notes
    -----
    Missing values are normalized to ``None`` to simplify downstream
    serialization.
    """
    return pd.read_csv(_COUNTY_DATA_FP).replace({np.nan: None})


def jurisdiction_websites(jurisdiction_info=None):
    """Build a mapping of jurisdiction identifiers to website URLs

    Parameters
    ----------
    jurisdiction_info : pandas.DataFrame, optional
        DataFrame containing jurisdiction names and websites. If
        ``None``, this info is loaded using
        :func:`load_all_jurisdiction_info`.
        By default, ``None``.

    Returns
    -------
    dict
        Mapping from jurisdiction FIPS codes to their primary website
        URLs.

    Notes
    -----
    The helper uses FIPS codes rather than string names to avoid
    collisions between same-named jurisdictions in different states.
    """
    if jurisdiction_info is None:
        jurisdiction_info = load_all_jurisdiction_info()

    return {
        row["FIPS"]: row["Website"] for __, row in jurisdiction_info.iterrows()
    }


def load_jurisdictions_from_fp(jurisdiction_fp):
    """Load jurisdiction metadata for entries listed in a CSV file

    This loader trims whitespace, deduplicates request rows, and filters
    out jurisdictions not present in the canonical data set.

    Parameters
    ----------
    jurisdiction_fp : path-like
        Path to csv file containing "County" and "State" columns that
        define the jurisdictions for which info should be loaded.

    Returns
    -------
    pandas.DataFrame
        Jurisdiction information, including FIPS codes and websites,
        for every matching entry in the lookup table.

    Raises
    ------
    COMPASSValueError
        If the input file is missing required columns (``State`` or
        ``Jurisdiction Type`` when subdivisions are provided).

    Notes
    -----
    Missing jurisdictions trigger warnings with a tabular summary.
    """
    jurisdictions = pd.read_csv(jurisdiction_fp).replace({np.nan: None})
    jurisdictions = _validate_jurisdiction_input(jurisdictions)

    all_jurisdiction_info = load_all_jurisdiction_info()
    merge_cols = ["County", "State"]
    if "Subdivision" in jurisdictions:
        merge_cols += ["Subdivision", "Jurisdiction Type"]
    else:
        all_jurisdiction_info = all_jurisdiction_info[
            all_jurisdiction_info["Subdivision"].isna()
        ].reset_index(drop=True)

    jurisdictions = (  # remove dupes
        jurisdictions.groupby(merge_cols, dropna=False)
        .first()
        .reset_index()
        .drop(columns="Unnamed: 0", errors="ignore")
        .replace({np.nan: None})
    )
    jurisdictions["jur_merge"] = jurisdictions.apply(
        _build_merge_col, axis=1, merge_cols=merge_cols
    )
    all_jurisdiction_info["jur_merge"] = all_jurisdiction_info.apply(
        _build_merge_col, axis=1, merge_cols=merge_cols
    )
    jurisdictions = jurisdictions.merge(
        all_jurisdiction_info,
        on="jur_merge",
        how="left",
        suffixes=["_user", ""],
    )

    jurisdictions = _filter_not_found_jurisdictions(jurisdictions, merge_cols)
    return _format_jurisdiction_df_for_output(jurisdictions)


def _validate_jurisdiction_input(jurisdictions):
    """Throw error if user is missing required columns"""
    if "State" not in jurisdictions:
        msg = "The jurisdiction input must have at least a 'State' column!"
        raise COMPASSValueError(msg)

    jurisdictions["State"] = jurisdictions["State"].str.strip()
    if "County" not in jurisdictions:
        jurisdictions["County"] = None
    else:
        jurisdictions["County"] = jurisdictions["County"].str.strip()

    if "Subdivision" in jurisdictions:
        if "Jurisdiction Type" not in jurisdictions:
            msg = (
                "The jurisdiction input must have a 'Jurisdiction Type' "
                "column if a 'Subdivision' column is provided (this helps "
                "avoid name clashes for certain subdivisions)!"
            )
            raise COMPASSValueError(msg)

        jurisdictions["Subdivision"] = jurisdictions["Subdivision"].str.strip()
        jurisdictions["Jurisdiction Type"] = (
            jurisdictions["Jurisdiction Type"].str.casefold().str.strip()
        )

    return jurisdictions


def _filter_not_found_jurisdictions(df, merge_cols):
    """Filter out jurisdictions with null FIPS codes"""
    _warn_about_missing_jurisdictions(df, merge_cols)
    return df[~df["FIPS"].isna()].copy()


def _warn_about_missing_jurisdictions(df, merge_cols):
    """Throw warning about jurisdictions that were not in the list"""
    not_found_jurisdictions = df[df["FIPS"].isna()]
    if len(not_found_jurisdictions):
        out_cols = {f"{col}_user": col for col in merge_cols}
        not_found_jurisdictions = not_found_jurisdictions[
            list(out_cols)
        ].rename(columns=out_cols)
        not_found_jurisdictions_str = not_found_jurisdictions[
            merge_cols
            # cspell: disable-next-line
        ].to_markdown(index=False, tablefmt="psql")
        msg = (
            "The following jurisdictions were not found! Please make sure to "
            "use proper spelling and capitalization.\n"
            f"{not_found_jurisdictions_str}"
        )
        warn(msg, COMPASSWarning)


def _format_jurisdiction_df_for_output(df):
    """Format jurisdiction DataFrame for output"""
    out_cols = [
        "County",
        "State",
        "Subdivision",
        "Jurisdiction Type",
        "FIPS",
        "Website",
    ]
    df["FIPS"] = df["FIPS"].astype(int)
    return df[out_cols].replace({np.nan: None}).reset_index(drop=True)


def _build_merge_col(row, merge_cols):
    """Build column to merge jurisdiction DataFrames on"""
    return " ".join(str(row[c]).casefold() for c in merge_cols)
