import logging
from django.db import connections
from rest_framework_simplejwt.authentication import JWTAuthentication

logger = logging.getLogger(__name__)

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

    def __call__(self, request):
        logger.info(f"===> [START] {request.method} {request.path}")

        # 1. Try to get user from standard Django session
        user = getattr(request, "user", None)

        # 2. Manual JWT Check (Crucial for DRF + Middleware)
        if not user or not user.is_authenticated:
            try:
                jwt_authenticator = JWTAuthentication()
                header = jwt_authenticator.get_header(request)
                if header:
                    raw_token = jwt_authenticator.get_raw_token(header)
                    validated_token = jwt_authenticator.get_validated_token(raw_token)
                    user = jwt_authenticator.get_user(validated_token)
                    request.user = user  # Manually attach user to request
                    logger.info(f"      Auth: {user.username} (Authenticated via JWT)")
            except Exception:
                # This will happen on Login or Public endpoints
                logger.info(f"      Auth: Anonymous - Staying on 'default' DB")

        # 3. Database Switching Logic
        if request.user and request.user.is_authenticated:
            try:
                # Handle different possible profile related_names
                profile = getattr(request.user, "userprofile", None) or getattr(request.user, "profile", None)
                
                if profile and profile.school:
                    school = profile.school
                    request.school = school
                    
                    # Log the intent
                    logger.info(f"      Target: School '{school.name}' | DB: {school.db_name}")

                    # Update connection settings dynamically
                    connections["school"].settings_dict.update({
                        "NAME": school.db_name,
                        "USER": getattr(school, 'db_user', 'postgres'),
                        "PASSWORD": getattr(school, 'db_password', '1234'),
                        "HOST": getattr(school, 'db_host', 'localhost'),
                        "PORT": getattr(school, 'db_port', '5432'),
                    })

                    # ⚠️ Close to force refresh on the next ORM query
                    connections["school"].close()
                    logger.info(f"      Result: Successfully routed 'school' alias to {school.db_name}")
                else:
                    logger.warning(f"      Result: User {request.user.username} has NO school assigned.")
            
            except Exception as e:
                logger.error(f"      CRITICAL: DB Switch Error: {str(e)}")

        response = self.get_response(request)
        logger.info(f"<=== [END] Status: {response.status_code}\n")
        return response