import json
import logging
import google.cloud.logging
from google.cloud import firestore
from typing import Callable
from datetime import datetime, timezone, timedelta

## Setup Logger
client = google.cloud.logging.Client()
client.setup_logging()
logger = logging.getLogger("uvicorn")
logger.setLevel(logging.INFO)


class GCPCache:
  cache_conn: firestore.Client | None = None

  def init_cache(self, project: str, database: str = "(default)"):
    """
    Initialize the Firestore client.
    Args:
        project: The GCP Project ID where Firestore is hosted.
        database: The Firestore database ID. Defaults to "(default)".
    """
    self.cache_conn = firestore.Client(project=project, database=database)

  def get_cache_instance(self):
    return self.cache_conn

  def _get_doc_ref(self, key: str):
    if not self.cache_conn:
        return None
    # Firestore document IDs cannot contain forward slashes '/'
    # GCP IAM keys like "organization_id:role:email" might have slashes in roles like "roles/owner".
    safe_key = key.replace("/", "_")
    return self.cache_conn.collection('_sdk_cache').document(safe_key)

  def is_exists(self, key: str, value: str, use_cache: bool = True) -> bool | None:
    """
    Read-only cache lookup: checks if a value exists in a cached JSON list.
    Returns True/False on cache hit, or None on cache miss (so the caller
    knows to fall through to the actual computation).

    If use_cache is False or cache_conn is not initialized, returns None.
    """
    if use_cache and self.cache_conn:
        doc_ref = self._get_doc_ref(key)
        doc = doc_ref.get() if doc_ref else None
        if doc and doc.exists:
            data = doc.to_dict()
            expires_at = data.get("expires_at")
            if expires_at and expires_at > datetime.now(timezone.utc):
                cached_value = data.get("value")
                if cached_value:
                    cached_list = json.loads(cached_value)
                    logger.info(f"Cache hit for {key}, verifying value against the cache")
                    return value in cached_list
    return None

  def compute_cache_aside(
      self, 
      key: str, 
      compute_fn: Callable[[], any],
      ex: int,
      use_cache: bool = True
  ) -> any:
    """
    Cache-aside pattern: attempts to retrieve a value from Firestore by key.
    On a cache hit, returns the cached value. On a miss, calls compute_fn
    to produce the result, stores it in Firestore with the given TTL, and returns it.

    If the Firestore connection was not initialized (cache_conn is None/unset),
    the cache is skipped entirely — compute_fn is called directly and its
    result is returned without any caching. This allows callers to use this
    method safely even before init_cache has been called.

    Args:
        key: The cache key to look up and store the result under.
        compute_fn: A callable (takes no arguments) that produces the value
                    to cache. Only invoked on a cache miss.
        ex: Expiration time in seconds for the cached entry.
        use_cache: Whether to use the cache. Defaults to True but is not enforced —
                    callers can pass False to bypass the cache and always invoke compute_fn.

    Returns:
        The cached value (on hit) or the freshly computed value (on miss).
    """
    if use_cache and self.cache_conn:
        doc_ref = self._get_doc_ref(key)
        doc = doc_ref.get() if doc_ref else None
        if doc and doc.exists:
            data = doc.to_dict()
            expires_at = data.get("expires_at")
            if expires_at and expires_at > datetime.now(timezone.utc):
                cached_value = data.get("value")
                if cached_value:
                    logger.info(f"Cache HIT for key: {key}")
                    return json.loads(cached_value)

    result = compute_fn()

    if use_cache and self.cache_conn and result:
        doc_ref = self._get_doc_ref(key)
        expires_at = datetime.now(timezone.utc) + timedelta(seconds=ex)
        if doc_ref:
            doc_ref.set({
                "value": json.dumps(result),
                "expires_at": expires_at
            })
            logger.info(f"Cached result for key: {key} (TTL limit={expires_at.isoformat()})")

    return result


cache_instance: GCPCache = GCPCache()