#!/usr/bin/env python3
import argparse
from collections import defaultdict
from datetime import datetime, timedelta
from pathlib import Path
import logging
import re

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)


def apply_retention_policy(backup_files: list, now: datetime) -> list:
    """Apply grandfather-father-son retention policy"""

    daily_files = []
    weekly_files = []
    monthly_files = []
    for file, backup_date in backup_files:
        age = now - backup_date

        if age < timedelta(days=3):
            daily_files.append((file, backup_date))
        elif age < timedelta(days=30):
            weekly_files.append((file, backup_date))
        elif age < timedelta(days=90):
            monthly_files.append((file, backup_date))

    to_keep = set()

    # Keep 1 backup per day
    to_keep.update(file for file, _ in daily_files)

    # Keep one backup per week for weekly period
    grouped_by_week = defaultdict(list)
    for file, backup_date in weekly_files:
        week_key = backup_date.strftime("%Y-%W")
        grouped_by_week[week_key].append((file, backup_date))

    for week_backups in grouped_by_week.values():
        to_keep.add(max(week_backups, key=lambda x: x[1])[0])

    # Keep one backup per month for monthly period
    grouped_by_month = defaultdict(list)
    for file, backup_date in monthly_files:
        month_key = backup_date.strftime("%Y-%m")
        grouped_by_month[month_key].append((file, backup_date))

    for month_backups in grouped_by_month.values():
        to_keep.add(max(month_backups, key=lambda x: x[1])[0])

    return [file for file, _ in backup_files if file not in to_keep]


def prune_backups(backup_dir: Path) -> None:
    logger.info(f"Pruning old backups in {backup_dir}")

    if not backup_dir.exists():
        return

    backup_files = []

    for file in backup_dir.glob("*.zip"):
        match = re.search(r"^(\d{4}\D\d{2}\D\d{2}\D\d{2}\D\d{2}\D\d{2})", file.name)
        if match:
            try:
                date_str = match.group(1).replace("-", "_")
                backup_date = datetime.strptime(date_str, "%Y_%m_%d_%H_%M_%S")
                backup_files.append((file, backup_date))
            except ValueError:
                continue

    backup_files.sort(key=lambda x: x[1])

    now = datetime.now()
    to_delete = apply_retention_policy(backup_files, now)

    for file in to_delete:
        logger.info(f"Deleting old backup: {file}")
        file.unlink()

    kept_count = len(backup_files) - len(to_delete)
    logger.info(f"Kept {kept_count} backups, pruned {len(to_delete)} old backups")


def parse_args():
    parser = argparse.ArgumentParser(description="Prune old backups based on retention policy", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("--backup-dir", type=Path, required=True, help="Directory containing backup files")
    return parser.parse_args()


def main():
    args = parse_args()
    backup_dir = args.backup_dir

    if not backup_dir.exists():
        logger.error(f"Backup directory does not exist: {backup_dir}")
        return

    prune_backups(backup_dir)
    logger.info("Backup pruning completed successfully")


if __name__ == "__main__":
    main()