import fs from "fs";
import path from "path";
import os from "os";
import crypto from "node:crypto";
import express from "express";
import cors from "cors";
import { SSEServerTransport } from "@modelcontextprotocol/sdk/server/sse.js";
import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js";
import { MCP_PORT, JWT_SECRET } from "./config/index.js";
import {
    sessionCredentials,
    sessionApiClients,
    sessionServers,
    clearSessionCredentials,
    setSessionUserToken,
    setMcpSessionCredentials,
    getMcpSessionCredentials
} from "./auth/index.js";
import { fetchProjectInfo } from "./auth/localClient.js";
import { createSessionServer } from "./server.js";

// Active sessions - stores { transport, server, type, heartbeatInterval }
const activeSessions = new Map();

const ACCESS_TOKEN_TTL = 3600; // seconds

const base64url = (value) => Buffer.from(value).toString("base64url");

const previewToken = (value) => {
    if (!value) return "<empty>";
    if (value.length <= 16) return value;
    return `${value.slice(0, 12)}...${value.slice(-8)} (len=${value.length})`;
};

const signJwt = (payload, expiresInSeconds = ACCESS_TOKEN_TTL) => {
    const header = { alg: "HS256", typ: "JWT" };
    const now = Math.floor(Date.now() / 1000);
    const enrichedPayload = { iat: now, exp: now + expiresInSeconds, ...payload };

    const headerEncoded = base64url(JSON.stringify(header));
    const payloadEncoded = base64url(JSON.stringify(enrichedPayload));
    const signature = crypto
        .createHmac("sha256", JWT_SECRET)
        .update(`${headerEncoded}.${payloadEncoded}`)
        .digest("base64url");

    return `${headerEncoded}.${payloadEncoded}.${signature}`;
};

const verifyJwt = (token) => {
    try {
        const [headerEncoded, payloadEncoded, signature] = token.split(".");
        if (!headerEncoded || !payloadEncoded || !signature) {
            return null;
        }
        const expected = crypto
            .createHmac("sha256", JWT_SECRET)
            .update(`${headerEncoded}.${payloadEncoded}`)
            .digest("base64url");
        if (expected !== signature) {
            return null;
        }
        const payload = JSON.parse(Buffer.from(payloadEncoded, "base64url").toString());
        if (payload.exp && Math.floor(Date.now() / 1000) > payload.exp) {
            return null;
        }
        return payload;
    } catch (error) {
        console.error(`[MCP HTTP] JWT verification error: ${error.message}`);
        return null;
    }
};

const resolveProjectCredentials = async (projectName) => {
    try {
        const info = await fetchProjectInfo(projectName);
        if (!info.success) {
            throw new Error(info.error || "Failed to resolve project info");
        }
        return {
            token: info.token,
            tokenHash: info.tokenHash || null,
            website: info.domain,
            web_url: info.web_url,
            api_web_url: info.api_web_url || info.web_url,
            forge_host: info.forge_host || null,
            mode: info.mode || "local",
        };
    } catch (error) {
        throw new Error(`Failed to resolve project '${projectName}': ${error.message}`);
    }
};

/**
 * Configure credentials from request headers/query params for a session
 */
const configureSessionCredentials = async (sessionId, { token, tokenHash, website, web_url, userToken, projectName }) => {
    // Priority 1: Resolve via project name from local Python server
    if (projectName) {
        try {
            const projectCreds = await resolveProjectCredentials(projectName);
            sessionCredentials.set(sessionId, {
                token: projectCreds.token,
                tokenHash: projectCreds.tokenHash || null,
                website: projectCreds.website,
                web_url: projectCreds.web_url,
                api_web_url: projectCreds.api_web_url || projectCreds.web_url,
                forge_host: projectCreds.forge_host || null,
                mode: projectCreds.mode || "local",
                profileName: 'project-' + projectName,
                role: 'developer',
            });
            console.log(`[MCP] Session ${sessionId} authenticated via project '${projectName}' - web_url: ${projectCreds.web_url}`);
            return true;
        } catch (error) {
            console.error(`[MCP] Failed to resolve project '${projectName}': ${error.message}`);
            // Fall through to try other auth methods
        }
    }

    // Priority 2: Direct credentials (legacy header-based)
    if (token && website) {
        sessionCredentials.set(sessionId, {
            token,
            tokenHash: tokenHash || null,
            website,
            web_url: web_url || `https://${website}`,
            api_web_url: web_url || `https://${website}`,
            forge_host: null,
            profileName: 'http-session',
            role: 'developer',
        });
        console.log(`[MCP] Session ${sessionId} authenticated - website: ${website}`);
        return true;
    } else if (userToken) {
        setSessionUserToken(userToken, sessionId);
        console.log(`[MCP] Session ${sessionId} with userToken for auto-login`);
        return true;
    }
    return false;
};

