# integration_hub.py
import asyncio
import json
import logging
import os
from datetime import datetime
from typing import Any, Dict, List, Optional

import httpx
import redis.asyncio as aioredis
from qdrant_client import QdrantClient, models
import asyncpg

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("IntegrationHub")

# Load environment variables (using dotenv if available)
try:
    from dotenv import load_dotenv

    load_dotenv()
except ImportError:
    pass  # dotenv not installed, rely on system environment variables

# --- Configuration ---
OLLAMA_URL = os.getenv("OLLAMA_URL", "http://152.53.201.152:23405/api/generate")
REDIS_URL = os.getenv("REDIS_URL", "redis://:e2ZyYYr4oWRdASI2CaLc-@redis-genesis-u50607.vm.elestio.app:26379")
POSTGRES_URL = os.getenv("POSTGRES_URL", "postgresql://user:password@host:port/database")  # Replace with your actual URL
QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
N8N_URL = os.getenv("N8N_URL", "http://localhost:5678/webhook")
AIVA_MODEL = os.getenv("AIVA_MODEL", "qwen-long")
AIVA_TIMEOUT = float(os.getenv("AIVA_TIMEOUT", "300"))  # 5 min default
RATE_LIMIT_PER_MINUTE = int(os.getenv("RATE_LIMIT_PER_MINUTE", "60"))  # Default 60 requests/minute

# --- Helper Functions ---
async def call_external_api(url: str, method: str = "GET", headers: Dict[str, str] = None, data: Dict[str, Any] = None) -> Dict[str, Any]:
    """Generic function to call external APIs with retries."""
    max_retries = 3
    retry_delay = 2  # seconds

    async with httpx.AsyncClient() as client:
        for attempt in range(max_retries):
            try:
                if method == "GET":
                    response = await client.get(url, headers=headers, timeout=10)  # Reduced timeout for external APIs
                elif method == "POST":
                    response = await client.post(url, headers=headers, json=data, timeout=10)
                else:
                    raise ValueError(f"Unsupported method: {method}")

                response.raise_for_status()  # Raise HTTPError for bad responses (4xx or 5xx)
                return response.json()
            except httpx.HTTPError as e:
                logger.warning(f"API call failed (attempt {attempt + 1}/{max_retries}): {e}")
                if attempt < max_retries - 1:
                    await asyncio.sleep(retry_delay)
                else:
                    logger.error(f"API call failed after {max_retries} attempts: {e}")
                    raise  # Re-raise the exception if retries are exhausted
            except Exception as e:
                logger.exception(f"Unexpected error during API call: {e}")
                raise
        return {} # should not reach here


