#!/usr/bin/env node
import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
import {
  CallToolRequestSchema,
  ErrorCode,
  ListToolsRequestSchema,
  McpError,
} from '@modelcontextprotocol/sdk/types.js';
import { PrismaClient } from '@prisma/client';
import { fileURLToPath } from 'url';
import { dirname, join } from 'path';
import { randomUUID } from 'crypto';
import express from 'express';
import cors from 'cors';

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

const app = express();
const PORT = process.env.PORT || 3001;

// Middleware
app.use(cors());
app.use(express.json());

// === USER SESSION MANAGEMENT ===

// Active user sessions
const userSessions = {};

// Middleware pour l'authentification
app.use((req, res, next) => {
  const apiKey = req.headers['x-api-key'];
  const userId = req.headers['x-user-id'];
  
  // Mode dev: pas d'authentification requise
  if (process.env.MCP_REQUIRE_AUTH !== 'true') {
    req.userSession = {
      id: 'dev-user',
      name: 'Development User',
      connectedAt: new Date().toISOString(),
      isAuth: false
    };
    return next();
  }

  // Mode production: vérifier API Key
  if (!apiKey && !userId) {
    return res.status(401).json({
      error: 'Authentication required',
      message: 'Please provide x-api-key or x-user-id header'
    });
  }

  const sessionKey = userId || apiKey;
  
  if (userSessions[sessionKey]) {
    // Session existante, la réutiliser
    req.userSession = userSessions[sessionKey];
    req.userSession.lastSeen = new Date().toISOString();
    console.log(`👤 User session reused: ${req.userSession.name} (${req.userSession.id})`);
  } else {
    // Nouvelle session
    req.userSession = {
      id: randomUUID(),
      name: userId || `API-Key-${apiKey.substring(0, 8)}`,
      connectedAt: new Date().toISOString(),
      lastSeen: new Date().toISOString(),
      requestCount: 0,
      isAuth: true
    };
    userSessions[sessionKey] = req.userSession;
    console.log(`👤 New user session: ${req.userSession.name} (${req.userSession.id})`);
  }

  next();
});

// Middleware de logging des requêtes
app.use((req, res, next) => {
  if (req.userSession) {
    req.userSession.requestCount = (req.userSession.requestCount || 0) + 1;
    
    console.log(`📝 [${req.userSession.id.substring(0, 8)}] ${req.method} ${req.path} - Request #${req.userSession.requestCount}`);
  }
  
  next();
});

// Initialize Prisma Client
const prisma = new PrismaClient();

// Helper to parse JSON fields
function parseNote(dbNote) {
  return {
    ...dbNote,
    checkItems: dbNote.checkItems ? JSON.parse(dbNote.checkItems) : null,
    labels: dbNote.labels ? JSON.parse(dbNote.labels) : null,
    images: dbNote.images ? JSON.parse(dbNote.images) : null,
    links: dbNote.links ? JSON.parse(dbNote.links) : null,
  };
}

// Helper to parse Notebook fields
function parseNotebook(dbNotebook) {
  return {
    ...dbNotebook,
    labels: dbNotebook.labels || [],
  };
}

// Helper to parse note with lightweight format (no images, truncated content)
function parseNoteLightweight(dbNote) {
  return {
    id: dbNote.id,
    title: dbNote.title,
    content: dbNote.content.length > 200 ? dbNote.content.substring(0, 200) + '...' : dbNote.content,
    color: dbNote.color,
    type: dbNote.type,
    isPinned: dbNote.isPinned,
    isArchived: dbNote.isArchived,
    hasImages: !!dbNote.images,
    imageCount: dbNote.images ? JSON.parse(dbNote.images).length : 0,
    labels: dbNote.labels ? JSON.parse(dbNote.labels) : null,
    hasCheckItems: !!dbNote.checkItems,
    checkItemsCount: dbNote.checkItems ? JSON.parse(dbNote.checkItems).length : 0,
    reminder: dbNote.reminder,
    createdAt: dbNote.createdAt,
    updatedAt: dbNote.updatedAt,
    notebookId: dbNote.notebookId,
  };
}

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

// === USER ACTIVITY TRACKING ===

