// During startup configuration validation
if (config.UserDelegatedAuth?.Enabled == true)
{
    if (config.DatabaseType != "mssql")
        throw new ConfigurationException(
            "user-delegated-auth is only supported for database-type 'mssql'");
    
    if (string.IsNullOrWhiteSpace(config.UserDelegatedAuth.DatabaseAudience))
        throw new ConfigurationException(
            "database-audience is required when user-delegated-auth is enabled");
    
    if (config.UserDelegatedAuth.TokenCacheDurationMinutes is < 1 or > 59)
        throw new ConfigurationException(
            "token-cache-duration-minutes must be between 1 and 59");
}

// OboSqlTokenProvider.cs
using System;
using System.Threading.Tasks;
using Microsoft.Identity.Client;
using Microsoft.AspNetCore.Http;
using System.Collections.Concurrent;

/// <summary>
/// Provides Azure SQL access tokens via On-Behalf-Of flow for per-user database authentication.
/// </summary>
public sealed class OboSqlTokenProvider
{
    private readonly string[] _scope;
    private readonly IConfidentialClientApplication _cca;
    private readonly ConcurrentDictionary<string, (string token, DateTimeOffset exp)> _cache = new();
    private readonly TimeSpan _earlyRefresh;

    /// <summary>
    /// Initialize the OBO token provider.
    /// </summary>
    /// <param name="tenantId">Azure AD tenant ID</param>
    /// <param name="clientId">DAB application (client) ID</param>
    /// <param name="clientSecret">DAB application client secret</param>
    /// <param name="databaseAudience">Azure SQL resource identifier (e.g., https://database.windows.net)</param>
    /// <param name="tokenCacheDurationMinutes">Cache duration in minutes (default: 50)</param>
    public OboSqlTokenProvider(
        string tenantId, 
        string clientId, 
        string clientSecret, 
        string databaseAudience,
        int tokenCacheDurationMinutes = 50)
    {
        if (string.IsNullOrWhiteSpace(tenantId))
            throw new ArgumentNullException(nameof(tenantId));
        if (string.IsNullOrWhiteSpace(clientId))
            throw new ArgumentNullException(nameof(clientId));
        if (string.IsNullOrWhiteSpace(clientSecret))
            throw new ArgumentNullException(nameof(clientSecret));
        if (string.IsNullOrWhiteSpace(databaseAudience))
            throw new ArgumentNullException(nameof(databaseAudience));
        if (tokenCacheDurationMinutes is < 1 or > 59)
            throw new ArgumentOutOfRangeException(nameof(tokenCacheDurationMinutes));

        _scope = new[] { $"{databaseAudience.TrimEnd('/')}/.default" };
        _earlyRefresh = TimeSpan.FromMinutes(Math.Max(5, 60 - tokenCacheDurationMinutes));
        
        _cca = ConfidentialClientApplicationBuilder
            .Create(clientId)
            .WithClientSecret(clientSecret)
            .WithAuthority($"https://login.microsoftonline.com/{tenantId}")
            .Build();
    }

    /// <summary>
    /// Acquire an Azure SQL access token on behalf of the authenticated user.
    /// </summary>
    /// <param name="ctx">Current HTTP context containing the user's bearer token</param>
    /// <returns>Azure SQL access token valid for the user</returns>
    /// <exception cref="InvalidOperationException">When bearer token is missing</exception>
    public async Task<string> GetAccessTokenAsync(HttpContext ctx)
    {
        string inboundJwt = ExtractBearerToken(ctx) 
            ?? throw new InvalidOperationException("No bearer token found in Authorization header.");
        
        // Cache key based on user identity (oid = object ID, tid = tenant ID)
        string oid = ctx.User.FindFirst("oid")?.Value 
            ?? throw new InvalidOperationException("Token missing 'oid' claim.");
        string tid = ctx.User.FindFirst("tid")?.Value 
            ?? throw new InvalidOperationException("Token missing 'tid' claim.");
        string cacheKey = $"{tid}:{oid}";

        // Check cache for unexpired token
        if (_cache.TryGetValue(cacheKey, out var cached) 
            && cached.exp - _earlyRefresh > DateTimeOffset.UtcNow)
        {
            return cached.token;
        }

        // Perform OBO exchange
        var result = await _cca
            .AcquireTokenOnBehalfOf(_scope, new UserAssertion(inboundJwt))
            .ExecuteAsync();
        
        _cache[cacheKey] = (result.AccessToken, result.ExpiresOn);
        return result.AccessToken;
    }

