# integration_hub.py
"""
Integration Hub - Connects AIVA components and external services.
"""

import os
import asyncio
import logging
from datetime import datetime
from typing import Optional, Dict, Any, List, Callable

import json  # Import json module

# --- Optional dependencies ---
try:
    import httpx
except ImportError:
    httpx = None

try:
    import aioredis
except ImportError:
    aioredis = None

try:
    import asyncpg
except ImportError:
    asyncpg = None

try:
    from qdrant_client import QdrantClient, models
except ImportError:
    QdrantClient = None
    models = None

# --- Project-specific dependencies ---
from pathlib import Path
ENV_PATH = Path(__file__).parent.parent / "AIVA" / "aiva_config.env"
if ENV_PATH.exists():
    try:
        from dotenv import load_dotenv
        load_dotenv(ENV_PATH)
    except ImportError:
        # Manual env loading fallback
        with open(ENV_PATH) as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith('#') and '=' in line:
                    key, value = line.split('=', 1)
                    os.environ.setdefault(key.strip(), value.strip())


# --- Logging setup ---
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("IntegrationHub")


class IntegrationHub:
    """
    Central hub for integrating various AIVA components and external services.
    """

    # --- Configuration ---
    DEFAULT_OLLAMA_URL = "http://152.53.201.152:23405/api/generate" # AIVA's ollama endpoint
    DEFAULT_REDIS_URL = "redis://:e2ZyYYr4oWRdASI2CaLc-@redis-genesis-u50607.vm.elestio.app:26379"
    DEFAULT_POSTGRES_URL = "postgresql://postgres:your_password@localhost:5432/aiva_db"  # Replace with your actual URL
    DEFAULT_QDRANT_URL = "http://localhost:6333"  # Replace with your Qdrant URL
    DEFAULT_N8N_URL = "http://localhost:5678/webhook"  # Replace with your n8n webhook URL

    def __init__(self):
        # Core services
        self.ollama_url = os.getenv("OLLAMA_URL", self.DEFAULT_OLLAMA_URL)
        self.redis_url = os.getenv("REDIS_URL", self.DEFAULT_REDIS_URL)
        self.postgres_url = os.getenv("POSTGRES_URL", self.DEFAULT_POSTGRES_URL)
        self.qdrant_url = os.getenv("QDRANT_URL", self.DEFAULT_QDRANT_URL)
        self.n8n_url = os.getenv("N8N_URL", self.DEFAULT_N8N_URL)

        # Ollama model
        self.model = os.getenv("AIVA_MODEL", "qwen-long")
        self.timeout = float(os.getenv("AIVA_TIMEOUT", "300"))  # 5 min default

        # Connection pools
        self._redis_pool = None
        self._postgres_pool = None
        self._http_client = None
        self._qdrant_client = None

        # Health status
        self.healthy = False

        # Rate limiting (example)
        self.reasoning_rate_limit = 5  # Requests per minute
        self._reasoning_last_called = None
        self._reasoning_lock = asyncio.Lock()

        # Initialize connections
        asyncio.create_task(self.initialize())

        logger.info("IntegrationHub initialized")

    async def initialize(self):
        """Initialize connections to all services."""
        try:
            await self._connect_redis()
            await self._connect_postgres()
            await self._connect_qdrant()
            self.healthy = True
            logger.info("IntegrationHub successfully connected to all services.")
        except Exception as e:
            logger.error(f"IntegrationHub initialization failed: {e}")
            self.healthy = False

    # --- Connection Management ---

    async def _get_http_client(self) -> "httpx.AsyncClient":
        """Get or create HTTP client."""
        if httpx is None:
            raise ImportError("httpx required: pip install httpx")
        if self._http_client is None or self._http_client.is_closed:
            self._http_client = httpx.AsyncClient(timeout=self.timeout)
        return self._http_client

    async def _connect_redis(self):
        """Establish Redis connection pool."""
        if aioredis is None:
            raise ImportError("aioredis required: pip install redis")
        try:
            self._redis_pool = aioredis.from_url(self.redis_url)
            await self._redis_pool.ping()  # Verify connection
            logger.info("Connected to Redis")
        except Exception as e:
            logger.error(f"Failed to connect to Redis: {e}")
            self._redis_pool = None

    async def _connect_postgres(self):
        """Establish PostgreSQL connection pool."""
        if asyncpg is None:
            raise ImportError("asyncpg required: pip install asyncpg")
        try:
            self._postgres_pool = await asyncpg.create_pool(self.postgres_url)
            logger.info("Connected to PostgreSQL")
        except Exception as e:
            logger.error(f"Failed to connect to PostgreSQL: {e}")
            self._postgres_pool = None

    async def _connect_qdrant(self):
        """Establish Qdrant client."""
        if QdrantClient is None:
            raise ImportError("qdrant_client required: pip install qdrant-client")
        try:
            self._qdrant_client = QdrantClient(url=self.qdrant_url)
            logger.info("Connected to Qdrant")
        except Exception as e:
            logger.error(f"Failed to connect to Qdrant: {e}")
            self._qdrant_client = None

    # --- Health Monitoring ---

    async def health_check(self) -> Dict[str, Any]:
        """Check health status of all components."""
        health = {
            "status": "unhealthy",
            "timestamp": datetime.now().isoformat(),
            "components": {}
        }

        # Check Ollama
        try:
            ollama_health = await self._check_ollama_health()
            health["components"]["ollama"] = ollama_health
        except Exception as e:
            health["components"]["ollama"] = {"status": "error", "error": str(e)}

        # Check Redis
        health["components"]["redis"] = {
            "status": "online" if self._redis_pool else "offline"
        }
        # Check Postgres
        health["components"]["postgres"] = {
            "status": "online" if self._postgres_pool else "offline"
        }
        # Check Qdrant
        health["components"]["qdrant"] = {
            "status": "online" if self._qdrant_client else "offline"
        }

        if all(component["status"] == "online" for component in health["components"].values() if isinstance(component, dict)):
            health["status"] = "healthy"
            self.healthy = True
        else:
            self.healthy = False

        return health

    async def _check_ollama_health(self) -> Dict[str, Any]:
        """Check AIVA/Ollama health status."""
        tags_url = self.ollama_url.replace("/api/generate", "/api/tags")
        try:
            client = await self._get_http_client()
            response = await client.get(tags_url, timeout=10.0)
            response.raise_for_status()
            data = response.json()

            models_available = [m.get("name") for m in data.get("models", [])]
            qwen_ready = self.model in models_available or f"{self.model}:latest" in models_available

            return {
                "status": "online" if qwen_ready else "model_missing",
                "models": models_available,
                "target_model": self.model,
                "model_ready": qwen_ready,
                "endpoint": self.ollama_url.split("/api")[0],
                "timestamp": datetime.now().isoformat()
            }
        except Exception as e:
            return {
                "status": "offline",
                "error": str(e),
                "endpoint": self.ollama_url.split("/api")[0],
                "timestamp": datetime.now().isoformat()
            }

    # --- Rate Limiting ---

    async def _check_reasoning_rate_limit(self) -> bool:
        """Check if reasoning rate limit has been exceeded."""
        async with self._reasoning_lock:
            now = datetime.now()
            if self._reasoning_last_called is not None:
                time_since_last_call = (now - self._reasoning_last_called).total_seconds()
                if time_since_last_call < (60 / self.reasoning_rate_limit):
                    return False  # Rate limit exceeded
            self._reasoning_last_called = now
            return True  # Rate limit not exceeded

    # --- Core Functionality ---

    async def reason(
        self,
        prompt: str,
        context: Optional[Dict[str, Any]] = None,
        system_prompt: Optional[str] = None,
        temperature: float = 0.7,
        max_tokens: int = 4096,
        strip_thinking: bool = False
    ) -> Dict[str, Any]:
        """
        Send reasoning request to AIVA via Ollama.

        Args:
            prompt: The question or task for AIVA
            context: Optional context dict to include
            system_prompt: Override default system prompt
            temperature: Sampling temperature (0.0-1.0)
            max_tokens: Maximum tokens to generate
            strip_thinking: Remove <think> blocks from response

        Returns:
            Dict with 'response', 'tokens', 'duration', 'success'
        """
        if not await self._check_reasoning_rate_limit():
            return {
                "success": False,
                "response": None,
                "error": "Reasoning rate limit exceeded.",
                "tokens": 0,
                "duration": 0
            }

        try:
            # Build full prompt with context
            full_prompt = prompt
            if context:
                context_str = json.dumps(context, indent=2)
                full_prompt = f"Context:\n