"""Secret rotation policies and management."""

from __future__ import annotations

from datetime import datetime, timedelta, timezone
from typing import Callable, Awaitable

from pydantic import BaseModel, Field

from vault.models import Secret, SecretMetadata
from vault.store import SecretStore


class RotationPolicy(BaseModel):
    """Defines when a secret should be rotated."""

    max_age_days: int = 90
    auto_rotate: bool = False


RotationHook = Callable[[str, int], Awaitable[None]]
"""Signature: async def hook(secret_name: str, new_version: int) -> None"""

ValueGenerator = Callable[[str], Awaitable[str]]
"""Signature: async def generator(secret_name: str) -> str  (produces new value)"""


class RotationManager:
    """Checks for expiring secrets and handles rotation."""

    def __init__(
        self,
        store: SecretStore,
        default_policy: RotationPolicy | None = None,
    ) -> None:
        self._store = store
        self._default_policy = default_policy or RotationPolicy()
        self._hooks: list[RotationHook] = []
        self._generators: dict[str, ValueGenerator] = {}

    # -- Hook registration --

    def on_rotation(self, hook: RotationHook) -> RotationHook:
        """Register a callback invoked after every rotation (can be used as decorator)."""
        self._hooks.append(hook)
        return hook

    def register_generator(self, secret_name: str, generator: ValueGenerator) -> None:
        """Register a custom value generator for a specific secret."""
        self._generators[secret_name] = generator

    # -- Core logic --

    async def check_expiring_secrets(
        self,
        *,
        policy: RotationPolicy | None = None,
        within_days: int | None = None,
    ) -> list[SecretMetadata]:
        """Return metadata for secrets that are within the rotation window."""
        p = policy or self._default_policy
        window = within_days if within_days is not None else p.max_age_days
        cutoff = datetime.now(timezone.utc) - timedelta(days=window)
        all_secrets = await self._store.list_secrets()
        expiring: list[SecretMetadata] = []
        for meta in all_secrets:
            if meta.updated_at <= cutoff:
                expiring.append(meta)
            elif meta.expires_at and meta.expires_at <= datetime.now(timezone.utc) + timedelta(days=7):
                expiring.append(meta)
        return expiring

    async def rotate_secret(
        self,
        name: str,
        new_value: str | None = None,
        *,
        created_by: str = "rotation-manager",
    ) -> Secret | None:
        """Rotate a secret to a new value.

        If *new_value* is not supplied, the registered generator for *name* is
        called.  If no generator exists, the rotation fails and returns None.
        """
        if new_value is None:
            generator = self._generators.get(name)
            if generator is None:
                return None
            new_value = await generator(name)

        rotated = await self._store.rotate(name, new_value, created_by=created_by)
        if rotated is not None:
            for hook in self._hooks:
                await hook(rotated.name, rotated.version)
        return rotated

    async def auto_rotate_all(
        self,
        *,
        policy: RotationPolicy | None = None,
        created_by: str = "auto-rotation",
    ) -> list[Secret]:
        """Rotate every secret that has a registered generator and is past its age limit."""
        expiring = await self.check_expiring_secrets(policy=policy)
        rotated: list[Secret] = []
        for meta in expiring:
            if meta.name in self._generators:
                result = await self.rotate_secret(meta.name, created_by=created_by)
                if result is not None:
                    rotated.append(result)
        return rotated
