using System.Net.Http.Headers;
using System.Security.Claims;
using Microsoft.AspNetCore.Authentication.JwtBearer;
using Microsoft.IdentityModel.Tokens;
using ModelContextProtocol.AspNetCore.Authentication;
using ModelContextProtocol.Types.Authentication;
using ProtectedMCPServer.Tools;

var builder = WebApplication.CreateBuilder(args);

var serverUrl = "http://localhost:7071/";
var tenantId = "a2213e1c-e51e-4304-9a0d-effe57f31655";
var instance = "https://login.microsoftonline.com/";

builder.Services.AddAuthentication(options =>
{
    options.DefaultChallengeScheme = McpAuthenticationDefaults.AuthenticationScheme;
    options.DefaultAuthenticateScheme = JwtBearerDefaults.AuthenticationScheme;
})
.AddJwtBearer(options =>
{
    options.Authority = $"{instance}{tenantId}/v2.0";
    options.TokenValidationParameters = new TokenValidationParameters
    {
        ValidateIssuer = true,
        ValidateAudience = true,
        ValidateLifetime = true,
        ValidateIssuerSigningKey = true,
        ValidAudience = "167b4284-3f92-4436-92ed-38b38f83ae08",
        ValidIssuer = $"{instance}{tenantId}/v2.0",
        NameClaimType = "name",
        RoleClaimType = "roles"
    };

    options.MetadataAddress = $"{instance}{tenantId}/v2.0/.well-known/openid-configuration";

    options.Events = new JwtBearerEvents
    {
        OnTokenValidated = context =>
        {
            var name = context.Principal?.Identity?.Name ?? "unknown";
            var email = context.Principal?.FindFirstValue("preferred_username") ?? "unknown";
            Console.WriteLine($"Token validated for: {name} ({email})");
            return Task.CompletedTask;
        },
        OnAuthenticationFailed = context =>
        {
            Console.WriteLine($"Authentication failed: {context.Exception.Message}");
            return Task.CompletedTask;
        },
        OnChallenge = context =>
        {
            Console.WriteLine($"Challenging client to authenticate with Entra ID");
            return Task.CompletedTask;
        }
    };
})
.AddMcp(options =>
{
    options.ResourceMetadataProvider = context => 
    {
        var metadata = new ProtectedResourceMetadata
        {
            BearerMethodsSupported = { "header" },
            ResourceDocumentation = new Uri("https://docs.example.com/api/weather"),
            AuthorizationServers = { new Uri($"{instance}{tenantId}/v2.0") }
        };

        metadata.ScopesSupported.AddRange(new[] {
            "api://167b4284-3f92-4436-92ed-38b38f83ae08/weather.read" 
        });
        
        return metadata;
    };
});

builder.Services.AddAuthorization();

builder.Services.AddHttpContextAccessor();
builder.Services.AddMcpServer()
.WithTools<WeatherTools>()
.WithHttpTransport();

builder.Services.AddSingleton(_ =>
{
    var client = new HttpClient() { BaseAddress = new Uri("https://api.weather.gov") };
    client.DefaultRequestHeaders.UserAgent.Add(new ProductInfoHeaderValue("weather-tool", "1.0"));
    return client;
});

var app = builder.Build();

app.UseAuthentication();
app.UseAuthorization();

// Use the default MCP policy name that we've configured
app.MapMcp().RequireAuthorization();

Console.WriteLine($"Starting MCP server with authorization at {serverUrl}");
Console.WriteLine($"PRM Document URL: {serverUrl}.well-known/oauth-protected-resource");
Console.WriteLine("Press Ctrl+C to stop the server");

app.Run(serverUrl);

[McpServerTool, Description("Get weather alerts for a US state.")]
public static async Task<string> GetAlerts(
    HttpClient client,
    IHttpContextAccessor httpContextAccessor,
    [Description("The US state to get alerts for. Use the 2 letter abbreviation for the state (e.g. NY).")] string state)

using ModelContextProtocol.Client;
using ModelContextProtocol.Protocol.Transport;

namespace ProtectedMCPClient;

class Program
{
    static async Task Main(string[] args)
    {
        Console.WriteLine("Protected MCP Weather Server");
        Console.WriteLine();

        var serverUrl = "http://localhost:7071/sse";

        var tokenProvider = new BasicOAuthAuthorizationProvider(
            new Uri(serverUrl), 
            clientId: "6ad97b5f-7a7b-413f-8603-7a3517d4adb8",
            redirectUri: new Uri("http://localhost:1179/callback"),
            scopes: ["api://167b4284-3f92-4436-92ed-38b38f83ae08/weather.read"]
        );

        Console.WriteLine();
        Console.WriteLine($"Connecting to weather server at {serverUrl}...");

        try
        {
            var transportOptions = new SseClientTransportOptions
            {
                Endpoint = new Uri(serverUrl),
                Name = "Secure Weather Client"
            };

            var transport = new SseClientTransport(transportOptions, tokenProvider);
            var client = await McpClientFactory.CreateAsync(transport);

            var tools = await client.ListToolsAsync();
            if (tools.Count == 0)
            {
                Console.WriteLine("No tools available on the server.");
                return;
            }

            Console.WriteLine($"Found {tools.Count} tools on the server.");
            Console.WriteLine();

            if (tools.Any(t => t.Name == "GetAlerts"))
            {
                Console.WriteLine("Calling GetAlerts tool...");

                var result = await client.CallToolAsync(
                    "GetAlerts",
                    new Dictionary<string, object?> { { "state", "WA" } }
                );

                Console.WriteLine("Result: " + result.Content[0].Text);
                Console.WriteLine();
            }
        }
        catch (Exception ex)
        {
            Console.WriteLine($"Error: {ex.Message}");
            if (ex.InnerException != null)
            {
                Console.WriteLine($"Inner error: {ex.InnerException.Message}");
            }

            #if DEBUG
            Console.WriteLine($"Stack trace: {ex.StackTrace}");
            #endif
        }

        Console.WriteLine("Press any key to exit...");
        Console.ReadKey();
    }
}