from dataclasses import dataclass
from datetime import datetime, timezone
from threading import Lock

from app.core.constants import IDEMPOTENCY_TTL


@dataclass
class IdempotencyRecord:
    task_id: str
    created_at: datetime


class IdempotencyStore:
    def __init__(self) -> None:
        self._records: dict[str, IdempotencyRecord] = {}
        self._lock = Lock()

    def get_task_id(self, key: str) -> str | None:
        now = datetime.now(timezone.utc)
        with self._lock:
            self._cleanup_locked(now)
            record = self._records.get(key)
            return record.task_id if record else None

    def put(self, key: str, task_id: str) -> None:
        with self._lock:
            self._records[key] = IdempotencyRecord(
                task_id=task_id,
                created_at=datetime.now(timezone.utc),
            )

    def _cleanup_locked(self, now: datetime) -> None:
        expired = [
            key
            for key, rec in self._records.items()
            if now - rec.created_at > IDEMPOTENCY_TTL
        ]
        for key in expired:
            del self._records[key]