from channels.db import database_sync_to_async
from django.contrib.auth.models import AnonymousUser
from rest_framework.authtoken.models import Token
from django.db import close_old_connections, connection

@database_sync_to_async
def get_user(token_key):
    """
    Look up user by token key.

    IMPORTANT: Tokens are stored in the public schema, so we need to
    explicitly query from public schema in multi-tenant setup.
    """
    try:
        # Use public schema for token lookup (tokens are in shared apps)
        from django_tenants.utils import schema_context
        with schema_context('public'):
            token = Token.objects.select_related('user').get(key=token_key)
            return token.user
    except Token.DoesNotExist:
        return AnonymousUser()
    except Exception as e:
        print(f"TokenAuthMiddleware: Error looking up token: {e}", flush=True)
        return AnonymousUser()

class TokenAuthMiddleware:
    """
    Custom middleware to authenticate users via 'access_token' cookie
    using Django REST Framework's TokenAuthentication.
    """
    def __init__(self, inner):
        self.inner = inner

    async def __call__(self, scope, receive, send):
        print(f"TokenAuthMiddleware: Processing request for {scope.get('path')}", flush=True)
        # Close old database connections to prevent usage of timed out connections
        close_old_connections()

        # Check if user is already authenticated (e.g. by SessionMiddleware)
        if scope.get("user") and scope["user"].is_authenticated:
            print(f"TokenAuthMiddleware: User already authenticated: {scope['user']}", flush=True)
            return await self.inner(scope, receive, send)

        token_key = None

        # Try to authenticate with token from query string
        query_string = scope.get('query_string', b'').decode()
        if 'token=' in query_string:
            for param in query_string.split('&'):
                if param.startswith('token='):
                    token_key = param.split('=')[1]
                    print(f"TokenAuthMiddleware: Found token in query string: {token_key}", flush=True)
                    break

        # Try to authenticate with token from cookie if not found in query string
        if not token_key:
            headers = dict(scope['headers'])
            if b'cookie' in headers:
                cookies = headers[b'cookie'].decode().split('; ')
                print(f"TokenAuthMiddleware: Found cookies: {cookies}", flush=True)
                for cookie in cookies:
                    if cookie.strip().startswith('access_token='):
                        token_key = cookie.strip().split('=')[1]
                        break
            else:
                print("TokenAuthMiddleware: No cookie header found.", flush=True)
        
        if token_key:
            print(f"TokenAuthMiddleware: Found token key: {token_key}", flush=True)
            user = await get_user(token_key)
            if user:
                print(f"TokenAuthMiddleware: Authenticated user: {user}", flush=True)
                scope['user'] = user
            else:
                print("TokenAuthMiddleware: User not found for token.", flush=True)
        else:
            print("TokenAuthMiddleware: No access_token cookie or query param found.", flush=True)
        
        return await self.inner(scope, receive, send)
