from __future__ import annotations

from dataclasses import dataclass
import time
from typing import Any

from .codec import (
    DEFAULT_COMPATIBILITY_VERSION,
    DEFAULT_SCHEMA_VERSION,
    decode_state,
    encode_state,
)


@dataclass
class _MemoryStateEntry:
    value_blob: bytes
    expires_at: float | None


class MemoryStateBackend:
    def __init__(
        self,
        *,
        compatibility_version: int = DEFAULT_COMPATIBILITY_VERSION,
        schema_version: int = DEFAULT_SCHEMA_VERSION,
        now_fn=time.time,
    ) -> None:
        self._compatibility_version = int(compatibility_version)
        self._schema_version = int(schema_version)
        self._now_fn = now_fn
        self._entries: dict[tuple[str, str], _MemoryStateEntry] = {}

    def get(self, namespace: str, key: str) -> Any | None:
        self._prune_expired()
        entry = self._entries.get((namespace, key))
        if entry is None:
            return None
        return decode_state(
            entry.value_blob,
            expected_compatibility_version=self._compatibility_version,
        )

    def put(
        self,
        namespace: str,
        key: str,
        value: Any,
        ttl_seconds: float | None = None,
    ) -> None:
        expires_at = self._resolve_expiry(ttl_seconds)
        value_blob = encode_state(
            value,
            compatibility_version=self._compatibility_version,
            schema_version=self._schema_version,
        )
        self._entries[(namespace, key)] = _MemoryStateEntry(
            value_blob=value_blob,
            expires_at=expires_at,
        )
        self._prune_expired()

    def delete(self, namespace: str, key: str) -> None:
        self._entries.pop((namespace, key), None)

    def scan(
        self,
        namespace: str,
        *,
        prefix: str | None = None,
        limit: int | None = None,
    ) -> list[tuple[str, Any]]:
        self._prune_expired()
        effective_limit = None if limit is None else max(0, int(limit))
        result: list[tuple[str, Any]] = []
        for (entry_namespace, entry_key), entry in self._entries.items():
            if entry_namespace != namespace:
                continue
            if prefix is not None and not entry_key.startswith(prefix):
                continue
            value = decode_state(
                entry.value_blob,
                expected_compatibility_version=self._compatibility_version,
            )
            result.append((entry_key, value))
            if effective_limit is not None and len(result) >= effective_limit:
                break
        return result

    def close(self) -> None:
        self._entries.clear()

    def _prune_expired(self) -> None:
        now = self._now_fn()
        expired_keys = [
            key
            for key, entry in self._entries.items()
            if entry.expires_at is not None and entry.expires_at <= now
        ]
        for key in expired_keys:
            self._entries.pop(key, None)

    def _resolve_expiry(self, ttl_seconds: float | None) -> float | None:
        if ttl_seconds is None:
            return None
        ttl = float(ttl_seconds)
        if ttl <= 0:
            return self._now_fn()
        return self._now_fn() + ttl


__all__ = ["MemoryStateBackend"]