// Log user activities to database (optional - for analytics)
async function logUserActivity(userId, toolName, args) {
  try {
    // Could be stored in a separate UserActivity table
    // For now, just console log
    console.log(`🔍 Activity: User ${userId} called ${toolName}`);
  } catch (error) {
    console.error(`Error logging activity:`, error.message);
  }
}

// List available tools
server.setRequestHandler(ListToolsRequestSchema, async () => {
  return {
    tools: [
      // === NOTE TOOLS ===
      {
        name: 'create_note',
        description: 'Create a new note in Keep Notes',
        inputSchema: {
          type: 'object',
          properties: {
            title: {
              type: 'string',
              description: 'Note title (optional)',
            },
            content: {
              type: 'string',
              description: 'Note content',
            },
            color: {
              type: 'string',
              description: 'Note color (default, red, orange, yellow, green, teal, blue, purple, pink, gray)',
              default: 'default',
            },
            type: {
              type: 'string',
              enum: ['text', 'checklist'],
              description: 'Note type',
              default: 'text',
            },
            checkItems: {
              type: 'array',
              description: 'Checklist items (if type is checklist)',
              items: {
                type: 'object',
                properties: {
                  id: { type: 'string' },
                  text: { type: 'string' },
                  checked: { type: 'boolean' },
                },
                required: ['id', 'text', 'checked'],
              },
            },
            labels: {
              type: 'array',
              description: 'Note labels/tags',
              items: { type: 'string' },
            },
            isPinned: {
              type: 'boolean',
              description: 'Pin the note',
              default: false,
            },
            isArchived: {
              type: 'boolean',
              description: 'Archive the note',
              default: false,
            },
            images: {
              type: 'array',
              description: 'Note images as base64 encoded strings',
              items: { type: 'string' },
            },
            links: {
              type: 'array',
              description: 'Note links',
              items: { type: 'string' },
            },
            reminder: {
              type: 'string',
              description: 'Reminder date/time (ISO 8601 format)',
            },
            isReminderDone: {
              type: 'boolean',
              description: 'Mark reminder as done',
              default: false,
            },
            reminderRecurrence: {
              type: 'string',
              description: 'Reminder recurrence (daily, weekly, monthly, yearly)',
            },
            reminderLocation: {
              type: 'string',
              description: 'Reminder location',
            },
            isMarkdown: {
              type: 'boolean',
              description: 'Enable markdown support',
              default: false,
            },
            size: {
              type: 'string',
              enum: ['small', 'medium', 'large'],
              description: 'Note size',
              default: 'small',
            },
            notebookId: {
              type: 'string',
              description: 'Notebook ID to associate the note with',
            },
          },
          required: ['content'],
        },
      },
      {
        name: 'get_notes',
        description: 'Get all notes from Keep Notes (lightweight format: titles, truncated content, no images to reduce payload size)',
        inputSchema: {
          type: 'object',
          properties: {
            includeArchived: {
              type: 'boolean',
              description: 'Include archived notes',
              default: false,
            },
            search: {
              type: 'string',
              description: 'Search query to filter notes',
            },
            notebookId: {
              type: 'string',
              description: 'Filter notes by notebook ID',
            },
            fullDetails: {
              type: 'boolean',
              description: 'Return full note details including images (warning: large payload)',
              default: false,
            },
          },
        },
      },
      {
        name: 'get_note',
        description: 'Get a specific note by ID',
        inputSchema: {
          type: 'object',
          properties: {
            id: {
              type: 'string',
              description: 'Note ID',
            },
          },
          required: ['id'],
        },
      },
      {
        name: 'update_note',
        description: 'Update an existing note',
        inputSchema: {
          type: 'object',
          properties: {
            id: {
              type: 'string',
              description: 'Note ID',
            },
            title: {
              type: 'string',
              description: 'Note title',
            },
            content: {
              type: 'string',
              description: 'Note content',
            },
            color: {
              type: 'string',
              description: 'Note color',
            },
            type: {
              type: 'string',
              enum: ['text', 'checklist'],
              description: 'Note type',
            },
            checkItems: {
              type: 'array',
              description: 'Checklist items',
              items: {
                type: 'object',
                properties: {
                  id: { type: 'string' },
                  text: { type: 'string' },
                  checked: { type: 'boolean' },
                },
              },
            },
            labels: {
              type: 'array',
              description: 'Note labels',
              items: { type: 'string' },
            },
            isPinned: {
              type: 'boolean',
              description: 'Pin status',
            },
            isArchived: {
              type: 'boolean',
              description: 'Archive status',
            },
            images: {
              type: 'array',
              description: 'Note images as base64 encoded strings',
              items: { type: 'string' },
            },
            links: {
              type: 'array',
              description: 'Note links',
              items: { type: 'string' },
            },
            reminder: {
              type: 'string',
              description: 'Reminder date/time (ISO 8601 format)',
            },
            isReminderDone: {
              type: 'boolean',
              description: 'Mark reminder as done',
            },
            reminderRecurrence: {
              type: 'string',
              description: 'Reminder recurrence',
            },
            reminderLocation: {
              type: 'string',
              description: 'Reminder location',
            },
            isMarkdown: {
              type: 'boolean',
              description: 'Enable markdown support',
            },
            size: {
              type: 'string',
              enum: ['small', 'medium', 'large'],
              description: 'Note size',
            },
            notebookId: {
              type: 'string',
              description: 'Notebook ID to move the note to',
            },
          },
          required: ['id'],
        },
      },
      {
        name: 'delete_note',
        description: 'Delete a note by ID',
        inputSchema: {
          type: 'object',
          properties: {
            id: {
              type: 'string',
              description: 'Note ID',
            },
          },
          required: ['id'],
        },
      },
      {
        name: 'search_notes',
        description: 'Search notes by query',
        inputSchema: {
          type: 'object',
          properties: {
            query: {
              type: 'string',
              description: 'Search query',
            },
            notebookId: {
              type: 'string',
              description: 'Filter search by notebook ID',
            },
          },
          required: ['query'],
        },
      },
      {
        name: 'get_labels',
        description: 'Get all unique labels from notes (legacy method)',
        inputSchema: {
          type: 'object',
          properties: {},
        },
      },
      {
        name: 'toggle_pin',
        description: 'Toggle pin status of a note',
        inputSchema: {
          type: 'object',
          properties: {
            id: {
              type: 'string',
              description: 'Note ID',
            },
          },
          required: ['id'],
        },
      },
      {
        name: 'toggle_archive',
        description: 'Toggle archive status of a note',
        inputSchema: {
          type: 'object',
          properties: {
            id: {
              type: 'string',
              description: 'Note ID',
            },
          },
          required: ['id'],
        },
      },

      // === USER MANAGEMENT TOOLS (NEW) ===
      {
        name: 'get_current_user',
        description: 'Get information about the currently authenticated user',
        inputSchema: {
          type: 'object',
          properties: {},
        },
      },
      {
        name: 'get_all_users',
        description: 'Get list of all currently active users/sessions',
        inputSchema: {
          type: 'object',
          properties: {},
        },
      },
      {
        name: 'logout',
        description: 'Log out current user session',
        inputSchema: {
          type: 'object',
          properties: {
            sessionId: {
              type: 'string',
              description: 'Session ID to logout (optional, defaults to current session)',
            },
          },
        },
      },

      // === NOTEBOOK TOOLS ===
      {
        name: 'create_notebook',
        description: 'Create a new notebook',
        inputSchema: {
          type: 'object',
          properties: {
            name: {
              type: 'string',
              description: 'Notebook name',
            },
            icon: {
              type: 'string',
              description: 'Notebook icon (emoji)',
            },
            color: {
              type: 'string',
              description: 'Notebook color (hex code)',
            },
            order: {
              type: 'number',
              description: 'Notebook order',
            },
          },
          required: ['name'],
        },
      },
      {
        name: 'get_notebooks',
        description: 'Get all notebooks',
        inputSchema: {
          type: 'object',
          properties: {},
        },
      },
      {
        name: 'get_notebook',
        description: 'Get a specific notebook by ID with its notes',
        inputSchema: {
          type: 'object',
          properties: {
            id: {
              type: 'string',
              description: 'Notebook ID',
            },
          },
          required: ['id'],
        },
      },
      {
        name: 'update_notebook',
        description: 'Update an existing notebook',
        inputSchema: {
          type: 'object',
          properties: {
            id: {
              type: 'string',
              description: 'Notebook ID',
            },
            name: {
              type: 'string',
              description: 'Notebook name',
            },
            icon: {
              type: 'string',
              description: 'Notebook icon',
            },
            color: {
              type: 'string',
              description: 'Notebook color',
            },
            order: {
              type: 'number',
              description: 'Notebook order',
            },
          },
          required: ['id'],
        },
      },
      {
        name: 'delete_notebook',
        description: 'Delete a notebook by ID',
        inputSchema: {
          type: 'object',
          properties: {
            id: {
              type: 'string',
              description: 'Notebook ID',
            },
          },
          required: ['id'],
        },
      },

      // === LABEL TOOLS ===
      {
        name: 'create_label',
        description: 'Create a new label',
        inputSchema: {
          type: 'object',
          properties: {
            name: {
              type: 'string',
              description: 'Label name',
            },
            color: {
              type: 'string',
              description: 'Label color (red, orange, yellow, green, teal, blue, purple, pink, gray)',
            },
            notebookId: {
              type: 'string',
              description: 'Notebook ID to associate the label with',
            },
          },
          required: ['name', 'notebookId'],
        },
      },
      {
        name: 'get_labels_detailed',
        description: 'Get all labels with details',
        inputSchema: {
          type: 'object',
          properties: {
            notebookId: {
              type: 'string',
              description: 'Filter labels by notebook ID',
            },
          },
        },
      },
      {
        name: 'update_label',
        description: 'Update an existing label',
        inputSchema: {
          type: 'object',
          properties: {
            id: {
              type: 'string',
              description: 'Label ID',
            },
            name: {
              type: 'string',
              description: 'Label name',
            },
            color: {
              type: 'string',
              description: 'Label color',
            },
          },
          required: ['id'],
        },
      },
      {
        name: 'delete_label',
        description: 'Delete a label by ID',
        inputSchema: {
          type: 'object',
          properties: {
            id: {
              type: 'string',
              description: 'Label ID',
            },
          },
          required: ['id'],
        },
      },
    ],
  };
});