/**
 * Extract credentials from request (headers or query params)
 */
const extractCredentialsFromRequest = (req) => {
    const url = new URL(req.url ?? "/", `http://${req.headers.host}`);
    return {
        projectName: url.searchParams.get('project') || req.headers['x-project-name'],
        token: url.searchParams.get('token') || req.headers['x-acai-token'],
        tokenHash: url.searchParams.get('tokenHash') || req.headers['x-acai-token-hash'],
        website: url.searchParams.get('website') || req.headers['x-acai-website'],
        userToken: url.searchParams.get('userToken') || req.headers['x-user-token']
    };
};

/**
 * Create and start the MCP HTTP server with both Streamable HTTP and SSE transports
 */
export function startHttpServer() {
    const app = express();

    // Parse JSON and URL-encoded bodies
    app.use(express.json());
    app.use(express.urlencoded({ extended: true }));

    // Configure CORS
    app.use(cors({
        origin: '*',
        methods: ['GET', 'POST', 'DELETE', 'OPTIONS'],
        allowedHeaders: ['Content-Type', 'X-Acai-Token', 'X-Acai-Website', 'X-Acai-Token-Hash', 'X-User-Token', 'X-Project-Name', 'Authorization', 'Mcp-Session-Id'],
        exposedHeaders: ['Mcp-Session-Id'],
        credentials: true
    }));

    //=============================================================================
    // STREAMABLE HTTP TRANSPORT (PROTOCOL VERSION 2025-03-26)
    // This is the new recommended transport for MCP
    //=============================================================================
    app.all('/mcp', async (req, res) => {
        console.log(`[MCP Streamable] ${req.method} /mcp`);

        try {
            const mcpSessionId = req.headers['mcp-session-id'];
            let transport;

            if (mcpSessionId && activeSessions.has(mcpSessionId)) {
                // Reuse existing transport
                const session = activeSessions.get(mcpSessionId);
                if (session.type === 'streamable') {
                    transport = session.transport;
                    console.log(`[MCP Streamable] Reusing session ${mcpSessionId.substring(0, 8)}...`);
                } else {
                    res.status(400).json({
                        jsonrpc: '2.0',
                        error: {
                            code: -32000,
                            message: 'Bad Request: Session exists but uses a different transport protocol'
                        },
                        id: null
                    });
                    return;
                }
            } else if (!mcpSessionId && req.method === 'POST' && isInitializeRequest(req.body)) {
                // New initialization request - create new transport
                console.log(`[MCP Streamable] New initialization request`);

                // Extract credentials from request
                const credentials = extractCredentialsFromRequest(req);

                transport = new StreamableHTTPServerTransport({
                    sessionIdGenerator: () => crypto.randomUUID(),
                    onsessioninitialized: (sessionId) => {
                        console.log(`[MCP Streamable] Session initialized: ${sessionId.substring(0, 8)}...`);

                        // Store the transport
                        activeSessions.set(sessionId, {
                            transport,
                            server: null, // Will be set after connect
                            type: 'streamable',
                            startTime: Date.now()
                        });

                        // Configure credentials for this session (async, fire-and-forget)
                        configureSessionCredentials(sessionId, credentials).then((configured) => {
                            if (configured) {
                                // Also store credentials by MCP-Session-Id for persistence
                                const creds = sessionCredentials.get(sessionId);
                                if (creds) {
                                    setMcpSessionCredentials(sessionId, creds);
                                }
                            }
                        }).catch((err) => {
                            console.error(`[MCP Streamable] Error configuring credentials for session ${sessionId}:`, err.message);
                        });
                    },
                    onsessionclosed: (sessionId) => {
                        console.log(`[MCP Streamable] Session closed: ${sessionId.substring(0, 8)}...`);
                        activeSessions.delete(sessionId);
                        clearSessionCredentials(sessionId);
                    }
                });

                // Set up onclose handler
                transport.onclose = () => {
                    const sid = transport.sessionId;
                    if (sid && activeSessions.has(sid)) {
                        console.log(`[MCP Streamable] Transport closed for session ${sid.substring(0, 8)}...`);
                        activeSessions.delete(sid);
                        clearSessionCredentials(sid);
                    }
                };

                // Create session-specific server and connect
                const sessionServer = createSessionServer();
                await sessionServer.connect(transport);

                // Store session server for role filtering
                if (transport.sessionId) {
                    sessionServers.set(transport.sessionId, sessionServer);
                }

                // Update session with server reference
                if (transport.sessionId) {
                    const session = activeSessions.get(transport.sessionId);
                    if (session) {
                        session.server = sessionServer;
                    }
                }
            } else if (mcpSessionId) {
                // Session ID provided but session not found - try to recover credentials
                const savedCreds = getMcpSessionCredentials(mcpSessionId);
                if (savedCreds) {
                    console.log(`[MCP Streamable] Recovering credentials for session ${mcpSessionId.substring(0, 8)}...`);
                    // Session might have been lost due to server restart - need to re-initialize
                }

                res.status(400).json({
                    jsonrpc: '2.0',
                    error: {
                        code: -32000,
                        message: 'Bad Request: Session not found. Please reinitialize.'
                    },
                    id: null
                });
                return;
            } else {
                res.status(400).json({
                    jsonrpc: '2.0',
                    error: {
                        code: -32000,
                        message: 'Bad Request: No valid session ID provided and not an initialization request'
                    },
                    id: null
                });
                return;
            }

            // Handle the request with the transport
            await transport.handleRequest(req, res, req.body);
        } catch (error) {
            console.error('[MCP Streamable] Error:', error);
            if (!res.headersSent) {
                res.status(500).json({
                    jsonrpc: '2.0',
                    error: {
                        code: -32603,
                        message: 'Internal server error'
                    },
                    id: null
                });
            }
        }
    });

    //=============================================================================
    // DEPRECATED HTTP+SSE TRANSPORT (PROTOCOL VERSION 2024-11-05)
    // Kept for backwards compatibility with older clients
    //=============================================================================

    // SSE connection endpoint (GET /sse)
    app.get('/sse', async (req, res) => {
        console.log(`[MCP SSE] New SSE connection`);

        const credentials = extractCredentialsFromRequest(req);

        // Create SSE transport
        const transport = new SSEServerTransport("/message", res);
        const sessionId = transport.sessionId;

        console.log(`[MCP SSE] Session created: ${sessionId}`);

        // Create session-specific server
        const sessionServer = createSessionServer();

        // Store session server for role filtering
        sessionServers.set(sessionId, sessionServer);

        // Set up heartbeat
        const heartbeatInterval = setInterval(() => {
            try {
                if (res.writableEnded || res.destroyed) {
                    clearInterval(heartbeatInterval);
                    return;
                }
                res.write(': heartbeat\n\n');
            } catch (err) {
                console.error(`[MCP SSE] Heartbeat error for session ${sessionId}:`, err.message);
                clearInterval(heartbeatInterval);
            }
        }, 30000);

        // Configure credentials
        await configureSessionCredentials(sessionId, credentials);

        // Store session
        activeSessions.set(sessionId, {
            transport,
            server: sessionServer,
            type: 'sse',
            heartbeatInterval,
            startTime: Date.now()
        });

        // Handle close
        res.on('close', () => {
            console.log(`[MCP SSE] Connection closed for session ${sessionId}`);
        });

        transport.onclose = () => {
            const session = activeSessions.get(sessionId);
            if (!session) return;

            activeSessions.delete(sessionId);
            if (session.heartbeatInterval) {
                clearInterval(session.heartbeatInterval);
            }
            clearSessionCredentials(sessionId);
            console.log(`[MCP SSE] Session ${sessionId} cleaned up`);
        };

        try {
            await sessionServer.connect(transport);
            console.log(`[MCP SSE] Session ${sessionId} connected`);
        } catch (error) {
            console.error(`[MCP SSE] Connection error for session ${sessionId}:`, error.message);
            activeSessions.delete(sessionId);
            clearSessionCredentials(sessionId);
        }
    });

    // Message endpoint for SSE transport (POST /message)
    app.post('/message', async (req, res) => {
        const url = new URL(req.url ?? "/", `http://${req.headers.host}`);
        const sessionId = url.searchParams.get("sessionId");

        if (!sessionId) {
            res.status(400).json({ error: "Missing session ID" });
            return;
        }

        const session = activeSessions.get(sessionId);
        if (!session || session.type !== 'sse') {
            res.status(404).json({ error: "Session not found" });
            return;
        }

        const { transport, heartbeatInterval } = session;

        // Check if SSE connection is still alive
        const sseResponse = transport._sseResponse;
        if (!sseResponse || sseResponse.writableEnded || sseResponse.destroyed) {
            activeSessions.delete(sessionId);
            clearSessionCredentials(sessionId);
            if (heartbeatInterval) {
                clearInterval(heartbeatInterval);
            }
            res.status(410).json({ error: "SSE connection closed", code: "SSE_CLOSED" });
            return;
        }

        try {
            await transport.handlePostMessage(req, res, req.body);
        } catch (error) {
            console.error(`[MCP SSE] POST error for session ${sessionId}:`, error.message);
            if (!res.headersSent) {
                res.status(500).json({ error: error.message });
            }
        }
    });

    // Root path normalization (for clients that call "/" instead of /sse or /mcp)
    app.get('/', (req, res) => {
        // Redirect to SSE for backwards compatibility
        res.redirect('/sse');
    });

    app.post('/', async (req, res) => {
        // Check if it's a Streamable HTTP initialize request
        if (isInitializeRequest(req.body)) {
            // Forward to /mcp
            req.url = '/mcp';
            app.handle(req, res);
        } else {
            // Forward to /message (needs sessionId in query)
            req.url = '/message';
            app.handle(req, res);
        }
    });

    //=============================================================================
    // HEALTH CHECK
    //=============================================================================
    app.get('/health', (req, res) => {
        const sseCount = Array.from(activeSessions.values()).filter(s => s.type === 'sse').length;
        const streamableCount = Array.from(activeSessions.values()).filter(s => s.type === 'streamable').length;

        res.json({
            status: "ok",
            activeSessions: activeSessions.size,
            sse: sseCount,
            streamable: streamableCount,
            mode: "hybrid"
        });
    });

    //=============================================================================
    // OAUTH2 ENDPOINTS
    //=============================================================================

    // OAuth2 Authorization Server Metadata endpoint (per RFC8414)
    app.get('/.well-known/oauth-authorization-server', (req, res) => {
        const baseUrl = `https://${req.headers.host}`;
        res.json({
            issuer: baseUrl,
            authorization_endpoint: `${baseUrl}/authorize`,
            token_endpoint: `${baseUrl}/token`,
            registration_endpoint: `${baseUrl}/register`,
            grant_types_supported: ["authorization_code", "client_credentials"],
            response_types_supported: ["code"],
            token_endpoint_auth_methods_supported: ["client_secret_post", "none"],
            service_documentation: `${baseUrl}/docs`,
            code_challenge_methods_supported: ["S256"],
            mcp_endpoint: `${baseUrl}/mcp`
        });
    });

    // OAuth2 Dynamic Client Registration endpoint (per RFC7591)
    app.post('/register', (req, res) => {
        const clientInfo = req.body;
        console.error(`[MCP HTTP] POST /register - Client registration request:`, clientInfo);

        const clientId = `acai-${Date.now()}-${Math.random().toString(36).substring(7)}`;

        res.status(201).json({
            client_id: clientId,
            client_id_issued_at: Math.floor(Date.now() / 1000),
            redirect_uris: clientInfo.redirect_uris || [],
            token_endpoint_auth_method: "none",
            grant_types: ["authorization_code"],
            response_types: ["code"]
        });

        console.error(`[MCP HTTP] POST /register - Registered client: ${clientId}`);
    });

    // OAuth2 Authorization endpoint
    app.get('/authorize', (req, res) => {
        const { client_id: clientId, redirect_uri: redirectUri, state, code_challenge: codeChallenge, code_challenge_method: codeChallengeMethod = "S256" } = req.query;

        console.error(`[MCP HTTP] GET /authorize - client_id: ${clientId}, redirect_uri: ${redirectUri}`);

        if (!clientId || !redirectUri) {
            res.status(400).json({
                error: "invalid_request",
                error_description: "Missing client_id or redirect_uri"
            });
            return;
        }

        const code = Buffer.from(JSON.stringify({
            client_id: clientId,
            code_challenge: codeChallenge,
            code_challenge_method: codeChallengeMethod,
            timestamp: Date.now(),
            expires_at: Date.now() + (10 * 60 * 1000)
        })).toString('base64');

        const callback = new URL(redirectUri);
        callback.searchParams.set("code", code);
        if (state) callback.searchParams.set("state", state);

        console.error(`[MCP HTTP] GET /authorize - Redirecting to: ${callback.toString()}`);

        res.redirect(302, callback.toString());
    });

    // OAuth2 Token endpoint
    app.post('/token', async (req, res) => {
        console.error(`[MCP HTTP] POST /token - Received token request`);

        try {
            let params = req.body;

            // Handle form-encoded body
            const contentType = req.headers['content-type'] || '';
            if (contentType.includes('application/x-www-form-urlencoded') && typeof req.body === 'string') {
                const searchParams = new URLSearchParams(req.body);
                params = Object.fromEntries(searchParams);
            }

            const { grant_type, client_secret, code, code_verifier } = params;

            console.error(`[MCP HTTP] POST /token - grant_type: ${grant_type}, has_code: ${!!code}, has_client_secret: ${!!client_secret}`);

            if (grant_type === 'authorization_code') {
                if (!code) {
                    res.status(400).json({
                        error: "invalid_request",
                        error_description: "Missing authorization code"
                    });
                    return;
                }

                try {
                    const decodedCode = JSON.parse(Buffer.from(code, 'base64').toString());

                    if (decodedCode.expires_at && Date.now() > decodedCode.expires_at) {
                        res.status(400).json({
                            error: "invalid_grant",
                            error_description: "Authorization code has expired"
                        });
                        return;
                    }

                    if (decodedCode.code_challenge && code_verifier) {
                        const digest = crypto.createHash('sha256').update(code_verifier).digest('base64url');
                        if (digest !== decodedCode.code_challenge) {
                            res.status(400).json({
                                error: "invalid_grant",
                                error_description: "PKCE code verifier mismatch"
                            });
                            return;
                        }
                    }

                    // client_secret = project name to resolve via local server
                    if (!client_secret) {
                        res.status(400).json({
                            error: "invalid_request",
                            error_description: "Missing client_secret (should be project name)"
                        });
                        return;
                    }

                    const projectCreds = await resolveProjectCredentials(client_secret);
                    const accessToken = signJwt({
                        acaiToken: projectCreds.token,
                        acaiTokenHash: projectCreds.tokenHash,
                        website: projectCreds.website,
                        web_url: projectCreds.web_url,
                        clientId: decodedCode.client_id,
                        tokenType: "acai-credentials",
                    });

                    res.setHeader("Cache-Control", "no-store");
                    res.setHeader("Pragma", "no-cache");
                    res.json({
                        access_token: accessToken,
                        token_type: "Bearer",
                        expires_in: ACCESS_TOKEN_TTL
                    });
                } catch (error) {
                    console.error(`[MCP HTTP] POST /token (auth_code) - ERROR:`, error.message);
                    res.status(400).json({
                        error: "invalid_grant",
                        error_description: "Invalid authorization code"
                    });
                }
                return;
            }

            if (grant_type !== 'client_credentials') {
                res.status(400).json({
                    error: "unsupported_grant_type",
                    error_description: "Only 'client_credentials' and 'authorization_code' grant types are supported"
                });
                return;
            }

            // client_secret = project name to resolve via local server
            if (!client_secret) {
                res.status(400).json({
                    error: "invalid_request",
                    error_description: "Missing client_secret (should be project name)"
                });
                return;
            }

            const projectCreds = await resolveProjectCredentials(client_secret);
            const accessToken = signJwt({
                acaiToken: projectCreds.token,
                acaiTokenHash: projectCreds.tokenHash,
                website: projectCreds.website,
                web_url: projectCreds.web_url,
                clientId: params.client_id || "client_credentials",
                tokenType: "acai-credentials",
            });

            res.setHeader("Cache-Control", "no-store");
            res.setHeader("Pragma", "no-cache");
            res.json({
                access_token: accessToken,
                token_type: "Bearer",
                expires_in: ACCESS_TOKEN_TTL
            });
        } catch (error) {
            console.error(`[MCP HTTP] POST /token - ERROR:`, error.message);
            const isAuthError = error.message?.toLowerCase().includes("exchange");
            res.status(isAuthError ? 400 : 500).json({
                error: isAuthError ? "invalid_grant" : "server_error",
                error_description: error.message
            });
        }
    });

    //=============================================================================
    // STATIC FILE SERVING
    //=============================================================================

    // Serve Figma images
    app.get('/figma-images/:fileName', (req, res) => {
        const { fileName } = req.params;
        if (!fileName) {
            res.status(400).json({ error: "File name required" });
            return;
        }

        try {
            const figmaDir = process.env.FIGMA_IMAGES_DIR || '/app/figma_images';
            const filePath = path.join(figmaDir, fileName);

            const resolvedPath = path.resolve(filePath);
            const resolvedFigmaDir = path.resolve(figmaDir);
            if (!resolvedPath.startsWith(resolvedFigmaDir)) {
                res.status(403).json({ error: "Access denied" });
                return;
            }

            if (!fs.existsSync(filePath)) {
                res.status(404).json({ error: "File not found", path: fileName });
                return;
            }

            const ext = path.extname(fileName).toLowerCase();
            const contentTypes = {
                '.png': 'image/png',
                '.jpg': 'image/jpeg',
                '.jpeg': 'image/jpeg',
                '.gif': 'image/gif',
                '.webp': 'image/webp',
                '.svg': 'image/svg+xml'
            };
            const contentType = contentTypes[ext] || 'application/octet-stream';

            const fileBuffer = fs.readFileSync(filePath);
            res.setHeader("Content-Type", contentType);
            res.setHeader("Content-Length", fileBuffer.length);
            res.setHeader("Cache-Control", "public, max-age=3600");
            res.send(fileBuffer);
        } catch (error) {
            res.status(500).json({ error: error.message });
        }
    });

    // Serve generated images
    app.get('/generated-images/:fileId', (req, res) => {
        const { fileId } = req.params;
        if (!fileId) {
            res.status(400).json({ error: "File ID required" });
            return;
        }

        try {
            const tempDir = process.env.GENERATED_IMAGES_TEMP_DIR || path.join(os.tmpdir(), 'generated-images');
            const filename = `generated-${fileId}.jpg`;
            const filePath = path.join(tempDir, filename);

            const resolvedPath = path.resolve(filePath);
            const resolvedTempDir = path.resolve(tempDir);
            if (!resolvedPath.startsWith(resolvedTempDir)) {
                res.status(403).json({ error: "Access denied" });
                return;
            }

            if (!fs.existsSync(filePath)) {
                res.status(404).json({ error: "File not found" });
                return;
            }

            const fileBuffer = fs.readFileSync(filePath);
            res.setHeader("Content-Type", "image/jpeg");
            res.setHeader("Content-Length", fileBuffer.length);
            res.setHeader("Cache-Control", "public, max-age=3600");
            res.send(fileBuffer);
        } catch (error) {
            res.status(500).json({ error: error.message });
        }
    });

    //=============================================================================
    // START SERVER
    //=============================================================================
    const server = app.listen(MCP_PORT, '0.0.0.0', () => {
        console.error(`[MCP] Server listening on http://0.0.0.0:${MCP_PORT}`);
        console.error(`[MCP] Streamable HTTP endpoint: /mcp (recommended)`);
        console.error(`[MCP] Legacy SSE endpoint: /sse (backwards compatible)`);
        console.error(`[MCP] Provide credentials via headers: X-Acai-Token, X-Acai-Website, X-Acai-Token-Hash`);
    });

    server.on("error", (error) => {
        console.error(`[MCP] Server error:`, error);
    });

    // Handle shutdown
    process.on('SIGINT', async () => {
        console.log('[MCP] Shutting down server...');
        for (const [sessionId, session] of activeSessions.entries()) {
            try {
                console.log(`[MCP] Closing session ${sessionId}`);
                if (session.transport.close) {
                    await session.transport.close();
                }
                if (session.heartbeatInterval) {
                    clearInterval(session.heartbeatInterval);
                }
            } catch (error) {
                console.error(`[MCP] Error closing session ${sessionId}:`, error);
            }
        }
        activeSessions.clear();
        console.log('[MCP] Shutdown complete');
        process.exit(0);
    });

    return server;
}

export { activeSessions };