import time
import logging
from django.db import connection
from django.conf import settings

logger = logging.getLogger(__name__)

class SQLTimingMiddleware:
    """生产环境慢 SQL 检测中间件"""
    # 默认阈值 500ms，可在 settings 中配置
    SLOW_SQL_THRESHOLD_MS = getattr(settings, "SLOW_SQL_THRESHOLD_MS", 500)

    def __init__(self, get_response):
        self.get_response = get_response

    def __call__(self, request):
        request_url = request.get_full_path()
        request_method = request.method
        client_ip = request.META.get("HTTP_X_FORWARDED_FOR",
                                     request.META.get("REMOTE_ADDR", ""))
        response = None

        total_sql_count = 0
        total_sql_time = 0.0

        def wrapper_execute(execute, sql, params, many, context):
            nonlocal total_sql_count, total_sql_time
            start = time.time()
            try:
                return execute(sql, params, many, context)
            finally:
                duration = (time.time() - start) * 1000
                total_sql_count += 1
                total_sql_time += duration
                if duration > self.SLOW_SQL_THRESHOLD_MS:
                    logger.warning(
                        f"[SLOW SQL] {duration:.2f}ms - {request_method} {request_url} "
                        f"[IP: {client_ip}] - SQL: {sql}"
                    )

        with connection.execute_wrapper(wrapper_execute):
            response = self.get_response(request)

        # 可选：记录每个请求总 SQL 执行时间
        if total_sql_count > 0:
            logger.info(
                f"[SQL SUMMARY] {request_method} {request_url} [IP: {client_ip}] "
                f"{total_sql_count} queries, total {total_sql_time:.2f}ms"
            )

        return response