import express from "express";
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
import { fetchServerConfig, MCPAuth } from "mcp-auth";
import { WachtClient } from "@wacht/backend";

const app = express();
app.use(express.json());

const wacht = new WachtClient({
  apiKey: process.env.WACHT_BACKEND_API_KEY!,
  baseUrl: "https://api.wacht.dev",
});

function createMcpServer() {
  const server = new McpServer({
    name: "whoami",
    version: "1.0.0",
  });

  server.registerTool(
    "whoami",
    {
      description: "Return current authenticated subject",
      inputSchema: {},
    },
    async (_, { authInfo }) => {
      if (!authInfo.subject) {
        return {
          content: [{ type: "text", text: JSON.stringify({ error: "Not authenticated" }) }],
        };
      }

      return {
        content: [{ type: "text", text: `Hi ${authInfo.subject}` }],
      };
    },
  );

  return server;
}

async function main() {
  const oauthIssuer = process.env.WACHT_OAUTH_ISSUER!;
  const resource = process.env.MCP_RESOURCE!;

  const authServerConfig = await fetchServerConfig(oauthIssuer, { type: "oauth" });

  const mcpAuth = new MCPAuth({
    protectedResources: {
      metadata: {
        resource,
        authorizationServers: [authServerConfig],
      },
    },
  });

  app.use(mcpAuth.protectedResourceMetadataRouter());
  app.use(
    mcpAuth.bearerAuth(
      async (token) => {
        const exchange = await wacht.gateway.verifyOauthAccessTokenRequest(
          token,
          "POST",
          "mcp",
        );

        return {
          clientId: exchange.headers["x-wacht-oauth-client-id"],
          token,
          scopes: exchange.metadata?.scopes || [],
          resource: exchange.metadata?.oauth_resource,
          expiresAt: exchange.metadata?.expires_at,
          issuer: exchange.headers["x-wacht-oauth-issuer"] || oauthIssuer,
          subject: exchange.headers["x-wacht-granted-resource"],
        };
      },
      { resource },
    ),
  );

  app.post("/", async (req, res) => {
    const server = createMcpServer();
    const transport = new StreamableHTTPServerTransport({
      sessionIdGenerator: undefined,
    });

    await server.connect(transport);
    await transport.handleRequest(req, res, req.body);

    res.on("close", () => {
      transport.close();
      server.close();
    });
  });

  app.listen(6969);
}

main();

use wacht::{WachtClient, WachtConfig};
use wacht::gateway::{GatewayAuthzOptions, GatewayPrincipalType};

let api_key = std::env::var("WACHT_BACKEND_API_KEY").expect("WACHT_BACKEND_API_KEY is required");
let frontend_host = std::env::var("WACHT_FRONTEND_HOST").expect("WACHT_FRONTEND_HOST is required");
let client = WachtClient::new(WachtConfig::new(api_key, frontend_host)).expect("valid Wacht config");

// Simple verification helper
let _authz = client
    .gateway()
    .verify_oauth_access_token_request(access_token, "POST", "mcp")
    .await?;

// Optional strict verification with required permissions
let authz_with_perms = client
    .gateway()
    .check_authz_with_principal_type(
        GatewayPrincipalType::OauthAccessToken,
        access_token,
        "POST",
        "mcp",
        GatewayAuthzOptions {
            required_permissions: Some(vec!["mcp:invoke".to_string()]),
            ..Default::default()
        },
    )
    .await?;

if !authz_with_perms.allowed {
    // Reject request in your framework middleware/handler.
}

let principal = authz_with_perms.resolve_principal_context();
let scopes = principal.metadata.scopes;
let granted_resource = principal.metadata.granted_resource;