// Handle tool calls
server.setRequestHandler(CallToolRequestSchema, async (request) => {
  const { name, arguments: args } = request.params;
  const userSession = request.userSession;

  // Log activity
  if (userSession && userSession.id) {
    await logUserActivity(userSession.id, name, args);
  }

  try {
    switch (name) {
      // === NOTE TOOLS ===
      case 'create_note': {
        const note = await prisma.note.create({
          data: {
            title: args.title || null,
            content: args.content,
            color: args.color || 'default',
            type: args.type || 'text',
            checkItems: args.checkItems ? JSON.stringify(args.checkItems) : null,
            labels: args.labels ? JSON.stringify(args.labels) : null,
            isPinned: args.isPinned || false,
            isArchived: args.isArchived || false,
            images: args.images ? JSON.stringify(args.images) : null,
            links: args.links ? JSON.stringify(args.links) : null,
            reminder: args.reminder ? new Date(args.reminder) : null,
            isReminderDone: args.isReminderDone || false,
            reminderRecurrence: args.reminderRecurrence || null,
            reminderLocation: args.reminderLocation || null,
            isMarkdown: args.isMarkdown || false,
            size: args.size || 'small',
            notebookId: args.notebookId || null,
          },
        });
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(parseNote(note), null, 2),
            },
          ],
        };
      }

      case 'get_notes': {
        let where = {};
        if (!args.includeArchived) {
          where.isArchived = false;
        }
        if (args.search) {
          where.OR = [
            { title: { contains: args.search, mode: 'insensitive' } },
            { content: { contains: args.search, mode: 'insensitive' } },
          ];
        }
        if (args.notebookId) {
          where.notebookId = args.notebookId;
        }

        const notes = await prisma.note.findMany({
          where,
          orderBy: [
            { isPinned: 'desc' },
            { order: 'asc' },
            { updatedAt: 'desc' },
          ],
        });

        // Use lightweight format by default, full details only if requested
        const parsedNotes = args.fullDetails
          ? notes.map(parseNote)
          : notes.map(parseNoteLightweight);

        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(parsedNotes, null, 2),
            },
          ],
        };
      }

      case 'get_note': {
        const note = await prisma.note.findUnique({
          where: { id: args.id },
        });
        if (!note) {
          throw new McpError(ErrorCode.InvalidRequest, 'Note not found');
        }
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(parseNote(note), null, 2),
            },
          ],
        };
      }

      case 'update_note': {
        const updateData = { ...args };
        delete updateData.id;

        if ('checkItems' in args) {
          updateData.checkItems = args.checkItems
            ? JSON.stringify(args.checkItems)
            : null;
        }
        if ('labels' in args) {
          updateData.labels = args.labels ? JSON.stringify(args.labels) : null;
        }
        if ('images' in args) {
          updateData.images = args.images ? JSON.stringify(args.images) : null;
        }
        if ('links' in args) {
          updateData.links = args.links ? JSON.stringify(args.links) : null;
        }
        if ('reminder' in args) {
          updateData.reminder = args.reminder ? new Date(args.reminder) : null;
        }
        updateData.updatedAt = new Date();

        const note = await prisma.note.update({
          where: { id: args.id },
          data: updateData,
        });

        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(parseNote(note), null, 2),
            },
          ],
        };
      }

      case 'delete_note': {
        await prisma.note.delete({
          where: { id: args.id },
        });
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify({ success: true, message: 'Note deleted' }),
            },
          ],
        };
      }

      case 'search_notes': {
        const where = {
          isArchived: false,
          OR: [
            { title: { contains: args.query, mode: 'insensitive' } },
            { content: { contains: args.query, mode: 'insensitive' } },
          ],
        };
        if (args.notebookId) {
          where.notebookId = args.notebookId;
        }

        const notes = await prisma.note.findMany({
          where,
          orderBy: [
            { isPinned: 'desc' },
            { updatedAt: 'desc' },
          ],
        });

        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(notes.map(parseNote), null, 2),
            },
          ],
        };
      }

      case 'get_labels': {
        const notes = await prisma.note.findMany({
          select: { labels: true },
        });

        const labelsSet = new Set();
        notes.forEach((note) => {
          const labels = note.labels ? JSON.parse(note.labels) : null;
          if (labels) {
            labels.forEach((label) => labelsSet.add(label));
          }
        });

        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(Array.from(labelsSet).sort(), null, 2),
            },
          ],
        };
      }

      case 'toggle_pin': {
        const note = await prisma.note.findUnique({ where: { id: args.id } });
        if (!note) {
          throw new McpError(ErrorCode.InvalidRequest, 'Note not found');
        }
        const updated = await prisma.note.update({
          where: { id: args.id },
          data: { isPinned: !note.isPinned },
        });
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(parseNote(updated), null, 2),
            },
          ],
        };
      }

      case 'toggle_archive': {
        const note = await prisma.note.findUnique({ where: { id: args.id } });
        if (!note) {
          throw new McpError(ErrorCode.InvalidRequest, 'Note failed');
        }
        const updated = await prisma.note.update({
          where: { id: args.id },
          data: { isArchived: !note.isArchived },
        });
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(parseNote(updated), null, 2),
            },
          ],
        };
      }

      // === USER MANAGEMENT TOOLS ===
      case 'get_current_user': {
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(userSession || null, null, 2),
            },
          ],
        };
      }

      case 'get_all_users': {
        const users = Object.values(userSessions).map(session => ({
          id: session.id,
          name: session.name,
          connectedAt: session.connectedAt,
          lastSeen: session.lastSeen,
          requestCount: session.requestCount || 0,
          isAuth: session.isAuth || false
        }));

        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(users, null, 2),
            },
          ],
        };
      }

      case 'logout': {
        const sessionId = args.sessionId;
        
        if (sessionId && userSessions[sessionId]) {
          delete userSessions[sessionId];
          console.log(`👋 User logged out: ${sessionId}`);
          return {
            content: [
              {
                type: 'text',
                text: JSON.stringify({ success: true, message: 'Logged out successfully', sessionId }),
              },
            ],
          };
        }

        // Logout current session
        if (userSession) {
          const sessionKeys = Object.keys(userSessions);
          for (const key of sessionKeys) {
            if (userSessions[key].id === userSession.id) {
              delete userSessions[key];
              console.log(`👋 Current user logged out: ${userSession.name} (${userSession.id})`);
              break;
            }
          }
        }

        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify({ success: true, message: 'Logged out successfully' }),
            },
          ],
        };
      }

      // === NOTEBOOK TOOLS ===
      case 'create_notebook': {
        const highestOrder = await prisma.notebook.findFirst({
          orderBy: { order: 'desc' },
          select: { order: true }
        });

        const nextOrder = args.order !== undefined ? args.order : (highestOrder?.order ?? -1) + 1;

        const notebook = await prisma.notebook.create({
          data: {
            name: args.name.trim(),
            icon: args.icon || '📁',
            color: args.color || '#3B82F6',
            order: nextOrder,
          },
          include: {
            labels: true,
            _count: {
              select: { notes: true }
            }
          }
        });

        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify({
                ...notebook,
                notesCount: notebook._count.notes
              }, null, 2),
            },
          ],
        };
      }

      case 'get_notebooks': {
        const notebooks = await prisma.notebook.findMany({
          include: {
            labels: {
              orderBy: { name: 'asc' }
            },
            _count: {
              select: { notes: true }
            }
          },
          orderBy: { order: 'asc' }
        });

        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(
                notebooks.map(nb => ({
                  ...nb,
                  notesCount: nb._count.notes
                })),
                null, 2
              ),
            },
          ],
        };
      }

      case 'get_notebook': {
        const notebook = await prisma.notebook.findUnique({
          where: { id: args.id },
          include: {
            labels: true,
            notes: true,
            _count: {
              select: { notes: true }
            }
          }
        });

        if (!notebook) {
          throw new McpError(ErrorCode.InvalidRequest, 'Notebook not found');
        }

        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify({
                ...notebook,
                notes: notebook.notes.map(parseNote),
                notesCount: notebook._count.notes
              }, null, 2),
            },
          ],
        };
      }

      case 'update_notebook': {
        const updateData = { ...args };
        delete updateData.id;

        const notebook = await prisma.notebook.update({
          where: { id: args.id },
          data: updateData,
          include: {
            labels: true,
            _count: {
              select: { notes: true }
            }
          }
        });

        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify({
                ...notebook,
                notesCount: notebook._count.notes
              }, null, 2),
            },
          ],
        };
      }

      case 'delete_notebook': {
        await prisma.notebook.delete({
          where: { id: args.id },
        });
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify({ success: true, message: 'Notebook deleted' }),
            },
          ],
        };
      }

      // === LABEL TOOLS ===
      case 'create_label': {
        const COLORS = ['red', 'orange', 'yellow', 'green', 'teal', 'blue', 'purple', 'pink', 'gray'];

        const existing = await prisma.label.findFirst({
          where: {
            name: args.name.trim(),
            notebookId: args.notebookId
          }
        });

        if (existing) {
          throw new McpError(ErrorCode.InvalidRequest, 'Label already exists in this notebook');
        }

        const label = await prisma.label.create({
          data: {
            name: args.name.trim(),
            color: args.color || COLORS[Math.floor(Math.random() * COLORS.length)],
            notebookId: args.notebookId,
          }
        });

        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(label, null, 2),
            },
          ],
        };
      }

      case 'get_labels_detailed': {
        const where = {};
        if (args.notebookId) {
          where.notebookId = args.notebookId;
        }

        const labels = await prisma.label.findMany({
          where,
          include: {
            notebook: {
              select: { id: true, name: true }
            }
          },
          orderBy: { name: 'asc' }
        });

        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(labels, null, 2),
            },
          ],
        };
      }

      case 'update_label': {
        const updateData = { ...args };
        delete updateData.id;

        const label = await prisma.label.update({
          where: { id: args.id },
          data: updateData,
        });

        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(label, null, 2),
            },
          ],
        };
      }

      case 'delete_label': {
        await prisma.label.delete({
          where: { id: args.id },
        });
        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify({ success: true, message: 'Label deleted' }),
            },
          ],
        };
      }

      default:
        throw new McpError(
          ErrorCode.MethodNotFound,
          `Unknown tool: ${name}`
        );
    }
  } catch (error) {
    if (error instanceof McpError) {
      throw error;
    }
    throw new McpError(
      ErrorCode.InternalError,
      `Tool execution failed: ${error.message}`
    );
  }
});

