"""FastMCP middleware that injects TenantContext before every tool call (Phase 2A).

Resolution order for ``tenant_id``:
1. Tool argument ``tenant_id`` (explicit per-call override)
2. Environment variable ``DATABRIDGE_TENANT_ID``
3. Constructor default (``"default"``)
"""
from __future__ import annotations

import os
from typing import Any

from fastmcp.server.middleware import Middleware

from .context import TenantContext, reset_current_tenant, set_current_tenant
from .registry import TenantRegistry

TENANT_ENV_VAR = "DATABRIDGE_TENANT_ID"


class TenantMiddleware(Middleware):
    """Sets ``_current_tenant`` contextvar around every tool invocation."""

    def __init__(
        self,
        registry: TenantRegistry,
        default_tenant_id: str = "default",
    ):
        self.registry = registry
        self.default_tenant_id = default_tenant_id

    async def on_call_tool(self, context, call_next):
        # --- Resolve tenant_id ------------------------------------------
        tenant_id = self.default_tenant_id

        # Priority 2: env var
        env_tenant = os.environ.get(TENANT_ENV_VAR)
        if env_tenant:
            tenant_id = env_tenant

        # Priority 1: explicit tool argument (highest priority)
        message = context.message
        args = getattr(message, "arguments", None) or {}
        if "tenant_id" in args:
            tenant_id = args["tenant_id"]

        # --- Build context and inject -----------------------------------
        tenant_ctx = self.registry.resolve(tenant_id)
        token = set_current_tenant(tenant_ctx)
        try:
            result = await call_next(context)
        finally:
            reset_current_tenant(token)

        return result
