import pyodbc
import psycopg2
from psycopg2.extras import RealDictCursor
import asyncpg
import redis
from contextlib import contextmanager
import logging
import time
from typing import Optional
from .config import config

logger = logging.getLogger(__name__)

class DatabaseConnections:
    """Manage database connections with retry logic and timeouts"""
    
    def __init__(self):
        self.mssql_conn = None
        self.postgres_conn = None
        self.redis_client = None
        self.pg_pool = None
        self.max_retries = 3
        self.retry_delay = 2  # seconds
    
    def _retry_connection(self, connection_func, connection_type: str, max_retries: Optional[int] = None):
        """Retry connection with exponential backoff"""
        max_retries = max_retries or self.max_retries
        
        for attempt in range(max_retries):
            try:
                return connection_func()
            except Exception as e:
                if attempt == max_retries - 1:
                    logger.error(f"Failed to connect to {connection_type} after {max_retries} attempts: {e}")
                    raise
                
                wait_time = self.retry_delay * (2 ** attempt)
                logger.warning(f"{connection_type} connection failed (attempt {attempt + 1}/{max_retries}): {e}")
                logger.info(f"Retrying in {wait_time} seconds...")
                time.sleep(wait_time)
    
    @contextmanager
    def mssql_connection(self):
        """Context manager for MS SQL connection using pyodbc with retry logic"""
        def _connect():
            connection_string = (
                f"DRIVER={{ODBC Driver 17 for SQL Server}};"
                f"SERVER={config.MSSQL_HOST},{config.MSSQL_PORT};"
                f"DATABASE={config.MSSQL_DATABASE};"
                f"UID={config.MSSQL_USER};"
                f"PWD={config.MSSQL_PASSWORD};"
                f"TrustServerCertificate=yes;"
                f"Connection Timeout=30;"
                f"Command Timeout=300;"
            )
            return pyodbc.connect(connection_string)
        
        conn = self._retry_connection(_connect, "MSSQL")
        try:
            yield conn
        finally:
            try:
                conn.close()
            except Exception as e:
                logger.warning(f"Error closing MSSQL connection: {e}")
    
    @contextmanager
    def postgres_connection(self):
        """Context manager for PostgreSQL connection with retry logic"""
        def _connect():
            return psycopg2.connect(
                host=config.POSTGRES_HOST,
                port=config.POSTGRES_PORT,
                database=config.POSTGRES_DATABASE,
                user=config.POSTGRES_USER,
                password=config.POSTGRES_PASSWORD,
                cursor_factory=RealDictCursor,
                connect_timeout=30,
                options='-c statement_timeout=300000'  # 5 minutes
            )
        
        conn = self._retry_connection(_connect, "PostgreSQL")
        try:
            yield conn
        finally:
            try:
                conn.close()
            except Exception as e:
                logger.warning(f"Error closing PostgreSQL connection: {e}")
    
    async def create_pg_pool(self):
        """Create async PostgreSQL connection pool"""
        self.pg_pool = await asyncpg.create_pool(
            host=config.POSTGRES_HOST,
            port=config.POSTGRES_PORT,
            database=config.POSTGRES_DATABASE,
            user=config.POSTGRES_USER,
            password=config.POSTGRES_PASSWORD,
            min_size=10,
            max_size=20
        )
        return self.pg_pool
    
    def get_redis_client(self):
        """Get Redis client"""
        if not self.redis_client:
            self.redis_client = redis.Redis(
                host=config.REDIS_HOST,
                port=config.REDIS_PORT,
                db=config.REDIS_DB,
                decode_responses=True
            )
        return self.redis_client

def test_connections():
    """Test all database connections for health check"""
    try:
        # Test MSSQL connection (use master DB to avoid failures before restore)
        db = DatabaseConnections()
        mssql_master_conn_str = (
            f"DRIVER={{ODBC Driver 17 for SQL Server}};"
            f"SERVER={config.MSSQL_HOST},{config.MSSQL_PORT};"
            f"DATABASE=master;"
            f"UID={config.MSSQL_USER};"
            f"PWD={config.MSSQL_PASSWORD};"
            f"TrustServerCertificate=yes;"
        )
        import pyodbc as _pyodbc
        with _pyodbc.connect(mssql_master_conn_str) as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT 1")
            cursor.fetchone()
            logger.info("MSSQL connection successful (master)")
        
        # Test PostgreSQL connection
        with db.postgres_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("SELECT 1")
            cursor.fetchone()
            logger.info("PostgreSQL connection successful")
        
        # Test Redis connection
        redis_client = db.get_redis_client()
        redis_client.ping()
        logger.info("Redis connection successful")
        
        return True
    except Exception as e:
        logger.error(f"Connection test failed: {e}")
        return False

db_connections = DatabaseConnections()