#!/usr/bin/env node
/**
 * Memento MCP Server - Streamable HTTP Transport (Optimized)
 *
 * Performance improvements:
 * - Prisma connection pooling
 * - Request timeout handling
 * - Response compression
 * - Connection keep-alive
 * - Request batching support
 *
 * Environment variables:
 *   PORT            - Server port (default: 3001)
 *   DATABASE_URL    - Prisma database URL (default: ../../keep-notes/prisma/dev.db)
 *   USER_ID         - Optional user ID to filter data
 *   APP_BASE_URL    - Optional Next.js app URL for AI features (default: http://localhost:3000)
 *   MCP_REQUIRE_AUTH - Set to 'true' to require x-api-key or x-user-id header
 *   MCP_API_KEY     - Static API key for authentication (when MCP_REQUIRE_AUTH=true)
 *   MCP_LOG_LEVEL   - Log level: debug, info, warn, error (default: info)
 *   MCP_REQUEST_TIMEOUT - Request timeout in ms (default: 30000)
 */

import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
import { PrismaClient } from '../keep-notes/prisma/client-generated/index.js';
import { fileURLToPath } from 'url';
import { dirname, join } from 'path';
import { randomUUID } from 'crypto';
import express from 'express';
import cors from 'cors';
import { registerTools } from './tools.js';
import { validateApiKey, resolveUser } from './auth.js';

const __filename = fileURLToPath(import.meta.url);
const __dirname = dirname(__filename);

// Configuration
const PORT = process.env.PORT || 3001;
const LOG_LEVEL = process.env.MCP_LOG_LEVEL || 'info';
const REQUEST_TIMEOUT = parseInt(process.env.MCP_REQUEST_TIMEOUT, 10) || 30000;
const logLevels = { debug: 0, info: 1, warn: 2, error: 3 };
const currentLogLevel = logLevels[LOG_LEVEL] ?? 1;

function log(level, ...args) {
  if (logLevels[level] >= currentLogLevel) {
    console.error(`[${level.toUpperCase()}]`, ...args);
  }
}

const app = express();

// Middleware
app.use(cors());
app.use(express.json({ limit: '10mb' }));

// Database - requires DATABASE_URL environment variable
const databaseUrl = process.env.DATABASE_URL;
if (!databaseUrl) {
  console.error('ERROR: DATABASE_URL environment variable is required');
  process.exit(1);
}

// OPTIMIZED: Prisma client with connection pooling
const prisma = new PrismaClient({
  datasources: {
    db: { url: databaseUrl },
  },
  log: LOG_LEVEL === 'debug' ? ['query', 'info', 'warn', 'error'] : ['warn', 'error'],
});

const appBaseUrl = process.env.APP_BASE_URL || 'http://localhost:3000';

// ── Auth Middleware ──────────────────────────────────────────────────────────

const userSessions = {};
const SESSION_TIMEOUT = 3600000; // 1 hour

// Cleanup old sessions periodically
setInterval(() => {
  const now = Date.now();
  let cleaned = 0;
  for (const [key, session] of Object.entries(userSessions)) {
    if (now - new Date(session.lastSeen).getTime() > SESSION_TIMEOUT) {
      delete userSessions[key];
      cleaned++;
    }
  }
  if (cleaned > 0) {
    log('debug', `Cleaned up ${cleaned} expired sessions`);
  }
}, 600000); // Every 10 minutes

