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

# Optional dependencies (install if needed)
try:
    import httpx
except ImportError:
    httpx = None

try:
    import redis.asyncio as 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

# Load environment variables (using .env or os.environ)
ENV_PATH = Path(__file__).parent / "config.env"  # Or wherever your config is
if ENV_PATH.exists():
    try:
        from dotenv import load_dotenv
        load_dotenv(ENV_PATH)
    except ImportError:
        print("dotenv not installed. Please install it to load environment variables from .env file.")

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


class IntegrationHub:
    """
    Unified integration hub for connecting various services.
    """

    def __init__(self):
        # Core Services
        self.ollama_url = os.getenv("OLLAMA_URL", "http://localhost:11434/api/generate")  # Default Ollama URL
        self.ollama_model = os.getenv("OLLAMA_MODEL", "mistral")  # Default model
        self.redis_url = os.getenv("REDIS_URL", "redis://localhost:6379")  # Default Redis URL
        self.postgres_url = os.getenv("POSTGRES_URL", "postgresql://user:password@localhost:5432/database")  # Default PostgreSQL URL
        self.qdrant_url = os.getenv("QDRANT_URL", "http://localhost:6333") # Default Qdrant URL
        self.qdrant_api_key = os.getenv("QDRANT_API_KEY") # Optional Qdrant API key
        self.n8n_url = os.getenv("N8N_URL") # Optional N8N URL
        self.rate_limit_per_second = int(os.getenv("RATE_LIMIT_PER_SECOND", "10")) # Default rate limit
        self.timeout = float(os.getenv("AIVA_TIMEOUT", "300"))  # Default timeout for httpx requests

        # Connection pools and clients
        self._http_client = None
        self._redis_client = None
        self._postgres_pool = None
        self._qdrant_client = None
        self.is_connected = False  # Overall health status

        # Rate limiting
        self._request_queue = asyncio.Queue()
        for _ in range(self.rate_limit_per_second):
            self._request_queue.put_nowait(None)  # Fill the queue with "tokens"
        self._rate_limit_task = None # To be started in aconnect
        self._is_closing = False # Flag to stop background tasks gracefully

        logger.info("IntegrationHub initialized.")

    async def aconnect(self):
        """Asynchronously connect to all services."""
        try:
            await asyncio.gather(
                self._connect_redis(),
                self._connect_postgres(),
                self._connect_qdrant(),
                self._health_check()  # Initial health check
            )
            self.is_connected = True
            logger.info("Successfully connected to all services.")
            self._rate_limit_task = asyncio.create_task(self._rate_limit_replenish()) # Start rate limiting

        except Exception as e:
            self.is_connected = False
            logger.error(f"Failed to connect to all services: {e}")

    async def _rate_limit_replenish(self):
        """Background task to replenish the rate limit queue."""
        while not self._is_closing:
            try:
                await asyncio.sleep(1)  # Replenish every second
                for _ in range(self.rate_limit_per_second):
                    self._request_queue.put_nowait(None)
            except asyncio.CancelledError:
                logger.info("Rate limit replenishing task cancelled.")
                break
            except Exception as e:
                logger.error(f"Error replenishing rate limit: {e}")
                break

    async def _acquire_token(self):
        """Acquire a rate limit token."""
        await self._request_queue.get()
        self._request_queue.task_done()

    async def _release_token(self):
        """Release a rate limit token (not needed with current implementation)."""
        pass  # Token is automatically replenished

    # --- Core Service Connections ---

    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):
        """Connect to Redis."""
        if aioredis is None:
            logger.warning("Redis not installed. Skipping Redis connection.")
            return
        try:
            self._redis_client = aioredis.from_url(self.redis_url)
            await self._redis_client.ping()  # Verify connection
            logger.info("Connected to Redis.")
        except Exception as e:
            self._redis_client = None
            logger.error(f"Failed to connect to Redis: {e}")

    async def _connect_postgres(self):
        """Connect to PostgreSQL using a connection pool."""
        if asyncpg is None:
            logger.warning("asyncpg not installed. Skipping PostgreSQL connection.")
            return

        try:
            self._postgres_pool = await asyncpg.create_pool(self.postgres_url)
            async with self._postgres_pool.acquire() as conn:
                await conn.execute("SELECT 1;")  # Verify connection
            logger.info("Connected to PostgreSQL.")
        except Exception as e:
            self._postgres_pool = None
            logger.error(f"Failed to connect to PostgreSQL: {e}")

    async def _connect_qdrant(self):
        """Connect to Qdrant."""
        if QdrantClient is None or models is None:
            logger.warning("Qdrant client not installed. Skipping Qdrant connection.")
            return

        try:
            self._qdrant_client = QdrantClient(
                url=self.qdrant_url,
                api_key=self.qdrant_api_key,
            )
            self._qdrant_client.get_telemetry_data() # Verify connection
            logger.info("Connected to Qdrant.")
        except Exception as e:
            self._qdrant_client = None
            logger.error(f"Failed to connect to Qdrant: {e}")

    async def _health_check(self) -> Dict[str, Any]:
        """Check the health status of all connected services."""
        health_status = {
            "timestamp": datetime.now().isoformat(),
            "ollama": await self._check_ollama_health(),
            "redis": await self._check_redis_health(),
            "postgres": await self._check_postgres_health(),
            "qdrant": await self._check_qdrant_health(),
            "n8n": await self._check_n8n_health(),
        }
        self.is_connected = all(v.get("status") == "online" for k, v in health_status.items())
        return health_status

    async def _check_ollama_health(self) -> Dict[str, Any]:
        """Check Ollama health status."""
        tags_url = self.ollama_url.replace("/api/generate", "/api/tags")
        try:
            if httpx:
                client = await self._get_http_client()
                response = await client.get(tags_url, timeout=10.0)
                response.raise_for_status()
                data = response.json()
            else:
                # Fallback to urllib
                import urllib.request
                req = urllib.request.Request(tags_url)
                with urllib.request.urlopen(req, timeout=10) as resp:
                    data = json.loads(resp.read().decode())

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

            return {
                "status": "online" if qwen_ready else "model_missing",
                "models": models,
                "target_model": self.ollama_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 _check_redis_health(self) -> Dict[str, Any]:
        """Check Redis health status."""
        if self._redis_client is None:
            return {"status": "disabled"}
        try:
            await self._redis_client.ping()
            return {"status": "online"}
        except Exception as e:
            return {"status": "offline", "error": str(e)}

    async def _check_postgres_health(self) -> Dict[str, Any]:
        """Check PostgreSQL health status."""
        if self._postgres_pool is None:
            return {"status": "disabled"}
        try:
            async with self._postgres_pool.acquire() as conn:
                await conn.execute("SELECT 1;")
            return {"status": "online"}
        except Exception as e:
            return {"status": "offline", "error": str(e)}

    async def _check_qdrant_health(self) -> Dict[str, Any]:
        """Check Qdrant health status."""
        if self._qdrant_client is None:
            return {"status": "disabled"}
        try:
            self._qdrant_client.get_telemetry_data()
            return {"status": "online"}
        except Exception as e:
            return {"status": "offline", "error": str(e)}

    async def _check_n8n_health(self) -> Dict[str, Any]:
        """Check n8n health status. (Simple HTTP check)."""
        if not self.n8n_url:
            return {"status": "disabled"}
        try:
            if httpx:
                client = await self._get_http_client()
                response = await client.get(self.n8n_url, timeout=5.0)
                response.raise_for_status()  # Raise HTTPError for bad responses (4xx or 5xx)
                return {"status": "online", "code": response.status_code}
            else:
                import urllib.request
                req = urllib.request.Request(self.n8n_url)
                with urllib.request.urlopen(req, timeout=5) as resp:
                    return {"status": "online", "code": resp.getcode()}

        except Exception as e:
            return {"status": "offline", "error": str(e)}

    # --- Unified Interface Methods ---

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