"""Integration module for connecting PraisonAIUI with praisonai backend.

This module provides utilities to use PraisonAIUI's React frontend
with the praisonai WebSocketGateway backend instead of the standalone server.
"""

from __future__ import annotations

import logging
import os
from pathlib import Path
from typing import TYPE_CHECKING, Any, Dict, List, Optional

if TYPE_CHECKING:
    from praisonaiagents import Agent

logger = logging.getLogger(__name__)

# Check if praisonai is available
try:
    from praisonai.gateway import WebSocketGateway
    from praisonaiagents import Agent as PraisonAgent
    HAS_PRAISONAI = True
except ImportError:
    HAS_PRAISONAI = False
    WebSocketGateway = None
    PraisonAgent = None


def check_praisonai_available() -> bool:
    """Check if praisonai package is available."""
    return HAS_PRAISONAI


class AIUIGateway:
    """Extended WebSocketGateway with static file serving for PraisonAIUI.

    This class wraps the praisonai WebSocketGateway and adds:
    - Static file serving for the React SPA
    - REST API endpoints for UI configuration
    - Integration with PraisonAIUI's YAML compiler

    Example:
        from praisonaiui.integration import AIUIGateway
        from praisonaiagents import Agent

        gateway = AIUIGateway(port=8080, static_dir="./dist")
        agent = Agent(name="assistant", instructions="You are helpful.")
        gateway.register_agent(agent)

        await gateway.start()
    """

    def __init__(
        self,
        host: str = "127.0.0.1",
        port: int = 8080,
        static_dir: Optional[str] = None,
        ui_config: Optional[Dict[str, Any]] = None,
    ):
        """Initialize the AIUI Gateway.

        Args:
            host: Host to bind to
            port: Port to listen on
            static_dir: Directory containing the React SPA build
            ui_config: UI configuration dict (from compiler output)
        """
        if not HAS_PRAISONAI:
            raise ImportError(
                "praisonai package is required for integration mode. "
                "Install with: pip install praisonai"
            )

        self._host = host
        self._port = port
        self._static_dir = Path(static_dir) if static_dir else None
        self._ui_config = ui_config or {}

        # Internal gateway instance
        self._gateway: Optional[WebSocketGateway] = None
        self._is_running = False
        self._server = None

    def register_agent(self, agent: "Agent", agent_id: Optional[str] = None) -> str:
        """Register an agent with the gateway."""
        if self._gateway is None:
            self._init_gateway()
        return self._gateway.register_agent(agent, agent_id)

    def unregister_agent(self, agent_id: str) -> bool:
        """Unregister an agent from the gateway."""
        if self._gateway:
            return self._gateway.unregister_agent(agent_id)
        return False

    def list_agents(self) -> List[str]:
        """List all registered agent IDs."""
        if self._gateway:
            return self._gateway.list_agents()
        return []

    def set_ui_config(self, config: Dict[str, Any]) -> None:
        """Set the UI configuration."""
        self._ui_config = config

    def _init_gateway(self) -> None:
        """Initialize the internal WebSocketGateway."""
        from praisonaiagents.gateway import GatewayConfig

        config = GatewayConfig(host=self._host, port=self._port)
        self._gateway = WebSocketGateway(
            host=self._host,
            port=self._port,
            config=config,
        )

    async def start(self) -> None:
        """Start the gateway server with static file serving.

        Delegates to ``server.create_app()`` which provides all dashboard
        API endpoints, CORS, auth, sessions CRUD, feature protocols, and
        the pluggable provider system.  The gateway adds a ``/ws`` WebSocket
        route on top and registers its agents with the server's provider.
        """
        if self._is_running:
            logger.warning("Gateway already running")
            return

        if self._gateway is None:
            self._init_gateway()

        try:
            import uvicorn
            from starlette.routing import WebSocketRoute
            from starlette.websockets import WebSocket, WebSocketDisconnect
        except ImportError:
            raise ImportError(
                "Gateway requires starlette and uvicorn. "
                "Install with: pip install starlette uvicorn"
            )

        # ── Register gateway agents with the server's provider system ──
        from praisonaiui.server import create_app, register_agent, set_provider

        # Expose gateway agents to the server registry
        gateway_agents = []
        for aid in self._gateway.list_agents():
            agent = self._gateway.get_agent(aid)
            if agent:
                register_agent(aid, agent)
                gateway_agents.append(agent)

        # Set up the PraisonAI provider with gateway agents
        from praisonaiui.providers import PraisonAIProvider
        set_provider(PraisonAIProvider(agents=gateway_agents))

        # ── Build the app via server.create_app() ──────────────────────
        static_path = self._static_dir if self._static_dir and self._static_dir.exists() else None
        app = create_app(
            config=self._ui_config or None,
            static_dir=static_path,
        )

        # ── Add WebSocket route for the gateway ────────────────────────
        import uuid

        async def ws_handler(websocket: WebSocket):
            await websocket.accept()
            client_id = str(uuid.uuid4())
            self._gateway._clients[client_id] = websocket

            logger.info(f"Client connected: {client_id}")

            from praisonaiagents.gateway import EventType, GatewayEvent
            await self._gateway.emit(GatewayEvent(
                type=EventType.CONNECT,
                data={"client_id": client_id},
                source=client_id,
            ))

            try:
                while True:
                    data = await websocket.receive_json()
                    await self._gateway._handle_client_message(client_id, data)
            except WebSocketDisconnect:
                logger.info(f"Client disconnected: {client_id}")
            except Exception as e:
                logger.error(f"WebSocket error: {e}")
            finally:
                self._gateway._clients.pop(client_id, None)
                session_id = self._gateway._client_sessions.pop(client_id, None)
                if session_id:
                    self._gateway.close_session(session_id)

                await self._gateway.emit(GatewayEvent(
                    type=EventType.DISCONNECT,
                    data={"client_id": client_id},
                    source=client_id,
                ))

        # Prepend the /ws route so it takes priority over static catch-all
        app.routes.insert(0, WebSocketRoute("/ws", ws_handler))

        config = uvicorn.Config(
            app,
            host=self._host,
            port=self._port,
            log_level="info",
        )
        self._server = uvicorn.Server(config)

        self._is_running = True

        # Publish gateway reference so feature modules can access live state
        try:
            from praisonaiui.features._gateway_ref import set_gateway
            set_gateway(self._gateway)
        except ImportError:
            pass

        # Initialize unified YAML config store for feature persistence
        try:
            from praisonaiui.config_store import init_config_store
            default_data_dir = Path.home() / ".praisonaiui"
            init_config_store(default_data_dir / "config.yaml")
        except Exception:
            pass

        logger.info(f"AIUI Gateway started on http://{self._host}:{self._port}")

        await self._server.serve()

    async def stop(self) -> None:
        """Stop the gateway server."""
        if not self._is_running:
            return

        self._is_running = False

        if self._gateway:
            await self._gateway.stop()

        if self._server:
            self._server.should_exit = True

        logger.info("AIUI Gateway stopped")

        # Clear gateway reference
        try:
            from praisonaiui.features._gateway_ref import set_gateway
            set_gateway(None)
        except ImportError:
            pass