    private static string? ExtractBearerToken(HttpContext ctx)
    {
        string? authHeader = ctx.Request.Headers.Authorization.ToString();
        return authHeader?.StartsWith("Bearer ", StringComparison.OrdinalIgnoreCase) == true
            ? authHeader.Substring(7)
            : null;
    }
}

// MsSqlQueryExecutor.cs changes
using Microsoft.Data.SqlClient;
using System.Data.Common;

public class MsSqlQueryExecutor : BaseQueryExecutor
{
    private readonly OboSqlTokenProvider? _oboTokenProvider;
    private readonly bool _disablePooling;

    // Constructor injection
    public MsSqlQueryExecutor(
        RuntimeConfigProvider configProvider,
        IHttpContextAccessor httpContextAccessor,
        OboSqlTokenProvider? oboTokenProvider = null,  // injected when user-delegated-auth enabled
        ILogger<MsSqlQueryExecutor> logger = null)
        : base(configProvider, httpContextAccessor, logger)
    {
        _oboTokenProvider = oboTokenProvider;
        
        // Read pooling preference from config
        var config = configProvider.GetConfig();
        _disablePooling = config.DataSource?.UserDelegatedAuth?.DisableConnectionPooling ?? true;
    }

    public override async Task SetManagedIdentityAccessTokenIfAnyAsync(
        DbConnection conn, 
        string dataSourceName)
    {
        if (_oboTokenProvider is not null) // user-delegated mode
        {
            var httpContext = HttpContextAccessor.HttpContext 
                ?? throw new InvalidOperationException(
                    "HttpContext unavailable for user-delegated authentication.");
            
            // Acquire per-user token via OBO
            string userSqlToken = await _oboTokenProvider.GetAccessTokenAsync(httpContext);
            
            var sqlConn = (SqlConnection)conn;

            // Apply pooling configuration
            if (_disablePooling)
            {
                var builder = new SqlConnectionStringBuilder(sqlConn.ConnectionString) 
                { 
                    Pooling = false 
                };
                sqlConn.ConnectionString = builder.ConnectionString;
            }

            // Set per-user access token
            sqlConn.AccessToken = userSqlToken;
            return;
        }

        // Default behavior: application identity (managed identity or connection string auth)
        await base.SetManagedIdentityAccessTokenIfAnyAsync(conn, dataSourceName);
    }
}

// Startup.cs or Program.cs
public void ConfigureServices(IServiceCollection services)
{
    var config = LoadRuntimeConfig();
    
    // Register OBO provider conditionally
    if (config.DataSource?.UserDelegatedAuth?.Enabled == true)
    {
        var oboConfig = config.DataSource.UserDelegatedAuth;
        var tenantId = configuration["AzureAd:TenantId"] 
            ?? throw new InvalidOperationException("AzureAd:TenantId required for OBO");
        var clientId = configuration["AzureAd:ClientId"] 
            ?? throw new InvalidOperationException("AzureAd:ClientId required for OBO");
        var clientSecret = configuration["AzureAd:ClientSecret"] 
            ?? throw new InvalidOperationException("AzureAd:ClientSecret required for OBO");

        services.AddSingleton(new OboSqlTokenProvider(
            tenantId,
            clientId,
            clientSecret,
            oboConfig.DatabaseAudience,
            oboConfig.TokenCacheDurationMinutes ?? 50
        ));
    }
    else
    {
        // Register null when not in OBO mode (default behavior)
        services.AddSingleton<OboSqlTokenProvider?>(null);
    }

    services.AddScoped<MsSqlQueryExecutor>();
}