from .dates_add_days import CivilDate, days_in_month, format_iso_date, parse_iso_date


def add_months(iso: str, months: int) -> str:
    """Add whole months, keeping the day of the month and clamping it to the
    last day of the target month when that month is shorter (31 Jan + 1 =
    28/29 Feb).

    The clamp is applied to the result only. Callers stepping through a
    schedule must add k months to the original anchor, not 1 month k times, or
    the 31st decays to the 28th after February.
    """
    if isinstance(months, bool) or not isinstance(months, int):
        raise TypeError("months must be an integer, received %r" % (months,))
    date = parse_iso_date(iso)
    # Count months from year 0 so the year and month fall out of one division.
    total = date.year * 12 + (date.month - 1) + months
    year = total // 12
    month = total - year * 12 + 1
    if year < 1 or year > 9999:
        raise ValueError(
            '"%s" plus %s months is outside the supported range 0001-01-01 to 9999-12-31' % (iso, months)
        )
    day = min(date.day, days_in_month(year, month))
    return format_iso_date(CivilDate(year=year, month=month, day=day))