from datetime import datetime, timedelta
from dateutil.relativedelta import relativedelta

def validate_partition(
    df,
    start_ds,
    end_ds=None,
    mode="once",
    ds_format="yyyymmdd",
    level="day",
    output="detail",
    col_validate = "ds"
):
    """
    Validate partition existence in a Spark DataFrame.

    Parameters:
    -----------
    df : pyspark.sql.DataFrame
        DataFrame to validate (must contain 'ds' or 'year', 'month', 'day' columns).
    start_ds : str
        Starting date in 'yyyymmdd' format.
    end_ds : str, optional
        Ending date in 'yyyymmdd' format (used for mode='range').
    mode : str
        'once'     → check one date only  
        'backward' → go backward until partition found  
        'range'    → check a date/month range  
        'empty'    → list partitions that are missing (❌ only)
    ds_format : str
        'yyyymmdd', 'yyyymm', or 'ymd'
    level : str
        'day', 'month', or 'month_end'
    output : str
        'detail', 'summary' (for "range" mode), or 'both' (for "range" mode)
    col_validate : str
        Name of column to validate
        
    Returns:
    --------
    str or dict or None
        - str if found (once/backward)
        - dict of date → True/False for range
        - None if not found
        
    Example:
    --------
    >>> validate_partition(df, start_ds="20250901", mode="once", ds_format="yyyymmdd", col_validate="event_date")
    >>> validate_partition(df, start_ds="20250901", mode="backward", ds_format="ymd", col_validate="ds")
    >>> validate_partition(df, "20250801", "20250805", mode="range", level="day", ds_format="yyyymmdd")
    >>> validate_partition(df, "20240101", "20240901", mode="range", level="month", ds_format="yyyymm")
    >>> validate_partition(df, "20240101", "20240901", mode="range", level="month_end", ds_format="ymd")

    """
    
    # --- parse start_date & end_date sesuai ds_format ---
    if ds_format == "yyyymmdd":
        start_date = datetime.strptime(start_ds, "%Y%m%d")
        end_date = datetime.strptime(end_ds, "%Y%m%d") if end_ds else start_date

    elif ds_format == "yyyymm":
        start_date = datetime.strptime(start_ds, "%Y%m")
        end_date = datetime.strptime(end_ds, "%Y%m") if end_ds else start_date

    elif ds_format == "ymd":
        start_date = datetime.strptime(start_ds, "%Y%m%d")
        end_date = datetime.strptime(end_ds, "%Y%m%d") if end_ds else start_date

    else:
        raise ValueError("ds_format must be 'yyyymmdd', 'yyyymm', or 'ymd'")

    results = {}
    current_date = start_date

    def check_partition(date_obj):
        y, m, d = date_obj.year, date_obj.month, date_obj.day
        key_ymd = date_obj.strftime("%Y%m%d")
        key_ym = date_obj.strftime("%Y%m")

        if ds_format == "yyyymmdd":
            count = df.filter(df[col_validate] == key_ymd).limit(1).count()
        elif ds_format == "yyyymm":
            count = df.filter(df[col_validate] == key_ym).limit(1).count()
        elif ds_format == "ymd":
            count = df.filter(
                (df["year"] == y) & (df["month"] == m) & (df["day"] == d)
            ).limit(1).count()
        else:
            raise ValueError("ds_format must be 'yyyymmdd', 'yyyymm', or 'ymd'")

        return count > 0

    if mode in ["once", "backward"]:
        while current_date >= datetime(2000, 1, 1):  # just a lower limit
            # tentukan key sesuai level
            if level == "month":
                key = current_date.strftime("%Y%m")
                date_check = current_date.replace(day=1)  # supaya konsisten cek partisi bulan
                status = check_partition(date_check)
            elif level == "month_end":
                last_day = current_date + relativedelta(day=31)
                key = last_day.strftime("%Y%m%d")
                status = check_partition(last_day)
            else:  # default: day
                key = current_date.strftime("%Y%m%d")
                status = check_partition(current_date)

            if status:
                print(f"Partition {key} ditemukan ✅")
                return key
            else:
                print(f"Partition {key} tidak ditemukan ❌")
                if mode == "once":
                    return None
                # backward step
                if level == "month":
                    current_date -= relativedelta(months=1)
                else:
                    current_date -= timedelta(days=1)

    elif mode == "range":
        while current_date <= end_date:
            if level == "day":
                key = current_date.strftime("%Y%m%d")
                status = check_partition(current_date)
                current_date += timedelta(days=1)

            elif level == "month":
                key = current_date.strftime("%Y%m")
                # Use 1st day of month to check month-level partition
                date_check = current_date.replace(day=1)
                status = check_partition(date_check)
                current_date += relativedelta(months=1)

            elif level == "month_end":
                last_day = current_date + relativedelta(day=31)
                key = last_day.strftime("%Y%m%d")
                status = check_partition(last_day)
                current_date += relativedelta(months=1)

            else:
                raise ValueError("level must be 'day', 'month', or 'month_end'")

            results[key] = status
            if output in ["detail", "both"]:
                print(f"{key}: {'✅' if status else '❌'}")

        if output in ["summary"]:
            total = len(results)
            available = sum(results.values())
            print(f"\nSummary: {available}/{total} partisi tersedia ✅")

        if output in ["both"]:
            total = len(results)
            available = sum(results.values())
            print(f"\nSummary: {available}/{total} partisi tersedia ✅")

    elif mode == "empty":
        missing = []
        while current_date <= end_date:
            if level == "day":
                key = current_date.strftime("%Y%m%d")
                status = check_partition(current_date)
                current_date += timedelta(days=1)

            elif level == "month":
                key = current_date.strftime("%Y%m")
                date_check = current_date.replace(day=1)
                status = check_partition(date_check)
                current_date += relativedelta(months=1)

            elif level == "month_end":
                last_day = current_date + relativedelta(day=31)
                key = last_day.strftime("%Y%m%d")
                status = check_partition(last_day)
                current_date += relativedelta(months=1)

            else:
                raise ValueError("level must be 'day', 'month', or 'month_end'")

            if not status:
                missing.append(key)
                if output in ["detail", "both"]:
                    print(f"{key}: ❌ (MISSING)")

        if output in ["summary"]:
            print(f"\nSummary: {len(missing)} partisi kosong ❌")
            return

        if output in ["both"]:
            print(f"\nSummary: {len(missing)} partisi kosong ❌")

        return missing
    
    else:
        raise ValueError("mode must be 'once', 'backward', 'range', or 'empty'")