"""
Unified Multi-Agent Orchestrator.

Consolidates Ag2Orchestrator and EnhancedAg2Orchestrator into a single
configurable orchestrator following DRY, SRP, and OCP principles.

Uses autogen-agentchat 0.7+ async API.

Modes:
- BASIC: Single engineer agent (fast, simple tasks)
- ENHANCED: Multi-agent team with selector (complex tasks)
"""
import os
from enum import Enum
from typing import Any, Callable, List, Optional

from rich.console import Console

from merlya.agents import autogen_tools, knowledge_tools
from merlya.agents.base_orchestrator import BaseOrchestrator
from merlya.agents.orchestrator_service.intent import IntentParser
from merlya.agents.orchestrator_service.planner import ExecutionPlanner
from merlya.llm.model_config import ModelConfig
from merlya.utils.config import ConfigManager
from merlya.utils.logger import logger

# Optional imports
try:
    from merlya.utils.verbosity import get_verbosity
    HAS_VERBOSITY = True
except ImportError:
    HAS_VERBOSITY = False

try:
    from merlya.knowledge.falkordb_client import FalkorDBClient, FalkorDBConfig
    HAS_FALKORDB = True
except ImportError:
    HAS_FALKORDB = False

# New autogen-agentchat 0.7+ imports
try:
    from autogen_ext.models.openai import OpenAIChatCompletionClient
    HAS_AUTOGEN = True
except ImportError:
    HAS_AUTOGEN = False
    logger.warning("⚠️ autogen-agentchat not installed. Orchestrator will not work.")


class OrchestratorMode(Enum):
    """Orchestrator operating modes."""
    BASIC = "basic"           # Single engineer agent
    ENHANCED = "enhanced"     # Multi-agent team


