#!/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());

// 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,
  };
}

// 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,
  };
}

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

// List available tools
server.setRequestHandler(ListToolsRequestSchema, async () => {
  return {
    tools: [
      {
        name: 'create_note',
        description: 'Create a new note in Memento',
        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' },
            },
          },
          required: ['content'],
        },
      },
      {
        name: 'get_notes',
        description: 'Get all notes from Memento (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',
            },
            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',
            },
            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' },
            },
          },
          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',
            },
          },
          required: ['query'],
        },
      },
      {
        name: 'get_labels',
        description: 'Get all unique labels from notes',
        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'],
        },
      },
    ],
  };
});

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

  try {
    switch (name) {
      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,
          },
        });
        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' } },
          ];
        }

        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;
        }
        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 notes = await prisma.note.findMany({
          where: {
            isArchived: false,
            OR: [
              { title: { contains: args.query, mode: 'insensitive' } },
              { content: { contains: args.query, mode: 'insensitive' } },
            ],
          },
          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 not found');
        }
        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),
            },
          ],
        };
      }

      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: 'Memento MCP SSE Server',
    version: '1.0.0',
    status: 'running',
    endpoints: {
      sse: '/sse',
      message: '/message',
    },
  });
});

// MCP endpoint - handles both GET and POST per Streamable HTTP spec
app.all('/sse', async (req, res) => {
  console.log(`Received ${req.method} request to /sse from:`, req.ip);
  
  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 the request
  await transport.handleRequest(req, res, req.body);
});

// Store active transports
const transports = {};

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

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

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

🛠️  Available Tools (9):
   1. create_note    - Create new note
   2. get_notes      - Get all notes
   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 all labels
   8. toggle_pin     - Pin/unpin note
   9. toggle_archive - Archive/unarchive note

📋 Database: ${join(__dirname, '../keep-notes/prisma/dev.db')}

🌐 For N8N configuration:
   Use SSE endpoint: http://YOUR_IP:${PORT}/sse

💡 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 MCP SSE server...');
  await prisma.$disconnect();
  process.exit(0);
});