class IntegrationHub:
    """
    Central hub for integrating various services:
    - Ollama/QwenLong (Reasoning)
    - Redis (Pub/Sub)
    - PostgreSQL (Persistent Storage)
    - Qdrant (Vector Embeddings)
    - n8n (Workflow Automation)
    - External APIs (GHL, Stripe, Telegram)
    """

    def __init__(self):
        self.ollama_url = OLLAMA_URL
        self.redis_url = REDIS_URL
        self.postgres_url = POSTGRES_URL
        self.qdrant_url = QDRANT_URL
        self.n8n_url = N8N_URL
        self.model = AIVA_MODEL
        self.timeout = AIVA_TIMEOUT
        self.rate_limit = RATE_LIMIT_PER_MINUTE

        self._redis_client = None
        self._postgres_pool = None
        self._qdrant_client = None
        self._http_client = None  # For Ollama and potentially n8n

        self._rate_limit_queue = asyncio.Queue(maxsize=self.rate_limit)
        self._populate_rate_limit_queue()

        logger.info("IntegrationHub initialized.")

    async def _populate_rate_limit_queue(self):
        """Pre-populate the rate limit queue."""
        now = datetime.now()
        for _ in range(self.rate_limit):
            await self._rate_limit_queue.put(now)

    async def _wait_for_rate_limit(self):
        """Wait until the rate limit allows the next request."""
        await self._rate_limit_queue.get()
        self._rate_limit_queue.put_nowait(datetime.now())

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

    async def _get_redis_client(self) -> aioredis.Redis:
        """Get or create Redis client."""
        if self._redis_client is None:
            try:
                self._redis_client = aioredis.from_url(self.redis_url)
                await self._redis_client.ping()  # Check connection
                logger.info("Connected to Redis.")
            except Exception as e:
                logger.error(f"Failed to connect to Redis: {e}")
                self._redis_client = None  # Ensure it's None on failure
                raise  # Re-raise so the application knows it can't connect

        return self._redis_client

    async def _get_postgres_pool(self) -> asyncpg.Pool:
        """Get or create PostgreSQL connection pool."""
        if self._postgres_pool is None:
            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
                raise
        return self._postgres_pool

    async def _get_qdrant_client(self) -> QdrantClient:
        """Get or create Qdrant client."""
        if self._qdrant_client is None:
            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
                raise
        return self._qdrant_client

    # --- Health Checks ---
    async def health_check(self) -> Dict[str, Any]:
        """Comprehensive health check for all services."""
        health = {
            "ollama": await self.ollama_health_check(),
            "redis": await self.redis_health_check(),
            "postgres": await self.postgres_health_check(),
            "qdrant": await self.qdrant_health_check(),
            "n8n": await self.n8n_health_check()
        }
        overall_status = "online" if all(v.get("status") == "online" for v in health.values()) else "degraded"
        return {"status": overall_status, "services": health, "timestamp": datetime.now().isoformat()}

    async def ollama_health_check(self) -> Dict[str, Any]:
        """Check Ollama/AIVA 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 = [m.get("name") for m in data.get("models", [])]
            qwen_ready = self.model in models or f"{self.model}:latest" in models

            return {
                "status": "online" if qwen_ready else "model_missing",
                "models": models,
                "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()
            }

    async def redis_health_check(self) -> Dict[str, Any]:
        """Check Redis connection."""
        try:
            redis = await self._get_redis_client()
            await redis.ping()
            return {"status": "online", "timestamp": datetime.now().isoformat()}
        except Exception as e:
            return {"status": "offline", "error": str(e), "timestamp": datetime.now().isoformat()}

    async def postgres_health_check(self) -> Dict[str, Any]:
        """Check PostgreSQL connection."""
        try:
            pool = await self._get_postgres_pool()
            async with pool.acquire() as conn:
                await conn.execute("SELECT 1")  # Simple query
            return {"status": "online", "timestamp": datetime.now().isoformat()}
        except Exception as e:
            return {"status": "offline", "error": str(e), "timestamp": datetime.now().isoformat()}

    async def qdrant_health_check(self) -> Dict[str, Any]:
        """Check Qdrant connection."""
        try:
            client = await self._get_qdrant_client()
            await client.get_telemetry_report()
            return {"status": "online", "timestamp": datetime.now().isoformat()}
        except Exception as e:
            return {"status": "offline", "error": str(e), "timestamp": datetime.now().isoformat()}

    async def n8n_health_check(self) -> Dict[str, Any]:
        """Check n8n connection (simple GET request)."""
        try:
            client = await self._get_http_client()
            response = await client.get(self.n8n_url, timeout=5)  # Very short timeout
            response.raise_for_status()
            return {"status": "online", "timestamp": datetime.now().isoformat()}
        except Exception as e:
            return {"status": "offline", "error": str(e), "timestamp": datetime.now().isoformat()}

    # --- 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.

        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'
        """
        await self._wait_for_rate_limit()  # Enforce rate limit
        try:
            # Build full prompt with context
            full_prompt = prompt
            if context:
                context_str = json.dumps(context, indent=2)
                full_prompt = f"Context:\n