from channels.db import database_sync_to_async


class TenantWebSocketMiddleware:
    def __init__(
        self,
        app,
        tenant_model,
        resolver,
        tenant_setter=None,
    ):
        self.app = app
        self.tenant_model = tenant_model
        self.resolver = resolver
        self.tenant_setter = tenant_setter

    async def __call__(self, scope, receive, send):

        headers = dict(scope["headers"])
        host = headers.get(b"host", b"").decode()

        tenant_name = self.resolver.resolve(host)

        tenant = None

        if tenant_name:
            tenant = await self.get_tenant(tenant_name)

        if tenant:
            scope["tenant"] = tenant

            if self.tenant_setter:
                await self.tenant_setter(tenant)

        return await self.app(scope, receive, send)

    @database_sync_to_async
    def get_tenant(self, tenant_name):
        try:
            return self.tenant_model.objects.get(schema_name=tenant_name)
        except self.tenant_model.DoesNotExist:
            return None
