import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
import { SSEServerTransport } from '@modelcontextprotocol/sdk/server/sse.js';
import express from 'express';
import crypto from 'crypto';
import { z } from 'zod';
import { QdrantMemoryService } from './qdrant.js';

async function main() {
    const server = new McpServer({
        name: '@mintel/memory-mcp',
        version: '1.0.0',
    });

    const qdrantService = new QdrantMemoryService(process.env.QDRANT_URL || 'http://localhost:6333');

    server.tool(
        'store_memory',
        'Store a new piece of knowledge/memory into the vector database. Use this to remember architectural decisions, preferences, aliases, etc.',
        {
            label: z.string().describe('A short, descriptive label or title for the memory (e.g., "Architektur-Entscheidungen")'),
            content: z.string().describe('The actual content to remember (e.g., "In diesem Projekt nutzen wir lieber Composition over Inheritance.")'),
        },
        async (args) => {
            const success = await qdrantService.storeMemory(args.label, args.content);
            if (success) {
                return {
                    content: [{ type: 'text', text: `Successfully stored memory: [${args.label}]` }],
                };
            } else {
                return {
                    content: [{ type: 'text', text: `Failed to store memory: [${args.label}]` }],
                    isError: true,
                };
            }
        }
    );

    server.tool(
        'retrieve_memory',
        'Retrieve relevant memories from the vector database based on a semantic search query.',
        {
            query: z.string().describe('The search query to find relevant memories.'),
            limit: z.number().optional().describe('Maximum number of results to return (default: 5)'),
        },
        async (args) => {
            const results = await qdrantService.retrieveMemory(args.query, args.limit || 5);

            if (results.length === 0) {
                return {
                    content: [{ type: 'text', text: 'No relevant memories found.' }],
                };
            }

            const formattedResults = results
                .map(r => `- [${r.label}] (Score: ${r.score.toFixed(3)}): ${r.content}`)
                .join('\n');

            return {
                content: [{ type: 'text', text: `Found ${results.length} memories:\n\n${formattedResults}` }],
            };
        }
    );

    const isStdio = process.argv.includes('--stdio');

    if (isStdio) {
        // Connect Stdio FIRST to avoid handshake timeouts while loading model
        const { StdioServerTransport } = await import('@modelcontextprotocol/sdk/server/stdio.js');
        const transport = new StdioServerTransport();
        await server.connect(transport);
        console.error('Memory MCP server is running on stdio');

        // Initialize dependency after connection
        try {
            await qdrantService.initialize();
        } catch (e) {
            console.error('Failed to initialize local dependencies:', e);
        }
    } else {
        const app = express();
        const transports = new Map<string, SSEServerTransport>();

        // Manual CORS middleware
        app.use((req, res, next) => {
            res.setHeader('Access-Control-Allow-Origin', '*');
            res.setHeader('Access-Control-Allow-Methods', 'GET, POST, OPTIONS');
            res.setHeader('Access-Control-Allow-Headers', 'Content-Type, Authorization, x-mcp-protocol-version');
            if (req.method === 'OPTIONS') {
                res.status(204).end();
                return;
            }
            next();
        });

        app.use((req, _res, next) => {
            console.error(`${req.method} ${req.url}`);
            next();
        });

        app.get('/sse', async (req, res) => {
            const transport = new SSEServerTransport('/message', res as any);
            const sessionId = transport.sessionId;
            console.error(`New SSE connection: ${sessionId}`);
            transports.set(sessionId, transport);

            const heartbeatInterval = setInterval(() => {
                res.write(": heartbeat\n\n");
            }, 15000);

            req.on('close', () => {
                console.error(`SSE connection closed: ${sessionId}`);
                clearInterval(heartbeatInterval);
                transports.delete(sessionId);
            });

            await server.connect(transport);
            
            await new Promise((resolve) => {
                req.on('close', resolve);
            });
        });

        app.post('/message', express.json(), async (req, res) => {
            const sessionId = req.query.sessionId as string;
            const transport = transports.get(sessionId);

            console.error(`POST /message/${sessionId} - Active transports: ${transports.size}`);

            if (!transport) {
                console.error(`No transport found for session: ${sessionId}`);
                res.status(400).send('No active SSE connection for this session');
                return;
            }
            await transport.handlePostMessage(req, res);
        });

        const PORT = Number(process.env.MEMORY_MCP_PORT) || 3002;
        app.listen(PORT, '0.0.0.0', async () => {
            console.error(`Memory MCP server running on http://0.0.0.0:${PORT}/sse`);
            // Initialize dependencies in SSE mode on startup
            try {
                await qdrantService.initialize();
            } catch (e) {
                console.error('Failed to initialize local dependencies:', e);
            }
        });
    }
}

main().catch((error) => {
    console.error('Fatal error:', error);
    process.exit(1);
});