from sqlmodel import create_engine, SQLModel, Session
from sqlalchemy.pool import StaticPool, QueuePool
from sqlalchemy import text
from sqlalchemy import event
from src.config.settings import DB_PATH, DATABASE_URL
import logging
import time
import os

logger = logging.getLogger(__name__)

# Slow query threshold in seconds
SLOW_QUERY_THRESHOLD = float(os.getenv("SLOW_QUERY_THRESHOLD", "1.0"))

# Connection pool configuration from environment variables
POOL_SIZE = int(os.getenv("DB_POOL_SIZE", "20"))
MAX_OVERFLOW = int(os.getenv("DB_MAX_OVERFLOW", "10"))
POOL_TIMEOUT = int(os.getenv("DB_POOL_TIMEOUT", "30"))
POOL_RECYCLE = int(os.getenv("DB_POOL_RECYCLE", "3600"))
POOL_PRE_PING = os.getenv("DB_POOL_PRE_PING", "true").lower() == "true"

# Create the database engine with optimized connection pooling
if DATABASE_URL:
    # PostgreSQL Configuration with configurable connection pool
    engine = create_engine(
        DATABASE_URL,
        echo=False,
        poolclass=QueuePool,
        pool_size=POOL_SIZE,
        max_overflow=MAX_OVERFLOW,
        pool_timeout=POOL_TIMEOUT,
        pool_recycle=POOL_RECYCLE,
        pool_pre_ping=POOL_PRE_PING,
        # Additional optimizations
        connect_args={
            "connect_timeout": 10,
            "options": "-c statement_timeout=120000"  # 120 second statement timeout for AI operations
        }
    )
    logger.info(
        f"PostgreSQL engine created with pool_size={POOL_SIZE}, "
        f"max_overflow={MAX_OVERFLOW}, pool_timeout={POOL_TIMEOUT}s, "
        f"pool_recycle={POOL_RECYCLE}s, pool_pre_ping={POOL_PRE_PING}"
    )
else:
    # SQLite Configuration (Fallback)
    engine = create_engine(
        f"sqlite:///{DB_PATH}",
        echo=False,
        connect_args={
            "check_same_thread": False,
            "timeout": 30  # 30 second busy timeout
        },
        poolclass=StaticPool,  # Use StaticPool for SQLite
    )
    logger.info(f"SQLite engine created with StaticPool at {DB_PATH}")


# Add slow query logging
@event.listens_for(engine, "before_cursor_execute")
def before_cursor_execute(conn, cursor, statement, parameters, context, executemany):
    """ 记录 query start time."""
    conn.info.setdefault("query_start_time", []).append(time.time())


@event.listens_for(engine, "after_cursor_execute")
def after_cursor_execute(conn, cursor, statement, parameters, context, executemany):
    """ 日志 slow queries."""
    total_time = time.time() - conn.info["query_start_time"].pop()
    
    if total_time > SLOW_QUERY_THRESHOLD:
        logger.warning(
            f"Slow query detected (took {total_time:.2f}s): {statement[:200]}",
            extra={
                "query_time": total_time,
                "query": statement[:500],
                "parameters": str(parameters)[:200] if parameters else None
            }
        )


def init_db():
    """ Initialize the database tables."""
    from sqlalchemy.exc import OperationalError

    # Try to create all tables
    try:
        SQLModel.metadata.create_all(engine)
    except OperationalError as e:
        if "already exists" in str(e):
            # Table or index already exists, skip creation
            logger.warning(f"Skipping table/index creation: {e}")
        else:
            raise


def get_session():
    """ 依赖 for getting a database session.
    Ensures proper cleanup and connection management.
    """
    with Session(engine) as session:
        try:
            yield session
        finally:
            # 会话 is automatically closed by context manager
            # This ensures connections are returned to the pool
            pass


def get_pool_status() -> dict:
    """ Get current connection pool status for monitoring.
    Returns pool statistics including size, checked out connections, etc.
    """
    pool = engine.pool
    
    status = {
        "pool_size": getattr(pool, "size", lambda: 0)(),
        "checked_out": getattr(pool, "checkedout", lambda: 0)(),
        "overflow": getattr(pool, "overflow", lambda: 0)(),
        "checked_in": getattr(pool, "checkedin", lambda: 0)(),
    }
    
    # Add pool-specific attributes if available
    if hasattr(pool, "_pool"):
        status["available"] = pool._pool.qsize()
    
    return status


def check_database_health() -> tuple[bool, str]:
    """
    Check database connectivity and health.
    Returns (is_healthy, message).
    """
    try:
        with Session(engine) as session:
            # Simple connectivity probe without model dependency.
            session.exec(text("SELECT 1"))
        return True, "Database connection healthy"
    except Exception as e:
        logger.error(f"Database health check failed: {e}")
        return False, f"Database connection failed: {str(e)}"


def check_pool_health() -> dict:
    """
    Comprehensive pool health check with recommendations.
    Returns detailed pool status and health indicators.
    """
    status = get_pool_status()
    pool = engine.pool

    # Calculate health metrics
    total_connections = status.get("pool_size", 0) + status.get("overflow", 0)
    checked_out = status.get("checked_out", 0)
    available = status.get("available", 0)

    # Health checks
    health = {
        "status": "healthy",
        "checks": {},
        "recommendations": [],
        "pool_status": status
    }

    # Check 1: Connection exhaustion
    if checked_out >= total_connections * 0.9:
        health["checks"]["exhaustion"] = "critical"
        health["status"] = "critical"
        health["recommendations"].append(
            f"Pool near exhaustion: {checked_out}/{total_connections} connections in use. "
            f"Consider increasing DB_POOL_SIZE or DB_MAX_OVERFLOW."
        )
    elif checked_out >= total_connections * 0.75:
        health["checks"]["exhaustion"] = "warning"
        health["status"] = "warning"
        health["recommendations"].append(
            f"Pool usage high: {checked_out}/{total_connections} connections in use."
        )
    else:
        health["checks"]["exhaustion"] = "healthy"

    # Check 2: Available connections
    if available == 0 and total_connections > 0:
        health["checks"]["availability"] = "warning"
        if health["status"] == "healthy":
            health["status"] = "warning"
    else:
        health["checks"]["availability"] = "healthy"

    # Check 3: Pool size configuration
    if total_connections < 5:
        health["recommendations"].append(
            "Pool size may be too small for production workloads. "
            "Consider DB_POOL_SIZE >= 10"
        )

    return health
