"""Async audit log manager for MongoDB storage with non-blocking writes."""

import asyncio
import hashlib
import logging
from datetime import datetime, timezone
from typing import Optional
from weakref import WeakSet
from motor.motor_asyncio import AsyncIOMotorCollection
import redis.asyncio as aioredis

from .models import (
    AuditLogEntry,
    AuditLogQuery,
    AuditActor,
    AuditTarget,
    AuditChange,
    GuardrailViolation,
)
from .enums import AuditAction, AuditResource, AuditResult, AuditSeverity

logger = logging.getLogger(__name__)


class AuditLogManager:
    """Async manager for audit log operations with non-blocking writes."""

    COLLECTION_NAME = "audit_logs"
    ACTIVITY_COLLECTION_NAME = "user_activity"

    def __init__(
        self,
        collection: AsyncIOMotorCollection,
        activity_collection: AsyncIOMotorCollection = None,
        redis_client: Optional[aioredis.Redis] = None,
    ) -> None:
        self.collection = collection
        self.activity_collection = activity_collection
        self.redis = redis_client
        self._buffer: list[dict] = []
        self._buffer_size = 10  # Flush after N entries
        self._buffer_lock = asyncio.Lock()
        # Track background tasks to prevent garbage collection
        self._pending_tasks: WeakSet[asyncio.Task] = WeakSet()

    @staticmethod
    def _hash_api_key(api_key: str) -> str:
        """Hash API key for storage (show last 4 chars for identification)."""
        if not api_key:
            return ""
        hashed = hashlib.sha256(api_key.encode()).hexdigest()[:8]
        suffix = api_key[-4:] if len(api_key) >= 4 else api_key
        return f"{hashed}...{suffix}"

    def _fire_and_forget(self, coro) -> None:
        """Schedule a coroutine to run in the background without blocking."""
        try:
            loop = asyncio.get_running_loop()
            task = loop.create_task(coro)
            self._pending_tasks.add(task)
            # Add callback to log errors silently
            task.add_done_callback(self._handle_task_result)
        except RuntimeError:
            # No running event loop - fall back to synchronous behavior
            logger.warning("No event loop available for background audit logging")

    def _handle_task_result(self, task: asyncio.Task) -> None:
        """Handle completed background task, log any errors."""
        try:
            task.result()
        except asyncio.CancelledError:
            pass
        except Exception as e:
            logger.error(f"Audit log write failed: {e}")

    async def _write_entry(self, doc: dict) -> Optional[str]:
        """Internal method to write a single entry to MongoDB."""
        try:
            result = await self.collection.insert_one(doc)
            return str(result.inserted_id)
        except Exception as e:
            logger.error(f"Failed to write audit log: {e}")
            return None

    def log_async(self, entry: AuditLogEntry) -> None:
        """
        Write a single audit log entry in the background (non-blocking).
        This is the recommended method for most use cases.
        """
        doc = entry.to_document()
        self._fire_and_forget(self._write_entry(doc))

    async def log(self, entry: AuditLogEntry) -> None:
        """
        Write a single audit log entry in the background (non-blocking).
        Returns immediately, does not wait for write to complete.
        """
        self.log_async(entry)

    async def log_sync(self, entry: AuditLogEntry) -> str:
        """
        Write a single audit log entry and wait for completion.
        Use only when you need confirmation that the log was written.
        """
        doc = entry.to_document()
        result = await self.collection.insert_one(doc)
        return str(result.inserted_id)

    def log_buffered(self, entry: AuditLogEntry) -> None:
        """
        Buffer an audit log entry for batch insert (non-blocking).
        Entries are flushed when buffer reaches _buffer_size.
        """
        self._fire_and_forget(self._add_to_buffer(entry))

    async def _add_to_buffer(self, entry: AuditLogEntry) -> None:
        """Add entry to buffer and flush if needed."""
        async with self._buffer_lock:
            self._buffer.append(entry.to_document())
            if len(self._buffer) >= self._buffer_size:
                await self._flush_buffer()

    async def _flush_buffer(self) -> None:
        """Flush buffered entries to database."""
        if not self._buffer:
            return
        entries = self._buffer.copy()
        self._buffer.clear()
        try:
            await self.collection.insert_many(entries, ordered=False)
        except Exception as e:
            logger.error(f"Batch audit log write failed: {e}")
            # On failure, try individual inserts
            for entry in entries:
                try:
                    await self.collection.insert_one(entry)
                except Exception:
                    pass

    async def flush(self) -> None:
        """Manually flush any buffered entries."""
        async with self._buffer_lock:
            await self._flush_buffer()

    async def shutdown(self) -> None:
        """
        Graceful shutdown - flush buffer and wait for pending writes.
        Call this on application shutdown.
        """
        # Flush any buffered entries
        await self.flush()
        # Wait for any pending background tasks (with timeout)
        pending = [t for t in self._pending_tasks if not t.done()]
        if pending:
            await asyncio.wait(pending, timeout=5.0)

    async def query(self, query: AuditLogQuery) -> list[dict]:
        """Query audit logs with filters."""
        filters = {}

        if query.org_id:
            filters["actor.org_id"] = query.org_id
        if query.user_id:
            filters["actor.user_id"] = query.user_id
        if query.action:
            filters["action"] = query.action
        if query.resource_type:
            filters["target.resource_type"] = query.resource_type
        if query.resource_id:
            filters["target.resource_id"] = query.resource_id
        if query.result:
            filters["result"] = query.result
        if query.severity:
            filters["severity"] = query.severity
        if query.session_id:
            filters["session_id"] = query.session_id

        # Time range filters
        if query.start_time or query.end_time:
            filters["timestamp"] = {}
            if query.start_time:
                filters["timestamp"]["$gte"] = query.start_time
            if query.end_time:
                filters["timestamp"]["$lte"] = query.end_time

        cursor = (
            self.collection.find(filters)
            .sort("timestamp", -1)
            .skip(query.offset)
            .limit(query.limit)
        )

        results = []
        async for doc in cursor:
            doc["_id"] = str(doc["_id"])
            results.append(doc)
        return results

    async def get_by_resource(
        self,
        org_id: str,
        resource_type: AuditResource,
        resource_id: str,
        limit: int = 50,
    ) -> list[dict]:
        """Get audit history for a specific resource."""
        return await self.query(
            AuditLogQuery(
                org_id=org_id,
                resource_type=resource_type,
                resource_id=resource_id,
                limit=limit,
            )
        )

    async def get_by_session(self, session_id: str, limit: int = 100) -> list[dict]:
        """Get all audit events for a session."""
        return await self.query(AuditLogQuery(session_id=session_id, limit=limit))

    async def count(self, query: AuditLogQuery) -> int:
        """Count matching audit log entries."""
        filters = {}
        if query.org_id:
            filters["actor.org_id"] = query.org_id
        if query.user_id:
            filters["actor.user_id"] = query.user_id
        if query.action:
            filters["action"] = query.action
        if query.resource_type:
            filters["target.resource_type"] = query.resource_type
        if query.resource_id:
            filters["target.resource_id"] = query.resource_id
        if query.result:
            filters["result"] = query.result
        if query.severity:
            filters["severity"] = query.severity
        if query.session_id:
            filters["session_id"] = query.session_id

        # Time range filters
        if query.start_time or query.end_time:
            filters["timestamp"] = {}
            if query.start_time:
                filters["timestamp"]["$gte"] = query.start_time
            if query.end_time:
                filters["timestamp"]["$lte"] = query.end_time

        return await self.collection.count_documents(filters)

    # Convenience methods for common audit scenarios (all non-blocking)

    def log_create(
        self,
        org_id: str,
        resource_type: AuditResource,
        resource_id: str,
        user_id: Optional[str] = None,
        api_key: Optional[str] = None,
        ip_address: Optional[str] = None,
        resource_name: Optional[str] = None,
        metadata: Optional[dict] = None,
        result: AuditResult = AuditResult.SUCCESS,
        error_message: Optional[str] = None,
    ) -> None:
        """Log a CREATE action (non-blocking)."""
        entry = AuditLogEntry(
            actor=AuditActor(
                org_id=org_id,
                user_id=user_id,
                api_key_hash=self._hash_api_key(api_key) if api_key else None,
                ip_address=ip_address,
            ),
            action=AuditAction.CREATE,
            target=AuditTarget(
                resource_type=resource_type,
                resource_id=resource_id,
                resource_name=resource_name,
            ),
            result=result,
            error_message=error_message,
            metadata=metadata,
        )
        self.log_async(entry)

    def log_update(
        self,
        org_id: str,
        resource_type: AuditResource,
        resource_id: str,
        changes: list[AuditChange],
        user_id: Optional[str] = None,
        api_key: Optional[str] = None,
        ip_address: Optional[str] = None,
        resource_name: Optional[str] = None,
        metadata: Optional[dict] = None,
        result: AuditResult = AuditResult.SUCCESS,
        error_message: Optional[str] = None,
    ) -> None:
        """Log an UPDATE action with changes (non-blocking)."""
        entry = AuditLogEntry(
            actor=AuditActor(
                org_id=org_id,
                user_id=user_id,
                api_key_hash=self._hash_api_key(api_key) if api_key else None,
                ip_address=ip_address,
            ),
            action=AuditAction.UPDATE,
            target=AuditTarget(
                resource_type=resource_type,
                resource_id=resource_id,
                resource_name=resource_name,
            ),
            changes=changes,
            result=result,
            error_message=error_message,
            metadata=metadata,
        )
        self.log_async(entry)

    def log_delete(
        self,
        org_id: str,
        resource_type: AuditResource,
        resource_id: str,
        user_id: Optional[str] = None,
        api_key: Optional[str] = None,
        ip_address: Optional[str] = None,
        resource_name: Optional[str] = None,
        metadata: Optional[dict] = None,
        result: AuditResult = AuditResult.SUCCESS,
        error_message: Optional[str] = None,
    ) -> None:
        """Log a DELETE action (non-blocking)."""
        entry = AuditLogEntry(
            actor=AuditActor(
                org_id=org_id,
                user_id=user_id,
                api_key_hash=self._hash_api_key(api_key) if api_key else None,
                ip_address=ip_address,
            ),
            action=AuditAction.DELETE,
            target=AuditTarget(
                resource_type=resource_type,
                resource_id=resource_id,
                resource_name=resource_name,
            ),
            result=result,
            error_message=error_message,
            metadata=metadata,
        )
        self.log_async(entry)

    def log_access_denied(
        self,
        org_id: str,
        resource_type: AuditResource,
        action: AuditAction,
        permission_required: str,
        user_id: Optional[str] = None,
        api_key: Optional[str] = None,
        ip_address: Optional[str] = None,
        resource_id: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> None:
        """Log an access denied event (non-blocking)."""
        entry = AuditLogEntry(
            actor=AuditActor(
                org_id=org_id,
                user_id=user_id,
                api_key_hash=self._hash_api_key(api_key) if api_key else None,
                ip_address=ip_address,
            ),
            action=AuditAction.ACCESS_DENIED,
            target=AuditTarget(
                resource_type=resource_type,
                resource_id=resource_id,
            ),
            result=AuditResult.BLOCKED,
            permission_required=permission_required,
            severity=AuditSeverity.HIGH,
            metadata={"attempted_action": action},
        )
        self.log_async(entry)

    def log_guardrail_violation(
        self,
        org_id: str,
        session_id: str,
        violations: list[GuardrailViolation],
        user_id: Optional[str] = None,
        api_key: Optional[str] = None,
        ip_address: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> None:
        """Log guardrail/RAI violations (non-blocking)."""
        # Determine severity based on violation types
        severity = AuditSeverity.MEDIUM
        for v in violations:
            if v.violation_type in ("injection", "nsfw"):
                severity = AuditSeverity.HIGH
                break

        entry = AuditLogEntry(
            actor=AuditActor(
                org_id=org_id,
                user_id=user_id,
                api_key_hash=self._hash_api_key(api_key) if api_key else None,
                ip_address=ip_address,
            ),
            action=AuditAction.EXECUTE,
            target=AuditTarget(resource_type=AuditResource.GUARDRAIL),
            result=AuditResult.BLOCKED,
            session_id=session_id,
            guardrail_violations=violations,
            severity=severity,
            metadata=metadata,
        )
        self.log_async(entry)

    def log_inference(
        self,
        org_id: str,
        agent_id: str,
        session_id: str,
        user_id: Optional[str] = None,
        api_key: Optional[str] = None,
        ip_address: Optional[str] = None,
        result: AuditResult = AuditResult.SUCCESS,
        error_message: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> None:
        """Log an inference/chat execution (non-blocking)."""
        entry = AuditLogEntry(
            actor=AuditActor(
                org_id=org_id,
                user_id=user_id,
                api_key_hash=self._hash_api_key(api_key) if api_key else None,
                ip_address=ip_address,
            ),
            action=AuditAction.EXECUTE,
            target=AuditTarget(
                resource_type=AuditResource.INFERENCE,
                resource_id=agent_id,
            ),
            session_id=session_id,
            result=result,
            error_message=error_message,
            metadata=metadata,
        )
        self.log_async(entry)

    def log_execute(
        self,
        org_id: str,
        resource_type: AuditResource,
        resource_id: str,
        user_id: Optional[str] = None,
        api_key: Optional[str] = None,
        ip_address: Optional[str] = None,
        resource_name: Optional[str] = None,
        result: AuditResult = AuditResult.SUCCESS,
        error_message: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> None:
        """Log an EXECUTE action (non-blocking)."""
        entry = AuditLogEntry(
            actor=AuditActor(
                org_id=org_id,
                user_id=user_id,
                api_key_hash=self._hash_api_key(api_key) if api_key else None,
                ip_address=ip_address,
            ),
            action=AuditAction.EXECUTE,
            target=AuditTarget(
                resource_type=resource_type,
                resource_id=resource_id,
                resource_name=resource_name,
            ),
            result=result,
            error_message=error_message,
            metadata=metadata,
        )
        self.log_async(entry)

    def log_parse(
        self,
        org_id: str,
        resource_type: AuditResource,
        resource_id: str,
        user_id: Optional[str] = None,
        api_key: Optional[str] = None,
        ip_address: Optional[str] = None,
        resource_name: Optional[str] = None,
        result: AuditResult = AuditResult.SUCCESS,
        error_message: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> None:
        """Log a PARSE action (non-blocking)."""
        entry = AuditLogEntry(
            actor=AuditActor(
                org_id=org_id,
                user_id=user_id,
                api_key_hash=self._hash_api_key(api_key) if api_key else None,
                ip_address=ip_address,
            ),
            action=AuditAction.PARSE,
            target=AuditTarget(
                resource_type=resource_type,
                resource_id=resource_id,
                resource_name=resource_name,
            ),
            result=result,
            error_message=error_message,
            metadata=metadata,
        )
        self.log_async(entry)

    def log_train(
        self,
        org_id: str,
        resource_type: AuditResource,
        resource_id: str,
        user_id: Optional[str] = None,
        api_key: Optional[str] = None,
        ip_address: Optional[str] = None,
        resource_name: Optional[str] = None,
        result: AuditResult = AuditResult.SUCCESS,
        error_message: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> None:
        """Log a TRAIN action (non-blocking)."""
        entry = AuditLogEntry(
            actor=AuditActor(
                org_id=org_id,
                user_id=user_id,
                api_key_hash=self._hash_api_key(api_key) if api_key else None,
                ip_address=ip_address,
            ),
            action=AuditAction.TRAIN,
            target=AuditTarget(
                resource_type=resource_type,
                resource_id=resource_id,
                resource_name=resource_name,
            ),
            result=result,
            error_message=error_message,
            metadata=metadata,
        )
        self.log_async(entry)

    def log_upload(
        self,
        org_id: str,
        resource_type: AuditResource,
        resource_id: str,
        user_id: Optional[str] = None,
        api_key: Optional[str] = None,
        ip_address: Optional[str] = None,
        resource_name: Optional[str] = None,
        result: AuditResult = AuditResult.SUCCESS,
        error_message: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> None:
        """Log an UPLOAD action (non-blocking)."""
        entry = AuditLogEntry(
            actor=AuditActor(
                org_id=org_id,
                user_id=user_id,
                api_key_hash=self._hash_api_key(api_key) if api_key else None,
                ip_address=ip_address,
            ),
            action=AuditAction.UPLOAD,
            target=AuditTarget(
                resource_type=resource_type,
                resource_id=resource_id,
                resource_name=resource_name,
            ),
            result=result,
            error_message=error_message,
            metadata=metadata,
        )
        self.log_async(entry)

    def log_download(
        self,
        org_id: str,
        resource_type: AuditResource,
        resource_id: str,
        user_id: Optional[str] = None,
        api_key: Optional[str] = None,
        ip_address: Optional[str] = None,
        resource_name: Optional[str] = None,
        result: AuditResult = AuditResult.SUCCESS,
        error_message: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> None:
        """Log a DOWNLOAD action (non-blocking)."""
        entry = AuditLogEntry(
            actor=AuditActor(
                org_id=org_id,
                user_id=user_id,
                api_key_hash=self._hash_api_key(api_key) if api_key else None,
                ip_address=ip_address,
            ),
            action=AuditAction.DOWNLOAD,
            target=AuditTarget(
                resource_type=resource_type,
                resource_id=resource_id,
                resource_name=resource_name,
            ),
            result=result,
            error_message=error_message,
            metadata=metadata,
        )
        self.log_async(entry)

    def log_reset(
        self,
        org_id: str,
        resource_type: AuditResource,
        resource_id: str,
        user_id: Optional[str] = None,
        api_key: Optional[str] = None,
        ip_address: Optional[str] = None,
        resource_name: Optional[str] = None,
        result: AuditResult = AuditResult.SUCCESS,
        error_message: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> None:
        """Log a RESET action (non-blocking)."""
        entry = AuditLogEntry(
            actor=AuditActor(
                org_id=org_id,
                user_id=user_id,
                api_key_hash=self._hash_api_key(api_key) if api_key else None,
                ip_address=ip_address,
            ),
            action=AuditAction.RESET,
            target=AuditTarget(
                resource_type=resource_type,
                resource_id=resource_id,
                resource_name=resource_name,
            ),
            result=result,
            error_message=error_message,
            metadata=metadata,
        )
        self.log_async(entry)
    
    def log_share(
        self,
        org_id: str,
        resource_type: AuditResource,
        resource_id: str,
        user_id: Optional[str] = None,
        api_key: Optional[str] = None,
        ip_address: Optional[str] = None,
        resource_name: Optional[str] = None,
        result: AuditResult = AuditResult.SUCCESS,
        error_message: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> None:
        """Log a SHARE action (non-blocking)."""
        entry = AuditLogEntry(
            actor=AuditActor(
                org_id=org_id,
                user_id=user_id,
                api_key_hash=self._hash_api_key(api_key) if api_key else None,
                ip_address=ip_address,
            ),
            action=AuditAction.SHARE,
            target=AuditTarget(
                resource_type=resource_type,
                resource_id=resource_id,
                resource_name=resource_name,
            ),
            result=result,
            error_message=error_message,
            metadata=metadata,
        )
        self.log_async(entry)

    def log_sign_in(
        self,
        org_id: str,
        resource_type: AuditResource,
        resource_id: str,
        user_id: Optional[str] = None,
        api_key: Optional[str] = None,
        ip_address: Optional[str] = None,
        resource_name: Optional[str] = None,
        result: AuditResult = AuditResult.SUCCESS,
        error_message: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> None:
        """Log a LOGIN action (non-blocking)."""
        entry = AuditLogEntry(
            actor=AuditActor(
                org_id=org_id,
                user_id=user_id,
                api_key_hash=self._hash_api_key(api_key) if api_key else None,
                ip_address=ip_address,
            ),
            action=AuditAction.LOGIN,
            target=AuditTarget(
                resource_type=resource_type,
                resource_id=resource_id,
                resource_name=resource_name,
            ),
            result=result,
            error_message=error_message,
            metadata=metadata,
        )
        self.log_async(entry)

    def log_sign_out(
        self,
        org_id: str,
        resource_type: AuditResource,
        resource_id: str,
        user_id: Optional[str] = None,
        api_key: Optional[str] = None,
        ip_address: Optional[str] = None,
        resource_name: Optional[str] = None,
        result: AuditResult = AuditResult.SUCCESS,
        error_message: Optional[str] = None,
        metadata: Optional[dict] = None,
    ) -> None:
        """Log a LOGOUT action (non-blocking)."""
        entry = AuditLogEntry(
            actor=AuditActor(
                org_id=org_id,
                user_id=user_id,
                api_key_hash=self._hash_api_key(api_key) if api_key else None,
                ip_address=ip_address,
            ),
            action=AuditAction.LOGOUT,
            target=AuditTarget(
                resource_type=resource_type,
                resource_id=resource_id,
                resource_name=resource_name,
            ),
            result=result,
            error_message=error_message,
            metadata=metadata,
        )
        self.log_async(entry)

    # ==================== Auth & DAU/MAU Tracking ====================

    @staticmethod
    def _get_user_identifier(
        user_id: Optional[str],
        api_key: Optional[str],
        org_id: str,
    ) -> str:
        """
        Get the best available user identifier for DAU/MAU tracking.
        Priority: user_id > api_key (hashed) > org_id (fallback)
        """
        if user_id:
            return f"user:{user_id}"
        if api_key:
            key_hash = hashlib.sha256(api_key.encode()).hexdigest()[:16]
            return f"apikey:{key_hash}"
        return f"org:{org_id}"

    async def _track_activity(
        self,
        org_id: str,
        user_identifier: str,
    ) -> bool:
        """Internal method to track activity using Redis dedup + MongoDB persistence."""
        now = datetime.now(timezone.utc)
        date_str = now.strftime("%Y-%m-%d")
        month_str = now.strftime("%Y-%m")

        try:
            dau_key = f"audit:{org_id}:dau:{date_str}"
            mau_key = f"audit:{org_id}:mau:{month_str}"

            # Redis SADD — returns 1 if new member, 0 if already exists
            pipe = self.redis.pipeline()
            pipe.sadd(dau_key, user_identifier)
            pipe.sadd(mau_key, user_identifier)
            dau_added, _ = await pipe.execute()

            # Set TTL only on new keys (nx=True won't reset existing TTL)
            pipe = self.redis.pipeline()
            pipe.expire(dau_key, 2 * 86400, nx=True)    # 2 days
            pipe.expire(mau_key, 35 * 86400, nx=True)   # 35 days
            await pipe.execute()

            # Only write to MongoDB if new unique user for today
            if dau_added:
                await self.activity_collection.update_one(
                    {
                        "org_id": org_id,
                        "user_identifier": user_identifier,
                        "date": date_str,
                    },
                    {
                        "$set": {
                            "month": month_str,
                            "last_seen_at": now,
                        },
                        "$setOnInsert": {
                            "first_seen_at": now,
                        },
                        "$inc": {
                            "request_count": 1,
                        },
                    },
                    upsert=True,
                )
            return True
        except Exception as e:
            logger.error(f"Failed to track activity: {e}")
            return False

    def log_auth(
        self,
        org_id: str,
        resource_type: AuditResource,
        user_id: Optional[str] = None,
        api_key: Optional[str] = None,
        ip_address: Optional[str] = None,
        metadata: Optional[dict] = None,
        result: AuditResult = AuditResult.SUCCESS,
        error_message: Optional[str] = None,
    ) -> None:
        """
        Log an AUTH action and track DAU/MAU (non-blocking).

        This method: Tracks unique user activity in user_activity collection for DAU/MAU

        Call this in your get_auth middleware for each authenticated request.
        """
        # Track DAU/MAU (non-blocking)
        user_identifier = self._get_user_identifier(user_id, api_key, org_id)
        self._fire_and_forget(self._track_activity(org_id, user_identifier))

    async def get_dau(
        self,
        org_id: str,
        date: Optional[datetime] = None,
    ) -> int:
        """
        Get Daily Active Users count for an organization.

        Args:
            org_id: Organization ID
            date: Date to query (defaults to today UTC)

        Returns:
            Number of unique active users for the day
        """
        if date is None:
            date = datetime.now(timezone.utc)
        date_str = date.strftime("%Y-%m-%d")

        try:
            # Try Redis first (fast O(1) lookup)
            dau_key = f"audit:{org_id}:dau:{date_str}"
            count = await self.redis.scard(dau_key)
            if count > 0:
                return count
            # Fallback to MongoDB for historical dates where Redis key expired
            return await self.activity_collection.count_documents({
                "org_id": org_id,
                "date": date_str,
            })
        except Exception as e:
            logger.error(f"Failed to get DAU: {e}")
            return 0

    async def get_mau(
        self,
        org_id: str,
        date: Optional[datetime] = None,
    ) -> int:
        """
        Get Monthly Active Users count for an organization.

        Args:
            org_id: Organization ID
            date: Date within the month to query (defaults to current month UTC)

        Returns:
            Number of unique active users for the month
        """
        if date is None:
            date = datetime.now(timezone.utc)
        month_str = date.strftime("%Y-%m")

        try:
            # Try Redis first (fast O(1) lookup)
            mau_key = f"audit:{org_id}:mau:{month_str}"
            count = await self.redis.scard(mau_key)
            if count > 0:
                return count
            # Fallback to MongoDB for historical months where Redis key expired
            pipeline = [
                {"$match": {"org_id": org_id, "month": month_str}},
                {"$group": {"_id": "$user_identifier"}},
                {"$count": "total"},
            ]
            result = await self.activity_collection.aggregate(pipeline).to_list(1)
            return result[0]["total"] if result else 0
        except Exception as e:
            logger.error(f"Failed to get MAU: {e}")
            return 0

    async def get_activity_metrics(
        self,
        org_id: str,
        date: Optional[datetime] = None,
    ) -> dict:
        """
        Get all activity metrics for an organization.

        Returns:
            Dict with dau, mau, and daily_active_rate
        """
        dau = await self.get_dau(org_id, date)
        mau = await self.get_mau(org_id, date)

        return {
            "dau": dau,
            "mau": mau,
            "daily_active_rate": round((dau / mau * 100), 2) if mau > 0 else 0.0,
        }

    async def get_dau_trend(
        self,
        org_id: str,
        days: int = 7,
    ) -> list[dict]:
        """
        Get DAU trend for the last N days.

        Args:
            org_id: Organization ID
            days: Number of days to include (default 7)

        Returns:
            List of {date, count} dicts, sorted by date ascending
        """
        from datetime import timedelta

        now = datetime.now(timezone.utc)
        start_date = now - timedelta(days=days - 1)
        start_str = start_date.strftime("%Y-%m-%d")

        try:
            pipeline = [
                {"$match": {"org_id": org_id, "date": {"$gte": start_str}}},
                {"$group": {"_id": "$date", "count": {"$sum": 1}}},
                {"$sort": {"_id": 1}},
                {"$project": {"date": "$_id", "count": 1, "_id": 0}},
            ]
            return await self.activity_collection.aggregate(pipeline).to_list(days)
        except Exception as e:
            logger.error(f"Failed to get DAU trend: {e}")
            return []

    async def get_mau_trend(
        self,
        org_id: str,
        months: int = 6,
    ) -> list[dict]:
        """
        Get MAU trend for the last N months.

        Args:
            org_id: Organization ID
            months: Number of months to include (default 6)

        Returns:
            List of {month, count} dicts, sorted by month ascending
        """
        from datetime import timedelta

        now = datetime.now(timezone.utc)
        start_date = now - timedelta(days=months * 30)
        start_month = start_date.strftime("%Y-%m")

        try:
            pipeline = [
                {"$match": {"org_id": org_id, "month": {"$gte": start_month}}},
                {"$group": {"_id": {"month": "$month", "user": "$user_identifier"}}},
                {"$group": {"_id": "$_id.month", "count": {"$sum": 1}}},
                {"$sort": {"_id": 1}},
                {"$project": {"month": "$_id", "count": 1, "_id": 0}},
            ]
            return await self.activity_collection.aggregate(pipeline).to_list(months)
        except Exception as e:
            logger.error(f"Failed to get MAU trend: {e}")
            return []

    async def ensure_indexes(self) -> None:
        """Create indexes for efficient querying."""
        # Audit logs indexes
        await self.collection.create_index([("timestamp", -1)])
        await self.collection.create_index([("actor.org_id", 1), ("timestamp", -1)])
        await self.collection.create_index([("actor.user_id", 1), ("timestamp", -1)])
        await self.collection.create_index([("target.resource_type", 1), ("target.resource_id", 1)])
        await self.collection.create_index([("session_id", 1)])
        await self.collection.create_index([("action", 1)])
        await self.collection.create_index([("result", 1)])
        await self.collection.create_index([("severity", 1)])

        # User activity indexes (for DAU/MAU)
        # Unique compound index - ensures one record per user per day
        await self.activity_collection.create_index(
            [("org_id", 1), ("user_identifier", 1), ("date", 1)],
            unique=True,
        )
        # Index for DAU queries
        await self.activity_collection.create_index([("org_id", 1), ("date", 1)])
        # Index for MAU queries
        await self.activity_collection.create_index([("org_id", 1), ("month", 1)])
        # TTL index for automatic cleanup (90 days retention)
        await self.activity_collection.create_index(
            [("last_seen_at", 1)],
            expireAfterSeconds=90 * 24 * 60 * 60,
        )
