#!/usr/bin/env node
/**
 * Memento MCP Server - Streamable HTTP Transport (Fast)
 *
 * - Prisma connection pooling
 * - Compact JSON output
 * - Bounded session cache
 * - Proper keep-alive & timeouts
 * - O(1) API key validation
 *
 * Environment:
 *   PORT              Server port (default: 3001)
 *   DATABASE_URL      Prisma database URL
 *   USER_ID           Optional user ID filter
 *   APP_BASE_URL      Next.js app URL (default: http://localhost:3000)
 *   MCP_REQUIRE_AUTH  Set 'true' to require authentication
 *   MCP_API_KEY       Static fallback API key
 *   MCP_LOG_LEVEL     debug, info, warn, error (default: info)
 *   MCP_REQUEST_TIMEOUT  Timeout in ms (default: 30000)
 */

import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
import { PrismaClient } from '@prisma/client';
import { randomUUID } from 'crypto';
import express from 'express';
import cors from 'cors';
import { registerTools } from './tools.js';
import { validateApiKey, resolveUser } from './auth.js';
import { requestContext } from './request-context.js';

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 MAX_SESSIONS = 500;
const SESSION_TTL = 3600000;

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 databaseUrl = process.env.DATABASE_URL;
if (!databaseUrl) {
  console.error('ERROR: DATABASE_URL is required');
  process.exit(1);
}

const isPostgres = databaseUrl.startsWith('postgresql://') || databaseUrl.startsWith('postgres://');

const prisma = new PrismaClient({
  datasources: { db: { url: databaseUrl } },
  ...(isPostgres ? { datasources: { db: { url: `${databaseUrl}${databaseUrl.includes('?') ? '&' : '?'}connection_limit=10&pool_timeout=10` } } } : {}),
  log: LOG_LEVEL === 'debug' ? ['query', 'info', 'warn', 'error'] : ['warn', 'error'],
});

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

// ── Bounded Session Cache ───────────────────────────────────────────────────

const sessions = new Map();

function cleanupSessions() {
  const now = Date.now();
  let cleaned = 0;
  for (const [key, s] of sessions) {
    if (now - s._lastSeen > SESSION_TTL) {
      sessions.delete(key);
      cleaned++;
    }
  }
  if (cleaned > 0) log('debug', `Cleaned ${cleaned} sessions`);
}

function pruneIfFull() {
  if (sessions.size < MAX_SESSIONS) return;
  const entries = [...sessions.entries()].sort((a, b) => a[1]._lastSeen - b[1]._lastSeen);
  for (let i = 0; i < Math.floor(MAX_SESSIONS / 4); i++) {
    sessions.delete(entries[i][0]);
  }
}

setInterval(cleanupSessions, 600000);

// ── Express ─────────────────────────────────────────────────────────────────

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

// ── Health (before auth middleware — used by Docker healthcheck) ────────────

app.get('/health', (req, res) => res.json({ ok: true, uptime: process.uptime() }));

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

app.use(async (req, res, next) => {
  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 or x-user-id header' });
  }

  if (apiKey) {
    const keyUser = await validateApiKey(prisma, apiKey);
    if (keyUser) {
      req.userSession = getOrCreateSession(`key:${keyUser.apiKeyId}`, {
        name: `${keyUser.userName} (${keyUser.apiKeyName})`,
        userId: keyUser.userId,
        userName: keyUser.userName,
        apiKeyId: keyUser.apiKeyId,
        authMethod: 'api-key',
      });
      return next();
    }

    if (process.env.MCP_API_KEY && apiKey === process.env.MCP_API_KEY) {
      req.userSession = getOrCreateSession(`static:${apiKey.substring(0, 8)}`, {
        name: 'Static API Key User',
        userId: process.env.USER_ID || null,
        authMethod: 'static-key',
      });
      return next();
    }

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

  if (headerUserId) {
    const user = await resolveUser(prisma, headerUserId);
    if (!user) {
      return res.status(401).json({ error: 'User not found' });
    }
    req.userSession = getOrCreateSession(`user:${user.id}`, {
      name: user.name,
      userId: user.id,
      userName: user.name,
      userEmail: user.email,
      userRole: user.role,
      authMethod: 'user-id',
    });
    return next();
  }

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

function getOrCreateSession(key, base) {
  const existing = sessions.get(key);
  if (existing) {
    existing._lastSeen = Date.now();
    existing.requestCount = (existing.requestCount || 0) + 1;
    return existing;
  }
  pruneIfFull();
  const s = {
    id: randomUUID(),
    ...base,
    connectedAt: new Date().toISOString(),
    requestCount: 1,
    isAuth: true,
    _lastSeen: Date.now(),
  };
  sessions.set(key, s);
  return s;
}

// ── Logging ─────────────────────────────────────────────────────────────────

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

// ── Timeout ─────────────────────────────────────────────────────────────────

app.use((req, res, next) => {
  req.setTimeout(REQUEST_TIMEOUT);
  res.setTimeout(REQUEST_TIMEOUT, () => {
    if (!res.headersSent) {
      res.status(504).json({ error: 'Gateway Timeout' });
    }
  });
  next();
});

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

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

registerTools(server, prisma);

// ── Routes ──────────────────────────────────────────────────────────────────

app.get('/', (req, res) => {
  res.json({
    name: 'Memento MCP Server',
    version: '3.2.0',
    status: 'running',
    endpoints: { mcp: '/mcp', health: '/health', sessions: '/sessions' },
    auth: { enabled: process.env.MCP_REQUIRE_AUTH === 'true' },
    tools: 22,
  });
});

app.get('/sessions', (req, res) => {
  const list = [...sessions.values()].map(s => ({
    id: s.id, name: s.name, connectedAt: s.connectedAt,
    requestCount: s.requestCount || 0, authMethod: s.authMethod,
  }));
  res.json({ activeUsers: list.length, sessions: list, 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 init: ${id}`);
        transports[id] = transport;
      },
    });

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

    await server.connect(transport);
  }

  const ctx = { userId: req.userSession?.userId || null };
  await requestContext.run(ctx, async () => {
    await transport.handleRequest(req, res, req.body);
  });
});

// Legacy /sse → /mcp redirect
app.all('/sse', (req, res) => {
  res.redirect(307, '/mcp');
});

const transports = {};

// ── Start ────────────────────────────────────────────────────────────────────

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

  Server:   http://localhost:${PORT}
  MCP:      http://localhost:${PORT}/mcp
  Auth:     ${process.env.MCP_REQUIRE_AUTH === 'true' ? 'ENABLED' : 'DISABLED (dev)'}
  Timeout:  ${REQUEST_TIMEOUT}ms
  Database: ${isPostgres ? 'PostgreSQL' : 'SQLite'}
  Tools:    22
`);
});

// ── Shutdown ─────────────────────────────────────────────────────────────────

async function shutdown() {
  log('info', 'Shutting down...');
  await prisma.$disconnect();
  process.exit(0);
}

process.on('SIGINT', shutdown);
process.on('SIGTERM', shutdown);
process.on('uncaughtException', (err) => log('error', 'Uncaught:', err.message));
process.on('unhandledRejection', (reason) => log('error', 'Unhandled rejection:', reason));