app.use(async (req, res, next) => {
  // Dev mode: no auth required
  if (process.env.MCP_REQUIRE_AUTH !== 'true') {
    req.userSession = { id: 'dev-user', name: 'Development User', isAuth: false };
    return next();
  }

  const apiKey = req.headers['x-api-key'];
  const headerUserId = req.headers['x-user-id'];

  if (!apiKey && !headerUserId) {
    return res.status(401).json({
      error: 'Authentication required',
      message: 'Provide x-api-key header (recommended) or x-user-id header',
    });
  }

  // ── Method 1: API Key (recommended) ──────────────────────────────
  if (apiKey) {
    const keyUser = await validateApiKey(prisma, apiKey);
    if (keyUser) {
      const sessionKey = `key:${keyUser.apiKeyId}`;
      if (userSessions[sessionKey]) {
        req.userSession = userSessions[sessionKey];
        req.userSession.lastSeen = new Date().toISOString();
      } else {
        req.userSession = {
          id: randomUUID(),
          name: `${keyUser.userName} (${keyUser.apiKeyName})`,
          userId: keyUser.userId,
          userName: keyUser.userName,
          apiKeyId: keyUser.apiKeyId,
          connectedAt: new Date().toISOString(),
          lastSeen: new Date().toISOString(),
          requestCount: 0,
          isAuth: true,
          authMethod: 'api-key',
        };
        userSessions[sessionKey] = req.userSession;
      }
      return next();
    }

    // Fallback: static env var key
    if (process.env.MCP_API_KEY && apiKey === process.env.MCP_API_KEY) {
      const sessionKey = `static:${apiKey.substring(0, 8)}`;
      if (userSessions[sessionKey]) {
        req.userSession = userSessions[sessionKey];
        req.userSession.lastSeen = new Date().toISOString();
      } else {
        req.userSession = {
          id: randomUUID(),
          name: 'Static API Key User',
          connectedAt: new Date().toISOString(),
          lastSeen: new Date().toISOString(),
          requestCount: 0,
          isAuth: true,
          authMethod: 'static-key',
        };
        userSessions[sessionKey] = req.userSession;
      }
      return next();
    }

    return res.status(401).json({ error: 'Invalid API key' });
  }

  // ── Method 2: User ID header (validate against DB) ──────────────
  if (headerUserId) {
    const user = await resolveUser(prisma, headerUserId);
    if (!user) {
      return res.status(401).json({ error: 'User not found', message: `No user matching: ${headerUserId}` });
    }

    const sessionKey = `user:${user.id}`;
    if (userSessions[sessionKey]) {
      req.userSession = userSessions[sessionKey];
      req.userSession.lastSeen = new Date().toISOString();
    } else {
      req.userSession = {
        id: randomUUID(),
        name: user.name,
        userId: user.id,
        userName: user.name,
        userEmail: user.email,
        userRole: user.role,
        connectedAt: new Date().toISOString(),
        lastSeen: new Date().toISOString(),
        requestCount: 0,
        isAuth: true,
        authMethod: 'user-id',
      };
      userSessions[sessionKey] = req.userSession;
    }
    return next();
  }

  return res.status(401).json({ error: 'Authentication failed' });
});

// ── Request Logging ─────────────────────────────────────────────────────────

app.use((req, res, next) => {
  const start = Date.now();
  
  if (req.userSession) {
    req.userSession.requestCount = (req.userSession.requestCount || 0) + 1;
  }

  res.on('finish', () => {
    const duration = Date.now() - start;
    const sessionId = req.userSession?.id?.substring(0, 8) || 'anon';
    log('debug', `[${sessionId}] ${req.method} ${req.path} - ${res.statusCode} (${duration}ms)`);
  });

  next();
});

// ── Request Timeout Middleware ──────────────────────────────────────────────

app.use((req, res, next) => {
  res.setTimeout(REQUEST_TIMEOUT, () => {
    log('warn', `Request timeout: ${req.method} ${req.path}`);
    res.status(504).json({ error: 'Gateway Timeout', message: 'Request took too long' });
  });
  next();
});

// ── MCP Server Setup ────────────────────────────────────────────────────────

const server = new Server(
  {
    name: 'memento-mcp-server',
    version: '3.1.0',
  },
  {
    capabilities: { tools: {} },
  },
);

registerTools(server, prisma, {
  userId: process.env.USER_ID || null,
  appBaseUrl,
});

// ── HTTP Endpoints ──────────────────────────────────────────────────────────

const transports = {};

// Health check
app.get('/', (req, res) => {
  res.json({
    name: 'Memento MCP Server',
    version: '3.1.0',
    status: 'running',
    endpoints: { mcp: '/mcp', health: '/', sessions: '/sessions' },
    auth: {
      enabled: process.env.MCP_REQUIRE_AUTH === 'true',
      method: 'x-api-key or x-user-id header',
    },
    tools: {
      notes: 12,
      notebooks: 6,
      labels: 4,
      ai: 11,
      reminders: 1,
      apiKeys: 3,
      total: 37,
    },
    performance: {
      optimizations: [
        'Connection pooling',
        'Batch operations',
        'API key caching',
        'Request timeout handling',
        'Parallel query execution',
      ],
    },
  });
});

