from datetime import datetime, timezone
from dateutil.relativedelta import relativedelta
from typing import Union
import calendar


def add_months_safe_utc(date: datetime, months: int) -> datetime:
    """
    Añade meses a una fecha de forma segura, manejando casos de desbordamiento de días.
    Por ejemplo: 31 de enero + 1 mes = 28 de febrero (último día del mes).
    """
    year = date.year
    month = date.month
    day = date.day

    total_months = month + months - 1  # -1 porque los meses son 1-indexed
    target_year = year + (total_months // 12)
    target_month = (total_months % 12) + 1

    # Obtener el último día del mes objetivo
    last_day = calendar.monthrange(target_year, target_month)[1]
    new_day = min(day, last_day)

    return datetime(
        target_year, target_month, new_day,
        date.hour, date.minute, date.second, date.microsecond,
        tzinfo=timezone.utc
    )


def count_monthly_occurrences(
        start: Union[str, datetime],
        end: Union[str, datetime],
        include_start: bool = False
) -> int:
    """
    Devuelve el número de ocurrencias mensuales entre start y end.

    Args:
        start: Fecha de inicio (string ISO o datetime)
        end: Fecha de fin (string ISO o datetime)
        include_start: Si False (default), empieza a contar desde start + 1 mes.
                      Si True, incluye start como primer mes.

    Returns:
        Número de ocurrencias mensuales
    """
    # Convertir a datetime si son strings
    if isinstance(start, str):
        try:
            start_dt = datetime.fromisoformat(start.replace('Z', '+00:00'))
        except ValueError:
            return 0
    else:
        start_dt = start

    if isinstance(end, str):
        try:
            end_dt = datetime.fromisoformat(end.replace('Z', '+00:00'))
        except ValueError:
            return 0
    else:
        end_dt = end

    # Asegurar que las fechas tengan timezone UTC
    if start_dt.tzinfo is None:
        start_dt = start_dt.replace(tzinfo=timezone.utc)
    if end_dt.tzinfo is None:
        end_dt = end_dt.replace(tzinfo=timezone.utc)

    # Validaciones
    if end_dt < start_dt:
        return 0

    count = 0
    cursor = datetime(
        start_dt.year, start_dt.month, start_dt.day,
        start_dt.hour, start_dt.minute, start_dt.second, start_dt.microsecond,
        tzinfo=timezone.utc
    )

    if not include_start:
        cursor = add_months_safe_utc(cursor, 1)

    while cursor <= end_dt:
        count += 1
        next_cursor = add_months_safe_utc(cursor, 1)

        # Seguridad: evitar loop infinito
        if next_cursor == cursor:
            break

        cursor = next_cursor

        # Límite de seguridad
        if count > 1000:
            break

    return count