#!/usr/bin/env node
import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.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';

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

// Initialize Prisma Client with correct database path
const prisma = new PrismaClient({
  datasources: {
    db: {
      url: 'file:D:/dev_new_pc/Keep/keep-notes/prisma/dev.db'
    }
  }
});

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

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

// 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',
        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',
            },
          },
        },
      },
      {
        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',
        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'],
        },
      },

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

  try {
    // === NOTE TOOLS ===
    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,
            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' },
          ],
        });

        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(notes.map(parseNote), 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 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),
            },
          ],
        };
      }

      // === NOTEBOOK TOOLS ===
      case 'create_notebook': {
        // Get the highest order value
        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'];
        
        // Check if label already exists in this notebook
        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}`
    );
  }
});

// Start server
async function main() {
  const transport = new StdioServerTransport();
  await server.connect(transport);
  console.error('Keep Notes MCP server running on stdio');
}

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