def create_gateway_from_yaml(
    config_path: str,
    static_dir: Optional[str] = None,
) -> AIUIGateway:
    """Create an AIUIGateway from a YAML configuration file.

    Args:
        config_path: Path to aiui.template.yaml or gateway.yaml
        static_dir: Directory containing the React SPA build

    Returns:
        Configured AIUIGateway instance
    """
    import yaml

    if not os.path.exists(config_path):
        raise FileNotFoundError(f"Config not found: {config_path}")

    with open(config_path) as f:
        config = yaml.safe_load(f)

    # Extract gateway settings
    gateway_cfg = config.get("gateway", {})
    host = gateway_cfg.get("host", "127.0.0.1")
    port = gateway_cfg.get("port", 8080)

    # Create gateway
    gateway = AIUIGateway(
        host=host,
        port=port,
        static_dir=static_dir,
        ui_config=config,
    )

    # G5: Resolve tool names to callables (same pattern as gateway G1 fix)
    tool_resolver = None
    try:
        from praisonai.tool_resolver import ToolResolver
        tool_resolver = ToolResolver()
    except ImportError:
        logger.debug("ToolResolver not available, agents will have no tools")

    # Create agents from config
    agents_cfg = config.get("agents", {})
    for agent_id, agent_def in agents_cfg.items():
        if not HAS_PRAISONAI:
            logger.warning("praisonai not available, skipping agent creation")
            break

        # G5: Resolve tools from config
        agent_tools = []
        for tool_name in agent_def.get("tools", []):
            if tool_resolver and isinstance(tool_name, str) and tool_name.strip():
                resolved = tool_resolver.resolve(tool_name.strip())
                if resolved:
                    agent_tools.append(resolved)
                else:
                    logger.warning(f"Tool '{tool_name}' not found for agent '{agent_id}'")

        agent = PraisonAgent(
            name=agent_id,
            instructions=agent_def.get("instructions", ""),
            llm=agent_def.get("model"),
            memory=agent_def.get("memory", False),
            tools=agent_tools if agent_tools else None,
            self_reflect=agent_def.get("reflection", False),
        )
        gateway.register_agent(agent, agent_id=agent_id)
        logger.info(f"Created agent '{agent_id}' (tools={len(agent_tools)}, reflection={agent_def.get('reflection', False)})")

    return gateway


async def run_with_praisonai(
    app_module: str,
    static_dir: str,
    host: str = "127.0.0.1",
    port: int = 8080,
) -> None:
    """Run PraisonAIUI with praisonai backend.

    This function:
    1. Imports the user's app module to get registered agents/callbacks
    2. Creates an AIUIGateway with static file serving
    3. Registers agents and starts the server

    Args:
        app_module: Path to user's app.py
        static_dir: Directory containing the React SPA build
        host: Host to bind to
        port: Port to listen on
    """
    import importlib.util

    # Load user's app module
    spec = importlib.util.spec_from_file_location("user_app", app_module)
    if spec and spec.loader:
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)

    # Create gateway
    gateway = AIUIGateway(
        host=host,
        port=port,
        static_dir=static_dir,
    )

    # Check if user defined agents
    if hasattr(module, "agent"):
        gateway.register_agent(module.agent)
    elif hasattr(module, "agents"):
        for agent in module.agents:
            gateway.register_agent(agent)

    # Start gateway
    await gateway.start()
