from __future__ import annotations

from collections.abc import Iterable
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta
from typing import Protocol


class RetentionError(ValueError):
    pass


class BackupLike(Protocol):
    id: str
    created_at: datetime
    pinned: bool
    tombstoned_at: datetime | None


@dataclass(frozen=True)
class RetentionPolicy:
    keep_last: int = 1
    keep_days: int = 0
    keep_daily: int = 0
    keep_weekly: int = 0
    keep_monthly: int = 0

    @classmethod
    def from_dict(cls, value: dict[str, object]) -> RetentionPolicy:
        allowed = {"keep_last", "keep_days", "keep_daily", "keep_weekly", "keep_monthly"}
        if set(value) - allowed:
            raise RetentionError("retention policy contains unknown keys")
        kwargs: dict[str, int] = {}
        for key in allowed:
            raw = value.get(key, 0 if key != "keep_last" else 1)
            if not isinstance(raw, int) or isinstance(raw, bool) or raw < 0:
                raise RetentionError("retention values must be non-negative integers")
            kwargs[key] = raw
        return cls(**kwargs)


def retained_ids(
    backups: Iterable[BackupLike], policy: RetentionPolicy, now: datetime | None = None
) -> set[str]:
    """Return union retention set; the newest non-tombstoned backup is always protected."""
    reference = (now or datetime.now(UTC)).astimezone(UTC)
    items = sorted(
        (item for item in backups if item.tombstoned_at is None),
        key=lambda item: item.created_at.astimezone(UTC),
        reverse=True,
    )
    if not items:
        return set()
    kept = {items[0].id}
    kept.update(item.id for item in items[: policy.keep_last])
    kept.update(item.id for item in items if item.pinned)
    if policy.keep_days:
        cutoff = reference - timedelta(days=policy.keep_days)
        kept.update(item.id for item in items if item.created_at.astimezone(UTC) >= cutoff)
    for count, key in (
        (policy.keep_daily, lambda stamp: stamp.date()),
        (policy.keep_weekly, lambda stamp: stamp.isocalendar()[:2]),
        (policy.keep_monthly, lambda stamp: (stamp.year, stamp.month)),
    ):
        buckets: set[object] = set()
        for item in items:
            stamp = item.created_at.astimezone(UTC)
            bucket = key(stamp)
            if len(buckets) >= count and bucket not in buckets:
                continue
            buckets.add(bucket)
            if bucket in buckets:
                kept.add(item.id)
    return kept