"""
Planning Service - Internal service for task planning with direct DB access

This service provides access to agents, teams, and context graph data using:
- Direct SQLAlchemy queries for agents/teams (no HTTP overhead)
- ContextGraphClient for external context graph API (graph.kubiya.ai)

All queries are org-scoped to ensure proper authorization.
"""

import structlog
from typing import List, Dict, Any, Optional
from sqlalchemy.orm import Session
from uuid import UUID

from control_plane_api.app.models.agent import Agent
from control_plane_api.app.models.team import Team, TeamStatus
from control_plane_api.app.services.context_graph_client import ContextGraphClient

logger = structlog.get_logger(__name__)


class PlanningService:
    """
    Internal service for accessing planning resources.

    Uses direct DB queries for agents/teams (fast, no HTTP overhead)
    Uses ContextGraphClient for context graph (external service, HTTP correct)
    """

    def __init__(
        self,
        db: Session,
        organization_id: str,
        api_token: str,
        context_graph_base: str = "https://graph.kubiya.ai"
    ):
        """
        Initialize planning service.

        Args:
            db: SQLAlchemy database session
            organization_id: Organization ID for filtering
            api_token: API token for context graph auth
            context_graph_base: Context graph API base URL
        """
        self.db = db
        self.organization_id = organization_id
        self.api_token = api_token

        # Initialize context graph client for external queries
        self.context_graph_client = ContextGraphClient(
            api_base=context_graph_base,
            api_token=api_token,
            organization_id=organization_id
        )

        logger.info(
            "planning_service_initialized",
            organization_id=organization_id[:8]
        )

    def _agent_to_dict(self, agent: Agent) -> Dict[str, Any]:
        """Convert Agent model to dictionary for planning (status field excluded)."""
        return {
            "id": str(agent.id),
            "name": agent.name,
            "description": agent.description,
            "model_id": agent.model_id or "default",
            "capabilities": agent.capabilities or [],
            "runtime": agent.runtime,
            "team_id": str(agent.team_id) if agent.team_id else None,
            "execution_environment": agent.execution_environment or {},
        }

    def _team_to_dict(self, team: Team, agents_map: Dict[UUID, Agent] = None) -> Dict[str, Any]:
        """Convert Team model to dictionary.

        Args:
            team: Team model instance
            agents_map: Optional pre-loaded map of agent_id -> Agent (avoids N+1 queries)
        """
        # Get agent count from configuration
        config = team.configuration or {}
        member_ids = config.get("member_ids", [])

        # Get agents if they exist
        agents = []
        if member_ids:
            # Use pre-loaded agents if provided (no query), otherwise fall back to query
            if agents_map is not None:
                db_agents = [
                    agents_map[UUID(mid) if isinstance(mid, str) else mid]
                    for mid in member_ids
                    if (UUID(mid) if isinstance(mid, str) else mid) in agents_map
                ]
            else:
                # Fallback to query for backward compatibility
                db_agents = self.db.query(Agent).filter(
                    Agent.id.in_([UUID(mid) if isinstance(mid, str) else mid for mid in member_ids]),
                    Agent.organization_id == self.organization_id
                ).all()

            agents = [
                {
                    "id": str(a.id),
                    "name": a.name,
                    "model_id": a.model_id or "default",
                    "status": a.status,
                }
                for a in db_agents
            ]

        return {
            "id": str(team.id),
            "name": team.name,
            "description": team.description,
            "status": team.status.value if hasattr(team.status, 'value') else str(team.status),
            "agent_count": len(agents),
            "agents": agents,
            "runtime": team.runtime,
            "execution_environment": team.execution_environment or {},
        }

    def list_agents(
        self,
        limit: int = 20,
        status: str = None  # Changed from "active" to None - don't filter by default
    ) -> List[Dict[str, Any]]:
        """
        List agents in the organization.

        Args:
            limit: Maximum agents to return
            status: Optional status filter (None = return all agents regardless of status)

        Returns:
            List of agent dictionaries
        """
        try:
            logger.info(
                "list_agents_called",
                organization_id=self.organization_id[:8],
                limit=limit,
                status=status or "all"
            )

            # DEBUG: Log actual query parameters
            logger.info(
                "list_agents_DEBUG_QUERY",
                full_org_id=self.organization_id,
                org_id_len=len(self.organization_id),
                status_filter=status or "no_filter",
                query_info=f"WHERE organization_id='{self.organization_id}'" + (f" AND status='{status}'" if status else " (no status filter)")
            )

            # Direct SQLAlchemy query with org filtering
            query = self.db.query(Agent).filter(
                Agent.organization_id == self.organization_id
            )

            # Only filter by status if explicitly provided
            if status:
                query = query.filter(Agent.status == status)

            agents = query.limit(limit).all()

            # DEBUG: Log what we got back
            logger.info(
                "list_agents_DEBUG_RESULTS",
                agents_found=len(agents),
                agent_names=[a.name for a in agents[:5]] if agents else []
            )

            result = [self._agent_to_dict(a) for a in agents]

            logger.info(
                "list_agents_completed",
                count=len(result),
                organization_id=self.organization_id[:8]
            )

            return result

        except Exception as e:
            logger.error("list_agents_error", error=str(e), exc_info=True)
            return []

    def search_agents_by_capability(
        self,
        capability: str,
        limit: int = 10
    ) -> List[Dict[str, Any]]:
        """
        Search for agents with a specific capability.

        Args:
            capability: Capability to search for
            limit: Maximum agents to return

        Returns:
            List of matching agent dictionaries
        """
        try:
            logger.info(
                "search_agents_by_capability_called",
                capability=capability,
                organization_id=self.organization_id[:8],
                limit=limit
            )

            # Filter in database with LIMIT (avoid loading all agents into memory)
            # Use ILIKE for case-insensitive search and OR conditions
            from sqlalchemy import or_, cast, String

            capability_pattern = f"%{capability}%"
            agents = self.db.query(Agent).filter(
                Agent.organization_id == self.organization_id,
                or_(
                    # Search in JSONB capabilities field (cast to text for ILIKE)
                    cast(Agent.capabilities, String).ilike(capability_pattern),
                    # Search in description
                    Agent.description.ilike(capability_pattern),
                    # Search in name
                    Agent.name.ilike(capability_pattern)
                )
            ).limit(limit * 2).all()  # Buffer for relevance scoring

            # Convert to dict (already filtered in DB)
            result = [self._agent_to_dict(a) for a in agents[:limit]]

            logger.info(
                "search_agents_by_capability_completed",
                capability=capability,
                count=len(result),
                organization_id=self.organization_id[:8]
            )

            return result

        except Exception as e:
            logger.error("search_agents_by_capability_error", error=str(e), exc_info=True)
            # Rollback transaction on error to prevent "invalid transaction" state
            try:
                self.db.rollback()
            except Exception:
                pass  # Ignore rollback errors
            return []

    def list_teams(
        self,
        limit: int = 20,
        status: str = None  # Don't filter by status by default
    ) -> List[Dict[str, Any]]:
        """
        List teams in the organization.

        Args:
            limit: Maximum teams to return
            status: Optional status filter (None = return all teams regardless of status)

        Returns:
            List of team dictionaries with agent details
        """
        try:
            logger.info(
                "list_teams_called",
                organization_id=self.organization_id[:8],
                limit=limit,
                status=status or "all"
            )

            # Direct SQLAlchemy query with org filtering
            query = self.db.query(Team).filter(
                Team.organization_id == self.organization_id
            )

            # Only filter by status if explicitly provided
            if status:
                query = query.filter(Team.status == status)

            teams = query.limit(limit).all()

            result = [self._team_to_dict(t) for t in teams]

            logger.info(
                "list_teams_completed",
                count=len(result),
                organization_id=self.organization_id[:8]
            )

            return result

        except Exception as e:
            logger.error("list_teams_error", error=str(e), exc_info=True)
            # Rollback transaction on error to prevent "invalid transaction" state
            try:
                self.db.rollback()
            except Exception:
                pass  # Ignore rollback errors
            return []

    def search_teams_by_capability(
        self,
        capability: str,
        limit: int = 5
    ) -> List[Dict[str, Any]]:
        """
        Search for teams with agents that have a specific capability.

        Args:
            capability: Capability to search for
            limit: Maximum teams to return

        Returns:
            List of matching team dictionaries
        """
        try:
            logger.info(
                "search_teams_by_capability_called",
                capability=capability,
                organization_id=self.organization_id[:8],
                limit=limit
            )

            # Get teams with LIMIT to avoid loading entire table (use buffer for filtering)
            teams = self.db.query(Team).filter(
                Team.organization_id == self.organization_id
            ).limit(limit * 2).all()

            # Collect all member_ids from all teams upfront
            all_member_ids = set()
            for team in teams:
                config = team.configuration or {}
                member_ids = config.get("member_ids", [])
                if member_ids:
                    all_member_ids.update([UUID(mid) if isinstance(mid, str) else mid for mid in member_ids])

            # Load ALL needed agents in ONE query (eliminates N+1)
            agents_map = {}
            if all_member_ids:
                agents = self.db.query(Agent).filter(
                    Agent.id.in_(list(all_member_ids)),
                    Agent.organization_id == self.organization_id
                ).all()
                agents_map = {agent.id: agent for agent in agents}

            # Filter teams using pre-loaded agents (no additional queries)
            capability_lower = capability.lower()
            matching = []

            for team in teams:
                # Check team description
                if capability_lower in (team.description or "").lower() or capability_lower in team.name.lower():
                    matching.append(team)
                    continue

                # Check team member capabilities using pre-loaded agents
                config = team.configuration or {}
                member_ids = config.get("member_ids", [])
                if member_ids:
                    for mid in member_ids:
                        mid_uuid = UUID(mid) if isinstance(mid, str) else mid
                        agent = agents_map.get(mid_uuid)
                        if agent and capability_lower in str(agent.capabilities).lower():
                            matching.append(team)
                            break

            # Convert to dict using pre-loaded agents (pass agents_map to avoid more queries)
            result = [self._team_to_dict(t, agents_map) for t in matching[:limit]]

            logger.info(
                "search_teams_by_capability_completed",
                capability=capability,
                count=len(result),
                organization_id=self.organization_id[:8]
            )

            return result

        except Exception as e:
            logger.error("search_teams_by_capability_error", error=str(e), exc_info=True)
            # Rollback transaction on error to prevent "invalid transaction" state
            try:
                self.db.rollback()
            except Exception:
                pass  # Ignore rollback errors
            return []

    def get_agent_details(
        self,
        agent_id: str
    ) -> Optional[Dict[str, Any]]:
        """
        Get complete details for a specific agent.

        Args:
            agent_id: Agent ID

        Returns:
            Agent dictionary with full details or None
        """
        try:
            logger.info(
                "get_agent_details_called",
                agent_id=agent_id,
                organization_id=self.organization_id[:8]
            )

            # Get agent with org filtering
            agent = self.db.query(Agent).filter(
                Agent.id == UUID(agent_id) if isinstance(agent_id, str) else agent_id,
                Agent.organization_id == self.organization_id
            ).first()

            if not agent:
                logger.warning(
                    "agent_not_found",
                    agent_id=agent_id,
                    organization_id=self.organization_id[:8]
                )
                return None

            # Import router helpers
            from control_plane_api.app.routers.agents_v2 import (
                get_agent_projects,
                get_agent_environments,
                get_agent_skills_with_inheritance
            )

            # Get related data
            result = {
                **self._agent_to_dict(agent),
                "projects": get_agent_projects(self.db, str(agent.id)),
                "environments": get_agent_environments(self.db, str(agent.id)),
                "skills": get_agent_skills_with_inheritance(
                    self.db,
                    self.organization_id,
                    str(agent.id),
                    str(agent.team_id) if agent.team_id else None
                )
            }

            logger.info(
                "get_agent_details_completed",
                agent_id=agent_id,
                agent_name=agent.name
            )

            return result

        except Exception as e:
            logger.error("get_agent_details_error", error=str(e), exc_info=True)
            return None

    def get_team_details(
        self,
        team_id: str
    ) -> Optional[Dict[str, Any]]:
        """
        Get complete details for a specific team.

        Args:
            team_id: Team ID

        Returns:
            Team dictionary with full details or None
        """
        try:
            logger.info(
                "get_team_details_called",
                team_id=team_id,
                organization_id=self.organization_id[:8]
            )

            # Get team with org filtering
            team = self.db.query(Team).filter(
                Team.id == UUID(team_id) if isinstance(team_id, str) else team_id,
                Team.organization_id == self.organization_id
            ).first()

            if not team:
                logger.warning(
                    "team_not_found",
                    team_id=team_id,
                    organization_id=self.organization_id[:8]
                )
                return None

            result = self._team_to_dict(team)

            logger.info(
                "get_team_details_completed",
                team_id=team_id,
                team_name=team.name
            )

            return result

        except Exception as e:
            logger.error("get_team_details_error", error=str(e), exc_info=True)
            return None

    def list_environments(
        self,
        status: str = "active",
        limit: int = 20
    ) -> List[Dict[str, Any]]:
        """
        List available environments in the organization.

        Args:
            status: Filter by status (default: "active")
            limit: Maximum number of environments to return

        Returns:
            List of environment dictionaries
        """
        try:
            from control_plane_api.app.models.environment import Environment

            logger.info(
                "list_environments_called",
                organization_id=self.organization_id[:8],
                status=status,
                limit=limit
            )

            query = self.db.query(Environment).filter(
                Environment.organization_id == self.organization_id
            )

            if status:
                query = query.filter(Environment.status == status)

            environments = query.order_by(Environment.created_at.desc()).limit(limit).all()

            result = [
                {
                    "id": str(env.id),
                    "name": env.name,
                    "display_name": env.display_name or env.name,
                    "description": env.description,
                    "status": env.status,
                    "tags": env.tags or [],
                    "created_at": env.created_at.isoformat() if env.created_at else None,
                }
                for env in environments
            ]

            logger.info(
                "list_environments_completed",
                count=len(result),
                organization_id=self.organization_id[:8]
            )

            return result

        except Exception as e:
            logger.error("list_environments_error", error=str(e), exc_info=True)
            return []

    def list_worker_queues(
        self,
        environment_id: Optional[str] = None,
        status: str = "active",
        limit: int = 20
    ) -> List[Dict[str, Any]]:
        """
        List available worker queues in the organization.

        CRITICAL: Uses Redis-based real-time worker validation to match execution requirements.
        Only returns queues with active workers currently heartbeating to Redis.

        Args:
            environment_id: Optional filter by environment
            status: Filter by status (default: "active")
            limit: Maximum number of queues to return

        Returns:
            List of worker queue dictionaries with active worker counts (from Redis)
        """
        try:
            from control_plane_api.app.models.worker import WorkerQueue
            from control_plane_api.app.routers.worker_queues import get_active_workers_from_redis
            import asyncio

            logger.info(
                "list_worker_queues_called",
                organization_id=self.organization_id[:8],
                environment_id=environment_id,
                status=status,
                limit=limit
            )

            # Build base query for worker queues
            query = self.db.query(WorkerQueue).filter(
                WorkerQueue.organization_id == self.organization_id,
                WorkerQueue.ephemeral == False  # Exclude ephemeral queues from planning
            )

            if environment_id:
                query = query.filter(WorkerQueue.environment_id == environment_id)

            if status:
                query = query.filter(WorkerQueue.status == status)

            query = query.order_by(WorkerQueue.created_at.desc()).limit(limit)

            candidate_queues = query.all()

            # CRITICAL: Use Redis-based validation (same as plan execution)
            # This ensures consistency between planning and execution
            # Batch all async calls to avoid creating/destroying event loop repeatedly
            async def get_all_workers_batch():
                """Fetch workers for all queues in parallel."""
                tasks = [
                    get_active_workers_from_redis(
                        org_id=self.organization_id,
                        queue_id=str(wq.id),
                        db=self.db
                    )
                    for wq in candidate_queues
                ]
                return await asyncio.gather(*tasks, return_exceptions=True)

            # Run all worker checks in parallel (single event loop)
            all_workers_results = asyncio.run(get_all_workers_batch())

            # Process results
            worker_queues = []
            for wq, active_workers in zip(candidate_queues, all_workers_results):
                try:
                    # Handle exceptions from gather
                    if isinstance(active_workers, Exception):
                        raise active_workers

                    active_count = len(active_workers) if active_workers else 0

                    # Only include queues with active workers
                    if active_count > 0:
                        worker_queues.append({
                            "id": str(wq.id),
                            "name": wq.name,
                            "display_name": wq.display_name or wq.name,
                            "description": wq.description,
                            "environment_id": str(wq.environment_id),
                            "status": wq.status,
                            "active_workers": active_count,
                            "created_at": wq.created_at.isoformat() if wq.created_at else None,
                        })
                        logger.debug(
                            "queue_has_active_workers",
                            queue_id=str(wq.id)[:8],
                            queue_name=wq.name,
                            active_workers=active_count
                        )
                    else:
                        logger.debug(
                            "queue_has_no_active_workers_skipped",
                            queue_id=str(wq.id)[:8],
                            queue_name=wq.name
                        )

                except Exception as e:
                    logger.warning(
                        "failed_to_check_queue_workers",
                        queue_id=str(wq.id)[:8],
                        error=str(e)
                    )
                    continue

            logger.info(
                "list_worker_queues_completed",
                count=len(worker_queues),
                total_queues_checked=len(candidate_queues),
                organization_id=self.organization_id[:8]
            )

            return worker_queues

        except Exception as e:
            logger.error("list_worker_queues_error", error=str(e), exc_info=True)
            # Rollback transaction on error to prevent "invalid transaction" state
            try:
                self.db.rollback()
            except Exception:
                pass  # Ignore rollback errors
            return []

    async def search_context_graph(
        self,
        query: str,
        label: Optional[str] = None,
        limit: int = 20
    ) -> Dict[str, Any]:
        """
        Search the context graph for resources.

        This uses ContextGraphClient which makes HTTP calls to external
        graph.kubiya.ai service - this is correct since it's external.

        Args:
            query: Search query text
            label: Optional node label filter
            limit: Maximum results

        Returns:
            Search results dictionary with nodes
        """
        try:
            logger.info(
                "search_context_graph_called",
                query=query[:100],
                label=label,
                organization_id=self.organization_id[:8],
                limit=limit
            )

            # Use ContextGraphClient (external service - HTTP is correct)
            result = await self.context_graph_client.search_by_text(
                search_text=query,
                label=label,
                limit=limit
            )

            node_count = len(result.get("nodes", [])) if isinstance(result, dict) else 0

            logger.info(
                "search_context_graph_completed",
                query=query[:100],
                node_count=node_count,
                organization_id=self.organization_id[:8]
            )

            return result if isinstance(result, dict) else {"nodes": [], "count": 0}

        except Exception as e:
            logger.error("search_context_graph_error", error=str(e), exc_info=True)
            return {"nodes": [], "count": 0}