class Orchestrator(BaseOrchestrator):
    """
    Unified Multi-Agent Orchestrator.

    Uses autogen-agentchat 0.7+ async API.

    Combines the functionality of Ag2Orchestrator and EnhancedAg2Orchestrator
    into a single, configurable class.

    Args:
        mode: Operating mode (BASIC or ENHANCED)
        env: Environment name (dev, staging, prod)
        language: Response language (en, fr)

    Example:
        # Simple mode for quick tasks
        orchestrator = Orchestrator(mode=OrchestratorMode.BASIC)

        # Enhanced mode with multi-agent team
        orchestrator = Orchestrator(mode=OrchestratorMode.ENHANCED)
    """

    def __init__(
        self,
        mode: OrchestratorMode = OrchestratorMode.BASIC,
        env: str = "dev",
        language: str = "en",
        console: Optional[Console] = None,
    ):
        super().__init__(env=env, language=language)

        if not HAS_AUTOGEN:
            raise ImportError(
                "autogen-agentchat is required. "
                "Run 'pip install pyautogen autogen-ext[openai]'."
            )

        self.mode = mode
        self.console = console or Console()
        self.config_manager = ConfigManager()

        # Verbosity (optional)
        self.verbosity = get_verbosity() if HAS_VERBOSITY else None

        # Intent Parser (uses local embeddings first, LLM router as fallback)
        self.intent_parser = IntentParser(
            self.console,
            self.verbosity,
            llm_router=self.llm_router
        )
        self.current_priority: Any = None

        # Knowledge graph (ENHANCED mode only)
        self.knowledge_db: FalkorDBClient | None = None
        if mode == OrchestratorMode.ENHANCED and HAS_FALKORDB:
            self._init_knowledge_db()

        # Initialize tools
        autogen_tools.initialize_autogen_tools(
            executor=self.executor,
            context_manager=self.context_manager,
            permissions=self.permissions,
            context_memory=getattr(self.context_manager, 'memory', None),
            credentials=self.credentials,
        )

        # Configure model client (legacy support)
        # We don't initialize a single client anymore, but keep the attribute for backward compatibility if needed
        self.model_client = None

        # Collect tools as callables
        self._tools = self._collect_tools()

        # Client cache for task-specific models (must be before planner creation)
        self._client_cache = {}

        # Execution Planner
        # Pass client factory instead of single client
        self.planner = ExecutionPlanner(
            client_factory=self._get_client_for_task,
            tools=self._tools,
            env=self.env,
            console=self.console
        )

        # Initialize agents
        self.planner.init_agents(mode.value, self.knowledge_db)

        logger.info(f"Orchestrator initialized in {mode.value} mode")

    # =========================================================================
    # Model Client Configuration (new API)
    # =========================================================================

    def _get_client_for_task(self, task: str = "synthesis") -> "OpenAIChatCompletionClient":
        """
        Get or create a model client for a specific task.

        Args:
            task: Task type (correction, planning, synthesis)

        Returns:
            Cached or new OpenAIChatCompletionClient
        """
        if task in self._client_cache:
            return self._client_cache[task]

        client = self._create_model_client(task)
        self._client_cache[task] = client
        return client

    def _create_model_client(self, task: Optional[str] = None) -> "OpenAIChatCompletionClient":
        """Create OpenAI-compatible model client for specific task.

        Priority: Environment variables > Config file > Defaults
        Provider is determined solely by ModelConfig (unified config).
        """
        # Load model config from ~/.merlya/config.json
        model_config = ModelConfig()
        config_provider = model_config.get_provider()

        # Determine provider: env var overrides config
        provider = os.getenv("MERLYA_PROVIDER", "").lower() or config_provider

        # ✅ FIX: Get model for THE ACTUAL PROVIDER, not config_provider
        # This ensures we use the correct model for the selected provider

        # Get task-specific model ID
        # Note: We pass the task to get_model to resolve aliases (haiku/sonnet/opus)
        # to the correct model ID for the provider

        # Ollama (local LLM)
        if provider == "ollama" or os.getenv("OLLAMA_MODEL"):
            model = os.getenv("OLLAMA_MODEL") or model_config.get_model("ollama", task=task)
            logger.info(f"✅ Using Ollama for {task or 'default'}: {model}")
            return OpenAIChatCompletionClient(
                model=model,
                api_key="ollama",
                base_url="http://localhost:11434/v1",
                model_info={
                    "vision": False,
                    "function_calling": True,
                    "json_output": True,
                    "family": "unknown",
                    "structured_output": True,
                },
            )

        if provider == "openrouter" or os.getenv("OPENROUTER_API_KEY"):
            api_key = os.getenv("OPENROUTER_API_KEY")
            if not api_key:
                raise ValueError("OPENROUTER_API_KEY not set.")
            # ✅ FIX: Get model for "openrouter", not config_provider
            model = os.getenv("OPENROUTER_MODEL") or model_config.get_model("openrouter", task=task)
            logger.info(f"✅ Using OpenRouter for {task or 'default'}: {model}")
            return OpenAIChatCompletionClient(
                model=model,
                api_key=api_key,
                base_url="https://openrouter.ai/api/v1",
                model_info={
                    "vision": False,
                    "function_calling": True,
                    "json_output": True,
                    "family": "unknown",
                    "structured_output": True,
                },
            )

        if provider == "anthropic" or os.getenv("ANTHROPIC_API_KEY"):
            api_key = os.getenv("ANTHROPIC_API_KEY")
            if not api_key:
                raise ValueError("ANTHROPIC_API_KEY not set.")
            # ✅ FIX: Get model for "anthropic", not config_provider
            model = os.getenv("ANTHROPIC_MODEL") or model_config.get_model("anthropic", task=task)
            logger.info(f"✅ Using Anthropic for {task or 'default'}: {model}")
            return OpenAIChatCompletionClient(
                model=model,
                api_key=api_key,
                base_url="https://api.anthropic.com/v1",
                model_info={
                    "vision": True,
                    "function_calling": True,
                    "json_output": True,
                    "family": "unknown",
                    "structured_output": True,
                },
            )

        if provider == "openai" or os.getenv("OPENAI_API_KEY"):
            api_key = os.getenv("OPENAI_API_KEY")
            if not api_key:
                raise ValueError("OPENAI_API_KEY not set.")
            # ✅ FIX: Get model for "openai", not config_provider
            model = os.getenv("OPENAI_MODEL") or model_config.get_model("openai", task=task)
            logger.info(f"✅ Using OpenAI for {task or 'default'}: {model}")
            return OpenAIChatCompletionClient(
                model=model,
                api_key=api_key,
            )

        raise ValueError(
            "No AI provider configured. Set MERLYA_PROVIDER and the corresponding API key, "
            "or use Ollama (MERLYA_PROVIDER=ollama with OLLAMA_MODEL)."
        )

    def reload_agents(self) -> None:
        """Reload agents with current configuration."""

        logger.info("Reloading agents...")

        # Close old model clients
        self.shutdown_sync()

        # Clear cache
        self._client_cache = {}

        # Re-initialize planner with new factory
        self.planner = ExecutionPlanner(
            client_factory=self._get_client_for_task,
            tools=self._tools,
            env=self.env,
            console=self.console
        )
        self.planner.init_agents(self.mode.value, self.knowledge_db)
        logger.info("Agents reloaded.")

    # =========================================================================
    # Tool Collection (new API - tools as callables)
    # =========================================================================

    def _collect_tools(self) -> List[Callable[..., Any]]:
        """Collect all tools as callable functions."""
        tools: List[Callable[..., Any]] = [
            # Core tools
            autogen_tools.scan_host,
            autogen_tools.execute_command,
            autogen_tools.check_permissions,
            autogen_tools.get_infrastructure_context,
            autogen_tools.list_hosts,
            autogen_tools.audit_host,
            autogen_tools.add_route,
            autogen_tools.analyze_security_logs,
            # File operations
            autogen_tools.read_remote_file,
            autogen_tools.write_remote_file,
            autogen_tools.tail_logs,
            autogen_tools.glob_files,
            autogen_tools.grep_files,
            autogen_tools.find_file,
            # Web tools
            autogen_tools.web_search,
            autogen_tools.web_fetch,
            # Interaction & Learning
            autogen_tools.ask_user,
            autogen_tools.get_user_variables,
            autogen_tools.get_variable_value,
            autogen_tools.remember_skill,
            autogen_tools.recall_skill,
            autogen_tools.request_elevation,
            autogen_tools.save_report,
            # Knowledge Graph tools
            knowledge_tools.record_incident,
            knowledge_tools.search_knowledge,
            knowledge_tools.get_solution_suggestion,
            knowledge_tools.graph_stats,
            # System info
            autogen_tools.disk_info,
            autogen_tools.memory_info,
            autogen_tools.network_connections,
            autogen_tools.process_list,
            autogen_tools.service_control,
            # Containers
            autogen_tools.docker_exec,
            autogen_tools.kubectl_exec,
        ]
        return tools

    # =========================================================================
    # Knowledge Graph (ENHANCED mode)
    # =========================================================================

    def _init_knowledge_db(self) -> None:
        """Initialize FalkorDB connection."""
        try:
            config = FalkorDBConfig(
                graph_name="merlya_knowledge",
                auto_start_docker=True,
            )
            self.knowledge_db = FalkorDBClient(config)
            if self.knowledge_db.connect():
                logger.info("FalkorDB connected")
                self._ensure_knowledge_schema()
            else:
                logger.warning("FalkorDB not available")
                self.knowledge_db = None
        except Exception as e:
            logger.warning(f"FalkorDB init failed: {e}")
            self.knowledge_db = None

    def _ensure_knowledge_schema(self) -> None:
        """Ensure knowledge graph indexes exist."""
        if not self.knowledge_db:
            return
        try:
            self.knowledge_db.query("CREATE INDEX IF NOT EXISTS FOR (h:Host) ON (h.hostname)")
            self.knowledge_db.query("CREATE INDEX IF NOT EXISTS FOR (i:Incident) ON (i.created_at)")
        except Exception as e:
            logger.warning(f"Failed to create schema: {e}")

    # =========================================================================
    # Session Management
    # =========================================================================

    def reset_session(self) -> None:
        """Reset the chat session."""
        # New API doesn't have agent.reset(), just reinitialize
        self.planner.init_agents(self.mode.value, self.knowledge_db)
        self.console.print("[dim]🔄 Session reset[/dim]")

    async def shutdown(self) -> None:
        """Clean shutdown of orchestrator resources.

        Closes all model clients to prevent 'Event loop is closed' errors
        from httpx connections during garbage collection.
        """
        for task, client in self._client_cache.items():
            try:
                await client.close()
                logger.debug(f"Model client for {task} closed successfully")
            except Exception as e:
                logger.debug(f"Error closing model client for {task}: {e}")

        self._client_cache = {}

    def shutdown_sync(self) -> None:
        """Synchronous wrapper for shutdown.

        Use this when not in an async context (e.g., REPL exit).
        Best effort cleanup - may not complete if event loop is in problematic state.
        """
        import asyncio

        if not self._client_cache:
            return

        try:
            # Try to get the current event loop
            loop = asyncio.get_event_loop()
            if loop.is_running():
                # Cannot synchronously wait while loop is running
                # Log and skip - the running code should handle cleanup
                logger.debug("Cannot shutdown_sync while event loop is running")
                return
            elif not loop.is_closed():
                # If loop exists but not running, use it
                loop.run_until_complete(self.shutdown())
                logger.debug("Model client closed successfully")
                return
        except RuntimeError:
            pass  # Fall through to sync fallback

        # Fallback: try to create a new loop or close synchronously
        try:
            asyncio.run(self.shutdown())
            logger.debug("Model client closed (new loop)")
        except RuntimeError:
            # Best-effort synchronous cleanup
            # Note: Relies on internal httpx structure - may not work on all clients
            self._close_client_sync()

    def _close_client_sync(self) -> None:
        """Best-effort synchronous client cleanup."""
        if not self._client_cache:
            return

        for task, client in self._client_cache.items():
            try:
                # Try common client structures
                inner_client = getattr(client, '_client', None)
                if inner_client is not None and hasattr(inner_client, 'close'):
                    inner_client.close()
                    logger.debug(f"Model client for {task} closed (sync fallback)")
            except Exception as e:
                logger.debug(f"Sync shutdown fallback failed for {task}: {e}")

    # =========================================================================
    # Error Handling
    # =========================================================================

    def _handle_execution_error(self, error: Exception) -> str:
        """
        Classify and format execution errors with actionable messages.

        Args:
            error: The exception that occurred

        Returns:
            User-friendly error message
        """
        error_str = str(error)

        # Check for function calling / tool use not supported
        if "No endpoints found that support tool use" in error_str or "404" in error_str:
            logger.error("❌ Model doesn't support function calling: {}", error, exc_info=True)
            return self._build_function_calling_error_message()

        # Generic error fallback
        logger.error("❌ Orchestrator failed: {}", error, exc_info=True)
        return f"❌ Error: {error_str}"

    def _build_function_calling_error_message(self) -> str:
        """
        Build actionable error message when model doesn't support function calling.

        Returns:
            User-friendly error message with specific solutions
        """
        model_config = ModelConfig()
        current_provider = model_config.get_provider()
        current_model = model_config.get_model(current_provider)

        message_parts = [
            "❌ Model Error: The current model doesn't support function calling (required for multi-agent mode)",
            "",
            "Current configuration:",
            f"  • Provider: {current_provider}",
            f"  • Model: {current_model}",
            "",
            "💡 Solutions:",
        ]

        # Suggest provider-specific solutions
        if current_provider == "openrouter":
            message_parts.extend([
                "  1. Switch to a compatible OpenRouter model:",
                "     /model set openrouter google/gemini-2.0-flash-exp:free",
                "     /model set openrouter anthropic/claude-3.5-sonnet",
                "     /model set openrouter qwen/qwen-2.5-72b-instruct",
                "",
                "  2. Switch to Ollama (if configured):",
                "     /model provider ollama",
            ])
        elif current_provider == "ollama":
            message_parts.extend([
                "  1. Verify Ollama model supports function calling",
                "  2. Try a different Ollama model:",
                "     /model set ollama qwen2.5-coder:latest",
                "",
                "  3. Switch to a cloud provider:",
                "     /model provider openrouter",
            ])
        elif current_provider == "anthropic":
            message_parts.extend([
                "  1. Use Claude without date suffix:",
                "     /model set anthropic claude-3-5-sonnet-20241022",
                "",
                "  2. Switch to OpenRouter for more model options:",
                "     /model provider openrouter",
            ])
        else:
            message_parts.extend([
                "  1. Check available models: /model list",
                "  2. Switch provider: /model provider <provider>",
                "  3. Set compatible model: /model set <provider> <model>",
            ])

        message_parts.extend([
            "",
            "📖 For more information:",
            "   /model help",
        ])

        return "\n".join(message_parts)

    # =========================================================================
    # Main Processing (new async API)
    # =========================================================================

    async def process_request(
        self,
        user_query: str,
        auto_confirm: bool = False,
        dry_run: bool = False,
        **kwargs
    ) -> str:
        """
        Process request using Multi-Agent conversation.

        Args:
            user_query: User's request
            auto_confirm: Auto-confirm critical actions
            dry_run: Preview only, don't execute
            conversation_history: Recent conversation context (list of {role, content})

        Returns:
            Agent response
        """
        if dry_run:
            return f"🔍 Dry run: Would process in {self.mode.value} mode"

        # Step 1: Full classification (priority + intent) - use AI when available
        # Use original_query for triage if provided (avoids exposing resolved credentials)
        system_state = kwargs.get("system_state")
        original_query = kwargs.get("original_query")
        triage_query = original_query if original_query else user_query

        try:
            triage_context = await self.intent_parser.classify_full_async(triage_query, system_state)
        except Exception as e:
            logger.error(f"❌ Classification failed: {e}", exc_info=True)
            # Fallback: process without triage context
            conversation_history = kwargs.get("conversation_history", [])
            try:
                return await self.planner.execute_basic(user_query, conversation_history)
            except Exception as fallback_error:
                return self._handle_execution_error(fallback_error)

        # Store for backward compatibility
        self.current_priority = triage_context.priority_result

        # Step 2: Display triage with intent
        self.intent_parser.display_full_triage(triage_context)

        # Log with defensive attribute access
        priority_name = getattr(
            getattr(triage_context.priority_result, "priority", None), "name", "UNKNOWN"
        )
        intent_value = getattr(triage_context.intent, "value", "unknown") if triage_context.intent else "unknown"
        logger.info(
            f"🎯 Request classified as {priority_name} "
            f"(intent: {intent_value}, mode: {self.mode.value})"
        )

        # Step 3: Get conversation history for context
        conversation_history = kwargs.get("conversation_history", [])

        # Step 4: Execute based on mode (with tool restrictions and intent)
        try:
            allowed_tools = triage_context.allowed_tools

            # Fetch knowledge context if in ENHANCED mode
            knowledge_context = None
            if self.mode == OrchestratorMode.ENHANCED and HAS_FALKORDB:
                try:
                    # Search for relevant knowledge using triage_query to avoid
                    # exposing resolved credentials in search
                    knowledge_context = knowledge_tools.search_knowledge(triage_query, limit=3)
                except Exception as e:
                    logger.warning(f"⚠️ Failed to fetch knowledge context: {e}")

            if self.mode == OrchestratorMode.ENHANCED:
                result = await self.planner.execute_enhanced(
                    user_query,
                    priority_name,  # Already extracted with defensive access above
                    conversation_history,
                    allowed_tools=allowed_tools,
                    intent=intent_value,  # Pass intent for behavior adaptation
                    knowledge_context=knowledge_context,
                )
            else:
                result = await self.planner.execute_basic(
                    user_query,
                    conversation_history,
                    allowed_tools=allowed_tools,
                    intent=intent_value,  # Pass intent for behavior adaptation
                    priority=priority_name,  # Pass priority for behavior adaptation
                )

            # Step 5: Implicit positive feedback - classification was used successfully
            # After ~3 successful uses, pattern becomes trusted (confidence >= 0.7)
            self.intent_parser.confirm_last_classification()

            return result

        except Exception as e:
            return self._handle_execution_error(e)

    async def chat_continue(self, message: str) -> str:
        """Continue conversation (BASIC mode compatibility)."""
        return await self.planner.execute_basic(message)


# =============================================================================
# Factory function for easy instantiation
# =============================================================================

def create_orchestrator(
    mode: str = "basic",
    env: str = "dev",
    language: str = "en",
    console: Optional[Console] = None,
) -> Orchestrator:
    """
    Factory function to create orchestrator.

    Args:
        mode: "basic" or "enhanced"
        env: Environment name
        language: Response language
        console: Optional existing console

    Returns:
        Configured Orchestrator instance
    """
    mode_enum = OrchestratorMode(mode.lower())
    return Orchestrator(mode=mode_enum, env=env, language=language, console=console)
