import { Elysia, t } from 'elysia';
import { randomUUID } from 'node:crypto';
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
import { isInitializeRequest } from '@modelcontextprotocol/sdk/types.js';

// Map untuk menyimpan transport berdasarkan sessionId
const transports: Record<string, StreamableHTTPServerTransport> = {};

const McpRoute = new Elysia()
    .post(
        '/mcp',
        async ({ request, body, set }) => {
            const sessionId = request.headers.get('mcp-session-id') ?? undefined;
            let transport: StreamableHTTPServerTransport;

            // Reuse existing session jika ada
            if (sessionId && transports[sessionId]) {
                transport = transports[sessionId];
            }
            // Jika ini permintaan inisialisasi MCP baru
            else if (!sessionId && isInitializeRequest(body)) {
                transport = new StreamableHTTPServerTransport({
                    sessionIdGenerator: () => randomUUID(),
                    onsessioninitialized: (sid) => {
                        transports[sid] = transport;
                        console.log(`🟢 Session initialized: ${sid}`);
                    },
                });

                transport.onclose = () => {
                    if (transport.sessionId) {
                        console.log(`🔴 Session closed: ${transport.sessionId}`);
                        delete transports[transport.sessionId];
                    }
                };

                // Buat instance MCP server
                const server = new McpServer({
                    name: 'elysia-mcp-server',
                    version: '1.0.0',
                });

                // Contoh: tambahkan dummy tool/resource di sini jika mau
                // server.addTool('ping', async () => 'pong');

                await server.connect(transport);

                // Tunggu hingga session ID terbentuk
                await new Promise<void>((resolve) => {
                    const wait = () => {
                        if (transport.sessionId) resolve();
                        else setTimeout(wait, 5);
                    };
                    wait();
                });

                // Kirim sessionId ke client
                set.headers['mcp-session-id'] = transport.sessionId!;
                set.status = 200;
                return { sessionId: transport.sessionId };
            }
            // Jika tidak valid
            else {
                set.status = 400;
                return {
                    jsonrpc: '2.0',
                    error: {
                        code: -32000,
                        message: 'Bad Request: No valid session ID provided',
                    },
                    id: null,
                };
            }

            // ✅ Gunakan interface Web (Bun/Elysia) langsung
            const webTransport = transport as any;
            if (typeof webTransport.handleRequestWeb === 'function') {
                // handleRequestWeb() adalah versi WebAPI (Request/Response)
                return await webTransport.handleRequestWeb(request);
            } else {
                // fallback – manual handle body
                return new Response(JSON.stringify({ ok: true }), {
                    headers: { 'Content-Type': 'application/json' },
                });
            }
        },
        { body: t.Any() }
    )
    // Server-sent events (SSE)
    .get('/mcp', async ({ request, set }) => {
        const sessionId = request.headers.get('mcp-session-id') ?? undefined;
        const transport = sessionId ? transports[sessionId] : undefined;

        if (!transport) {
            set.status = 400;
            return 'Invalid or missing session ID';
        }

        const webTransport = transport as any;
        if (typeof webTransport.handleRequestWeb === 'function') {
            return await webTransport.handleRequestWeb(request);
        }

        set.status = 200;
        return new Response('SSE not supported by this transport');
    })
    // Session cleanup
    .delete('/mcp', async ({ request, set }) => {
        const sessionId = request.headers.get('mcp-session-id') ?? undefined;
        const transport = sessionId ? transports[sessionId] : undefined;

        if (!transport) {
            set.status = 400;
            return 'Invalid or missing session ID';
        }

        const webTransport = transport as any;
        if (typeof webTransport.handleRequestWeb === 'function') {
            return await webTransport.handleRequestWeb(request);
        }

        set.status = 200;
        return new Response('Session deleted');
    })

export default McpRoute