"""Iceberg maintenance: expire old snapshots and rewrite small files.

Usage:
  python -m iceberg.maintenance expire [--retention-days N]
  python -m iceberg.maintenance rewrite

Notes:
  pyiceberg's maintenance surface (expire_snapshots, rewrite_data_files)
  is stable as of pyiceberg>=0.7.  If an older version is installed and
  raises AttributeError on either method, pin pyiceberg>=0.7 in
  requirements.txt and re-install.
"""

from __future__ import annotations

import argparse
from datetime import datetime, timedelta, timezone

from .tables import create_or_load_events


def expire_snapshots(retention_days: int = 7) -> None:
    table = create_or_load_events()
    older_than_ms = int(
        (datetime.now(tz=timezone.utc) - timedelta(days=retention_days)).timestamp()
        * 1000
    )
    result = (
        table.expire_snapshots()
        .expire_older_than(older_than_ms)
        .commit()
    )
    print(f"expired snapshots older than {retention_days} day(s): {result}")


def rewrite_data_files() -> None:
    table = create_or_load_events()
    table.rewrite_data_files().commit()
    print("rewrite_data_files completed")


def main() -> None:
    parser = argparse.ArgumentParser(description="Iceberg table maintenance")
    parser.add_argument("action", choices=["expire", "rewrite"])
    parser.add_argument(
        "--retention-days",
        type=int,
        default=7,
        help="Expire snapshots older than this many days (default: 7)",
    )
    args = parser.parse_args()
    if args.action == "expire":
        expire_snapshots(args.retention_days)
    else:
        rewrite_data_files()


if __name__ == "__main__":
    main()