// Health check endpoint
app.get('/', (req, res) => {
  res.json({
    name: 'Keep Notes MCP SSE Server',
    version: '2.0.0',
    status: 'running',
    endpoints: {
      sse: '/sse',
      message: '/message',
    },
    authentication: {
      enabled: process.env.MCP_REQUIRE_AUTH === 'true',
      mode: 'dev' !== 'true' ? 'Disabled' : 'Enabled',
      method: 'Provide x-api-key or x-user-id header'
    },
  });
});

// User session status endpoint
app.get('/sessions', (req, res) => {
  const sessions = Object.values(userSessions).map(session => ({
    id: session.id,
    name: session.name,
    connectedAt: session.connectedAt,
    lastSeen: session.lastSeen,
    requestCount: session.requestCount || 0,
    isAuth: session.isAuth || false
  }));

  res.json({
    activeUsers: sessions.length,
    sessions: sessions
  });
});

// MCP endpoint - handles both GET and POST per Streamable HTTP spec
app.all('/sse', async (req, res) => {
  const sessionId = req.headers['mcp-session-id'];
  let transport;

  if (sessionId && transports[sessionId]) {
    // Reuse existing transport
    transport = transports[sessionId];
  } else {
    // Create new transport with session management
    transport = new StreamableHTTPServerTransport({
      sessionIdGenerator: () => randomUUID(),
      onsessioninitialized: (id) => {
        console.log(`Session initialized: ${id}`);
        transports[id] = transport;
      }
    });

    // Set up close handler
    transport.onclose = () => {
      const sid = transport.sessionId;
      if (sid && transports[sid]) {
        console.log(`Transport closed for session ${sid}`);
        delete transports[sid];
      }
    };

    // Connect to MCP server
    await server.connect(transport);
  }

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

// Store active transports
const transports = {};

// Start server
app.listen(PORT, '0.0.0.0', () => {
  console.log(`
╔═══════════════════════════════════════════════════════╗
║  🎉 Keep Notes MCP SSE Server Started (v2.0.0)       ║
╚═════════════════════════════════════════════════════════╝

📡 Server running on:
   - Local:   http://localhost:${PORT}
   - Network: http://0.0.0.0:${PORT}

🔌 Endpoints:
   - Health:   GET  http://localhost:${PORT}/
   - Sessions: GET  http://localhost:${PORT}/sessions
   - SSE:      GET  http://localhost:${PORT}/sse
   - Message:  POST http://localhost:${PORT}/message

🔐 Authentication: ${process.env.MCP_REQUIRE_AUTH === 'true' ? 'ENABLED' : 'DISABLED (dev mode)'}

🛠️  Available Tools (22):
   === NOTES (9) ===
   1. create_note     - Create new note with full support
   2. get_notes       - Get all notes (supports filters)
   3. get_note        - Get note by ID
   4. update_note     - Update note
   5. delete_note     - Delete note
   6. search_notes    - Search notes
   7. get_labels      - Get unique labels (legacy)
   8. toggle_pin      - Pin/unpin note
   9. toggle_archive  - Archive/unarchive note

   === USER MANAGEMENT (3) === 🆕
   10. get_current_user - Get current authenticated user info
   11. get_all_users   - List all active user sessions
   12. logout          - Logout user session

   === NOTEBOOKS (5) ===
   13. create_notebook  - Create new notebook
   14. get_notebooks   - Get all notebooks
   15. get_notebook    - Get notebook with notes
   16. update_notebook - Update notebook
   17. delete_notebook - Delete notebook

   === LABELS (5) ===
   18. create_label      - Create label
   19. get_labels_detailed - Get labels with details
   20. update_label      - Update label
   21. delete_label      - Delete label

📋 Database: PostgreSQL (via DATABASE_URL env var)

🌐 For N8N configuration:
   Use SSE endpoint: http://YOUR_IP:${PORT}/sse
   Add headers: x-api-key or x-user-id

💡 Find your IP with: ipconfig (Windows) or ifconfig (Mac/Linux)

Press Ctrl+C to stop
`);
});

// Graceful shutdown
process.on('SIGINT', async () => {
  console.log('\n\n🛑 Shutting down Keep Notes MCP SSE server...');
  console.log(`👋 Active sessions: ${Object.keys(userSessions).length}`);
  await prisma.$disconnect();
  process.exit(0);
});
