import threading
import psutil
from django.http import JsonResponse
from django.utils.deprecation import MiddlewareMixin

# 使用 psutil 获取系统资源信息
total_memory_gb = psutil.virtual_memory().total / (1024 ** 3)
total_cpu_cores = psutil.cpu_count()

memory_per_request_gb = 128
cpu_per_request_percentage = 50

# 计算最大并发数量
max_concurrency_memory = total_memory_gb // memory_per_request_gb
max_concurrency_cpu = total_cpu_cores * (100 // cpu_per_request_percentage)
max_concurrency = min(max_concurrency_memory, max_concurrency_cpu)

# 设置信号量最大并发数
semaphore = threading.Semaphore(max_concurrency)

class ConcurrencyMiddleware(MiddlewareMixin):
    def process_request(self, request):
        if not semaphore.acquire(blocking=False):
            return JsonResponse({'error': 'Server is too busy. Please try again later.'}, status=503)

    def process_response(self, request, response):
        semaphore.release()
        return response

    def process_exception(self, request, exception):
        semaphore.release()
        return None