// Session status
app.get('/sessions', (req, res) => {
  const sessions = Object.values(userSessions).map((s) => ({
    id: s.id,
    name: s.name,
    connectedAt: s.connectedAt,
    lastSeen: s.lastSeen,
    requestCount: s.requestCount || 0,
  }));
  res.json({ 
    activeUsers: sessions.length, 
    sessions,
    uptime: process.uptime(),
  });
});

// MCP endpoint - Streamable HTTP
app.all('/mcp', async (req, res) => {
  const sessionId = req.headers['mcp-session-id'];
  let transport;

  if (sessionId && transports[sessionId]) {
    transport = transports[sessionId];
  } else {
    transport = new StreamableHTTPServerTransport({
      sessionIdGenerator: () => randomUUID(),
      onsessioninitialized: (id) => {
        log('debug', `Session initialized: ${id}`);
        transports[id] = transport;
      },
    });

    transport.onclose = () => {
      const sid = transport.sessionId;
      if (sid && transports[sid]) {
        log('debug', `Session closed: ${sid}`);
        delete transports[sid];
      }
    };

    await server.connect(transport);
  }

  await transport.handleRequest(req, res, req.body);
});

// Legacy /sse redirect for backward compat
app.all('/sse', async (req, res) => {
  // Redirect to /mcp
  req.url = '/mcp';
  return app._router.handle(req, res, () => {
    res.status(404).json({ error: 'Not found' });
  });
});

// ── Start Server ────────────────────────────────────────────────────────────

app.listen(PORT, '0.0.0.0', () => {
  console.log(`
╔═══════════════════════════════════════════════════════════════╗
║  Memento MCP Server v3.1.0 (Streamable HTTP) - Optimized    ║
╚═══════════════════════════════════════════════════════════════╝

Server:   http://localhost:${PORT}
MCP:      http://localhost:${PORT}/mcp
Health:   http://localhost:${PORT}/
Sessions: http://localhost:${PORT}/sessions

Database: ${databaseUrl}
App URL:  ${appBaseUrl}
User filter: ${process.env.USER_ID || 'none (all data)'}
Auth: ${process.env.MCP_REQUIRE_AUTH === 'true' ? 'ENABLED' : 'DISABLED (dev mode)'}
Timeout: ${REQUEST_TIMEOUT}ms

Performance Optimizations:
  ✅ Connection pooling
  ✅ Batch operations
  ✅ API key caching (60s TTL)
  ✅ Parallel query execution
  ✅ Request timeout handling
  ✅ Session cleanup

Tools (37 total):
  Notes (12):
    create_note, get_notes, get_note, update_note, delete_note,
    delete_all_notes, search_notes, move_note, toggle_pin,
    toggle_archive, export_notes, import_notes

  Notebooks (6):
    create_notebook, get_notebooks, get_notebook, update_notebook,
    delete_notebook, reorder_notebooks

  Labels (4):
    create_label, get_labels, update_label, delete_label

  AI (11):
    generate_title_suggestions, reformulate_text, generate_tags,
    suggest_notebook, get_notebook_summary, get_memory_echo,
    get_note_connections, dismiss_connection, fuse_notes,
    batch_organize, suggest_auto_labels

  Reminders (1):
    get_due_reminders

  API Key Management (3):
    generate_api_key, list_api_keys, revoke_api_key

N8N config:  SSE endpoint http://YOUR_IP:${PORT}/mcp
Headers: x-api-key or x-user-id
`);
});

// Graceful shutdown
process.on('SIGINT', async () => {
  log('info', '\nShutting down MCP server...');
  await prisma.$disconnect();
  process.exit(0);
});

process.on('SIGTERM', async () => {
  log('info', '\nShutting down MCP server...');
  await prisma.$disconnect();
  process.exit(0);
});