"""
SecretCache for per-task caching of decrypted secrets.

Designed to be instantiated at the start of a Cloud Task and discarded when done.
This avoids repeated KMS decrypt calls within a single chunk processing.

Usage:
    async with SecretCache.for_workspace(db, workspace_id) as cache:
        fullenrich_key = cache.get("fullenrich", "api_key")
        # Use key for provider API calls...
"""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Dict, Optional, Set, TYPE_CHECKING

import logging


logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    from sqlalchemy.ext.asyncio import AsyncSession


@dataclass
class CachedSecret:
    """A cached decrypted secret with metadata."""
    value: str
    integration: str
    key: str
    key_suffix: str
    cached_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))


class SecretCache:
    """
    In-memory cache for decrypted secrets within a single task.

    This class is designed to be short-lived - instantiate at the start of
    a Cloud Task and discard when the task completes. Do not persist across
    task boundaries.

    Thread Safety:
        This cache is NOT thread-safe. Each task/coroutine should have its own
        SecretCache instance. In practice, Cloud Tasks workers are single-threaded
        per task anyway.
    """

    def __init__(self, workspace_id: str):
        """
        Initialize a secret cache for a specific workspace.

        Args:
            workspace_id: The workspace ID this cache is scoped to
        """
        self.workspace_id = workspace_id
        self._cache: Dict[str, CachedSecret] = {}
        self._preloaded = False

    def _cache_key(self, integration: str, key: str) -> str:
        """Generate cache key from integration and key name."""
        return f"{integration}:{key}"

    def get(self, integration: str, key: str) -> Optional[str]:
        """
        Get a cached secret value.

        Args:
            integration: Integration name (e.g., 'fullenrich')
            key: Key name (e.g., 'api_key')

        Returns:
            Decrypted secret value if cached, None otherwise
        """
        cache_key = self._cache_key(integration, key)
        cached = self._cache.get(cache_key)
        if cached:
            return cached.value
        return None

    def get_secret(self, integration: str, key: str) -> Optional[CachedSecret]:
        """
        Get a cached secret with metadata.

        Args:
            integration: Integration name
            key: Key name

        Returns:
            CachedSecret if found, None otherwise
        """
        cache_key = self._cache_key(integration, key)
        return self._cache.get(cache_key)

    def set(
        self,
        integration: str,
        key: str,
        value: str,
        key_suffix: Optional[str] = None,
    ) -> None:
        """
        Cache a decrypted secret.

        Args:
            integration: Integration name
            key: Key name
            value: Decrypted secret value
            key_suffix: Last 4 chars for display (optional)
        """
        cache_key = self._cache_key(integration, key)
        self._cache[cache_key] = CachedSecret(
            value=value,
            integration=integration,
            key=key,
            key_suffix=key_suffix or value[-4:] if len(value) >= 4 else "****",
        )

    def has(self, integration: str, key: str) -> bool:
        """Check if a secret is cached."""
        return self._cache_key(integration, key) in self._cache

    def remove(self, integration: str, key: str) -> bool:
        """Remove a secret from cache."""
        cache_key = self._cache_key(integration, key)
        if cache_key in self._cache:
            del self._cache[cache_key]
            return True
        return False

    def clear(self) -> None:
        """Clear all cached secrets."""
        self._cache.clear()
        self._preloaded = False

    def cached_integrations(self) -> Set[str]:
        """Get set of integrations that have cached secrets."""
        return {cached.integration for cached in self._cache.values()}

    @property
    def is_preloaded(self) -> bool:
        """Check if secrets have been preloaded from DB."""
        return self._preloaded

    async def preload(
        self,
        db: "AsyncSession",
        integrations: Optional[Set[str]] = None,
    ) -> int:
        """
        Preload secrets from the database into the cache.

        This is the recommended way to populate the cache - call once at
        the start of a task to decrypt all needed secrets in one batch.

        Args:
            db: Async database session
            integrations: Optional set of integrations to preload.
                         If None, preloads all available secrets.

        Returns:
            Number of secrets loaded
        """
        from .secret_storage import get_saas_secret_backend

        backend = get_saas_secret_backend()
        secrets = await backend.list_secrets(db, self.workspace_id)

        count = 0
        for secret_info in secrets:
            integration = secret_info["integration"]
            key = secret_info["key_name"]

            # Skip if filtering by integrations and not in set
            if integrations and integration not in integrations:
                continue

            # Get the full secret (with decryption)
            secret = await backend.get_secret(db, integration, key, self.workspace_id)
            if secret:
                self.set(
                    integration=integration,
                    key=key,
                    value=secret.value,
                    key_suffix=secret_info.get("key_suffix"),
                )
                count += 1

        self._preloaded = True
        logger.debug(
            f"SecretCache preloaded {count} secrets for workspace {self.workspace_id}"
        )
        return count

    async def __aenter__(self) -> "SecretCache":
        """Async context manager entry."""
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        """Async context manager exit - clears secrets from memory."""
        self.clear()

    @classmethod
    async def for_workspace(
        cls,
        db: "AsyncSession",
        workspace_id: str,
        integrations: Optional[Set[str]] = None,
    ) -> "SecretCache":
        """
        Factory method to create and preload a SecretCache.

        Args:
            db: Async database session
            workspace_id: Workspace ID
            integrations: Optional set of integrations to preload

        Returns:
            Preloaded SecretCache instance
        """
        cache = cls(workspace_id)
        await cache.preload(db, integrations)
        return cache


class SecretCacheContext:
    """
    Context manager wrapper for SecretCache with automatic preloading.

    Example:
        async with SecretCacheContext(db, workspace_id, {"fullenrich"}) as cache:
            api_key = cache.get("fullenrich", "api_key")
            # Use api_key...
    """

    def __init__(
        self,
        db: "AsyncSession",
        workspace_id: str,
        integrations: Optional[Set[str]] = None,
    ):
        self.db = db
        self.workspace_id = workspace_id
        self.integrations = integrations
        self._cache: Optional[SecretCache] = None

    async def __aenter__(self) -> SecretCache:
        self._cache = await SecretCache.for_workspace(
            self.db, self.workspace_id, self.integrations
        )
        return self._cache

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        if self._cache:
            self._cache.clear()
            self._cache = None
