"""
Entitlements and team context data loader.

Orchestrates loading of organization entitlements, user context, and team data
from MongoDB with Redis caching.
"""

from typing import Any, Dict, List, Optional

from ..cache.redis_client import redis_client
from ..utils.logger import logger
from .mongo_client import mongo_client


class EntitlementsLoader:
    """
    Loads entitlements and team context data with caching.

    Provides a high-level interface for loading organization entitlements,
    user context, and team details with intelligent caching strategies.
    """

    async def load_organization_entitlements(
        self, stytch_org_id: str
    ) -> Optional[Dict[str, Any]]:
        """
        Load organization entitlements with 1-hour caching.

        Args:
            stytch_org_id: Stytch organization identifier

        Returns:
            Dict containing:
                - entitlements: List[str]
                - subscription_tier: str
                - subscription_limits: Dict[str, int]
                - mongo_organization_id: str (MongoDB ObjectId as string)
            Returns None if MongoDB is not configured or organization not found
        """
        try:
            # Try cache first
            cached_data = await redis_client.get_cached_organization_entitlements(
                stytch_org_id
            )
            if cached_data:
                logger.debug(
                    f"Using cached organization entitlements for org: {stytch_org_id}"
                )
                return cached_data

            # Fallback to MongoDB
            org_doc = await mongo_client.get_organization(stytch_org_id)
            if not org_doc:
                logger.debug(
                    f"Organization not found in MongoDB for org: {stytch_org_id}"
                )
                return None

            # Extract entitlements data
            raw_id = org_doc.get("_id")
            entitlements_data = {
                "entitlements": org_doc.get("entitlements", []),
                "subscription_tier": org_doc.get("subscription_tier"),
                "subscription_limits": org_doc.get("subscription_limits", {}),
                "mongo_organization_id": str(raw_id) if raw_id is not None else None,
            }

            # Cache for 1 hour
            await redis_client.cache_organization_entitlements(
                stytch_org_id, entitlements_data
            )

            logger.debug(
                f"Loaded organization entitlements from MongoDB for org: {stytch_org_id}"
            )
            return entitlements_data

        except Exception as e:
            logger.warning(
                f"Failed to load organization entitlements: {str(e)}. "
                f"Continuing without entitlements data."
            )
            return None

    async def load_user_context(
        self, stytch_member_id: str, stytch_org_id: str
    ) -> Optional[Dict[str, Any]]:
        """
        Load user context with team information (v2.3.0+).

        This method uses a single aggregation query to fetch all teams the user
        has access to in the current organization, eliminating the need for
        multi-query validation and auto-correction logic.

        Args:
            stytch_member_id: Stytch member identifier
            stytch_org_id: Stytch organization identifier (REQUIRED in v2.3.0+)

        Returns:
            Dict containing:
                - available_teams: List of team info dicts (v2.3.0+)
                - mongo_user_id: str (MongoDB ObjectId as string)
            Returns None if MongoDB is not configured or user not found

        Note:
            v2.3.0+: stytch_org_id is now required. The deprecated fields
            current_team_id and current_team_name are no longer populated as
            the users.current_team_id field has been removed from the schema.
        """
        try:
            # Pass org context for multi-org user lookup
            user_doc = await mongo_client.get_user(stytch_member_id, stytch_org_id)
            if not user_doc:
                logger.debug(
                    f"User not found in MongoDB for member: {stytch_member_id} "
                    f"in org: {stytch_org_id}"
                )
                return None

            user_object_id = user_doc.get("_id")
            mongo_user_id_str = str(user_object_id) if user_object_id else None

            # Load all teams for this user in this org (v2.3.0+)
            available_teams = await self._load_available_teams(
                user_object_id, stytch_org_id
            )

            user_context = {
                "available_teams": available_teams,  # v2.3.0+
                "mongo_user_id": mongo_user_id_str,
                # Deprecated fields (v2.3.0+) - always None, kept for backwards compat
                "current_team_id": None,
                "current_team_name": None,
            }

            logger.debug(
                f"Loaded user context from MongoDB for member: {stytch_member_id} "
                f"in org: {stytch_org_id} (found {len(available_teams)} teams)"
            )
            return user_context

        except Exception as e:
            logger.warning(
                f"Failed to load user context: {str(e)}. "
                f"Continuing without user context data."
            )
            return None

    async def _load_available_teams(
        self, user_id: Any, stytch_org_id: str
    ) -> List[Dict[str, Any]]:
        """
        Load all teams user has access to in the specified organization (v2.3.0+).

        Uses a single MongoDB aggregation pipeline to fetch teams with roles and
        permissions in one query. This replaces the old validation logic that
        required 3-4 queries.

        Args:
            user_id: MongoDB ObjectId of the user
            stytch_org_id: Stytch organization identifier

        Returns:
            List of team info dicts, each containing:
                - id: str (team ObjectId as string)
                - name: str (team name)
                - organization_id: str (org ObjectId as string)
                - role: str or None (user's role in team)
                - permissions: List[str] (team-scoped permissions)
        """
        try:
            db = await mongo_client._get_client()
            if db is None:
                return []

            # Get MongoDB organization ObjectId
            org_doc = await mongo_client.get_organization(stytch_org_id)
            if not org_doc:
                logger.warning(
                    f"Organization not found for stytch_org_id: {stytch_org_id}"
                )
                return []

            mongo_org_id = org_doc.get("_id")

            # Single aggregation query to get all teams with roles/permissions
            teams_pipeline = [
                # Match active team memberships for this user in this org
                {
                    "$match": {
                        "user_id": user_id,
                        "organization_id": mongo_org_id,
                        "status": "active",
                    }
                },
                # Join with teams collection
                {
                    "$lookup": {
                        "from": "teams",
                        "localField": "team_id",
                        "foreignField": "_id",
                        "as": "team_doc",
                    }
                },
                {"$unwind": "$team_doc"},
                # Join with roles collection
                {
                    "$lookup": {
                        "from": "roles",
                        "localField": "role_id",
                        "foreignField": "_id",
                        "as": "role_doc",
                    }
                },
                {
                    "$unwind": {
                        "path": "$role_doc",
                        "preserveNullAndEmptyArrays": True,  # Some memberships may not have roles
                    }
                },
                # Project only needed fields
                {
                    "$project": {
                        "team_id": "$team_doc._id",
                        "team_name": "$team_doc.name",
                        "organization_id": "$team_doc.organization",
                        "role_name": "$role_doc.name",
                        "permissions": "$role_doc.permissions",
                    }
                },
            ]

            team_memberships = (
                await db["user_team_memberships"]
                .aggregate(teams_pipeline)
                .to_list(length=None)
            )

            # Build team info list
            available_teams = [
                {
                    "id": str(m["team_id"]),
                    "name": m["team_name"],
                    "organization_id": str(m["organization_id"]),
                    "role": m.get("role_name"),
                    "permissions": m.get("permissions")
                    or [],  # Handle None from roles without permissions
                }
                for m in team_memberships
            ]

            # Sort by name for consistent ordering
            available_teams.sort(key=lambda t: t["name"])

            logger.debug(
                f"Found {len(available_teams)} teams for user {user_id} "
                f"in org {stytch_org_id}"
            )

            return available_teams

        except Exception as e:
            logger.error(f"Error loading available teams: {e}", exc_info=True)
            return []

    async def load_complete_session_data(
        self, stytch_org_id: str, stytch_member_id: str
    ) -> Dict[str, Any]:
        """
        Load both organization entitlements and user context in one call (v2.3.0+).

        Args:
            stytch_org_id: Stytch organization identifier
            stytch_member_id: Stytch member identifier

        Returns:
            Dict containing all entitlements and user context data including
            available_teams list (v2.3.0+).
            Returns empty dict values if MongoDB is not configured.
        """
        # Load organization entitlements and user context in parallel
        import asyncio

        org_data_task = self.load_organization_entitlements(stytch_org_id)
        # Pass organization context to user context loader for team loading
        user_data_task = self.load_user_context(stytch_member_id, stytch_org_id)

        results = await asyncio.gather(
            org_data_task, user_data_task, return_exceptions=True
        )

        # Handle exceptions from parallel loading
        org_data: Optional[Dict[str, Any]] = None
        user_data: Optional[Dict[str, Any]] = None

        if isinstance(results[0], Exception):
            logger.warning(f"Exception loading organization data: {results[0]}")
        elif results[0] is not None:
            org_data = results[0]  # type: ignore[assignment]

        if isinstance(results[1], Exception):
            logger.warning(f"Exception loading user data: {results[1]}")
        elif results[1] is not None:
            user_data = results[1]  # type: ignore[assignment]

        # Combine results with safe defaults
        return {
            # Organization entitlements (default to None if not available)
            "entitlements": org_data.get("entitlements") if org_data else None,
            "subscription_tier": (
                org_data.get("subscription_tier") if org_data else None
            ),
            "subscription_limits": (
                org_data.get("subscription_limits") if org_data else None
            ),
            "mongo_organization_id": (
                org_data.get("mongo_organization_id") if org_data else None
            ),
            # User context (v2.3.0+ includes available_teams)
            "available_teams": (
                user_data.get("available_teams", []) if user_data else []
            ),
            "mongo_user_id": user_data.get("mongo_user_id") if user_data else None,
            # Deprecated fields (v2.3.0+) - always None, kept for backwards compat during transition
            "current_team_id": None,
            "current_team_name": None,
        }


# Global entitlements loader instance
entitlements_loader = EntitlementsLoader()
