"""Client for Probability of Default data."""

from datetime import date, datetime
from typing import Literal

import pandas as pd

from tellimer._models import Result
from tellimer._utils import parse_date
from tellimer.types import MakeRequestFunc

VALID_TENORS = {"1y", "3y", "5y"}
PD_COLUMNS = {
    "1y": "probability_of_default",
    "3y": "probability_of_default_3y",
    "5y": "probability_of_default_5y",
}


class ProbabilityDefaultClient:
    """
    A client for the Tellimer API's probability of default endpoints.
    """

    def __init__(self, make_request: MakeRequestFunc):
        """
        Initialize the client.

        Args:
            make_request: The function to use to make the request.
        """
        self._make_request = make_request

    def available_countries(self) -> list[str]:
        """
        List all available countries for probability of default data.

        Returns:
            A list of country codes.
        """
        url = "api/v1/probability-default/list-countries"
        response = self._make_request("get", url)
        return response.json()

    def get(
        self,
        countries: list[str] | str | None = None,
        include_contributions: bool = False,
        start_date: str | date | datetime | None = None,
        end_date: str | date | datetime | None = None,
        tenors: list[Literal["1y", "3y", "5y"]] | Literal["1y", "3y", "5y"] = "1y",
    ) -> Result:
        """
        Get probability of default data.

        Args:
            countries: Optional country or list of countries to get data for.
                      If not provided, fetches data for all available countries.
            include_contributions: Whether to include contribution breakdowns.
            start_date: Optional start date filter (ISO format).
            end_date: Optional end date filter (ISO format).
            tenors: Tenor horizon(s) to fetch. A single tenor string or a list
                    of tenors. Defaults to "1y". Valid values: "1y", "3y", "5y".

        Returns:
            A Result containing a pandas DataFrame with columns: date, country,
            probability_of_default, and optionally probability_of_default_3y /
            probability_of_default_5y depending on requested tenors.
        """
        if isinstance(countries, str):
            countries = [countries]

        if isinstance(tenors, str):
            tenors = [tenors]

        invalid = set(tenors) - VALID_TENORS
        if invalid:
            raise ValueError(
                f"Invalid tenor(s): {sorted(invalid)}. Valid tenors: {sorted(VALID_TENORS)}"
            )

        if countries:
            available = self.available_countries()
            for country in countries:
                if country not in available:
                    raise ValueError(
                        f"Country {country} is not available for probability of default data"
                    )

        url = "api/v1/probability-default"
        payload = {
            "countries": countries,
            "include_contributions": include_contributions,
            "start_date": parse_date(start_date) if start_date is not None else None,
            "end_date": parse_date(end_date) if end_date is not None else None,
            "tenors": tenors,
        }
        response = self._make_request("post", url, json=payload)
        data = response.json()
        df = pd.DataFrame(data)
        for tenor in tenors:
            col = PD_COLUMNS[tenor]
            if col in df.columns:
                df[col] = pd.to_numeric(df[col])
        # Reorder columns: date, country, PD columns (by tenor), then everything else
        pd_cols = [PD_COLUMNS[t] for t in tenors if PD_COLUMNS[t] in df.columns]
        leading = ["date", "country"] + pd_cols
        remaining = [c for c in df.columns if c not in leading]
        df = df[leading + remaining]
        return Result(metadata=None, data=df)
