import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import {
  CallToolRequestSchema,
  ListToolsRequestSchema,
  Tool,
} from '@modelcontextprotocol/sdk/types.js';
import { NotionClient } from './clients/notion.js';
import type {
  PageId,
  DatabaseId,
  BlockId,
  UserId,
  Block,
  PageProperty,
  DatabaseProperty,
  Filter,
  Sort,
} from './types/index.js';

export interface NotionServerConfig {
  apiKey: string;
}

export class NotionServer {
  private server: Server;
  private client: NotionClient;

  constructor(config: NotionServerConfig) {
    this.client = new NotionClient({ auth: config.apiKey });

    this.server = new Server(
      {
        name: '@mcpengine/notion',
        version: '0.1.0',
      },
      {
        capabilities: {
          tools: {},
        },
      }
    );

    this.setupHandlers();
  }

  private setupHandlers(): void {
    this.server.setRequestHandler(ListToolsRequestSchema, async () => {
      const tools: Tool[] = [
        // Pages module
        {
          name: 'notion_get_page',
          description: 'Retrieve a Notion page by ID',
          inputSchema: {
            type: 'object',
            properties: {
              page_id: {
                type: 'string',
                description: 'The ID of the page to retrieve',
              },
            },
            required: ['page_id'],
          },
        },
        {
          name: 'notion_create_page',
          description: 'Create a new page in a database or as a child of another page',
          inputSchema: {
            type: 'object',
            properties: {
              parent_type: {
                type: 'string',
                enum: ['database_id', 'page_id'],
                description: 'Type of parent',
              },
              parent_id: {
                type: 'string',
                description: 'ID of the parent database or page',
              },
              properties: {
                type: 'object',
                description: 'Page properties as JSON object',
              },
              children: {
                type: 'array',
                description: 'Array of block objects to append to the page',
              },
            },
            required: ['parent_type', 'parent_id', 'properties'],
          },
        },
        {
          name: 'notion_update_page',
          description: 'Update page properties or archive a page',
          inputSchema: {
            type: 'object',
            properties: {
              page_id: {
                type: 'string',
                description: 'The ID of the page to update',
              },
              properties: {
                type: 'object',
                description: 'Properties to update',
              },
              archived: {
                type: 'boolean',
                description: 'Whether to archive the page',
              },
            },
            required: ['page_id'],
          },
        },

        // Databases module
        {
          name: 'notion_get_database',
          description: 'Retrieve a database by ID',
          inputSchema: {
            type: 'object',
            properties: {
              database_id: {
                type: 'string',
                description: 'The ID of the database',
              },
            },
            required: ['database_id'],
          },
        },
        {
          name: 'notion_query_database',
          description: 'Query a database with filters and sorting',
          inputSchema: {
            type: 'object',
            properties: {
              database_id: {
                type: 'string',
                description: 'The ID of the database to query',
              },
              filter: {
                type: 'object',
                description: 'Filter object (optional)',
              },
              sorts: {
                type: 'array',
                description: 'Array of sort objects (optional)',
              },
              page_size: {
                type: 'number',
                description: 'Number of results to return (max 100)',
              },
            },
            required: ['database_id'],
          },
        },
        {
          name: 'notion_create_database',
          description: 'Create a new database',
          inputSchema: {
            type: 'object',
            properties: {
              parent_page_id: {
                type: 'string',
                description: 'ID of the parent page',
              },
              title: {
                type: 'string',
                description: 'Database title',
              },
              properties: {
                type: 'object',
                description: 'Database schema properties',
              },
            },
            required: ['parent_page_id', 'title', 'properties'],
          },
        },

        // Blocks module
        {
          name: 'notion_get_block',
          description: 'Retrieve a block by ID',
          inputSchema: {
            type: 'object',
            properties: {
              block_id: {
                type: 'string',
                description: 'The ID of the block',
              },
            },
            required: ['block_id'],
          },
        },
        {
          name: 'notion_get_block_children',
          description: 'Retrieve children blocks of a block',
          inputSchema: {
            type: 'object',
            properties: {
              block_id: {
                type: 'string',
                description: 'The ID of the parent block',
              },
              page_size: {
                type: 'number',
                description: 'Number of results to return',
              },
            },
            required: ['block_id'],
          },
        },
        {
          name: 'notion_append_block_children',
          description: 'Append child blocks to a parent block',
          inputSchema: {
            type: 'object',
            properties: {
              block_id: {
                type: 'string',
                description: 'The ID of the parent block',
              },
              children: {
                type: 'array',
                description: 'Array of block objects to append',
              },
            },
            required: ['block_id', 'children'],
          },
        },
        {
          name: 'notion_delete_block',
          description: 'Delete (archive) a block',
          inputSchema: {
            type: 'object',
            properties: {
              block_id: {
                type: 'string',
                description: 'The ID of the block to delete',
              },
            },
            required: ['block_id'],
          },
        },

        // Users module
        {
          name: 'notion_get_user',
          description: 'Retrieve a user by ID',
          inputSchema: {
            type: 'object',
            properties: {
              user_id: {
                type: 'string',
                description: 'The ID of the user',
              },
            },
            required: ['user_id'],
          },
        },
        {
          name: 'notion_list_users',
          description: 'List all users in the workspace',
          inputSchema: {
            type: 'object',
            properties: {
              page_size: {
                type: 'number',
                description: 'Number of results to return',
              },
            },
          },
        },

        // Comments module
        {
          name: 'notion_create_comment',
          description: 'Add a comment to a page or block',
          inputSchema: {
            type: 'object',
            properties: {
              parent_type: {
                type: 'string',
                enum: ['page_id', 'block_id'],
                description: 'Type of parent',
              },
              parent_id: {
                type: 'string',
                description: 'ID of the parent page or block',
              },
              text: {
                type: 'string',
                description: 'Comment text content',
              },
            },
            required: ['parent_type', 'parent_id', 'text'],
          },
        },
        {
          name: 'notion_list_comments',
          description: 'Retrieve comments for a block',
          inputSchema: {
            type: 'object',
            properties: {
              block_id: {
                type: 'string',
                description: 'The ID of the block',
              },
            },
            required: ['block_id'],
          },
        },

        // Search module
        {
          name: 'notion_search',
          description: 'Search all pages and databases',
          inputSchema: {
            type: 'object',
            properties: {
              query: {
                type: 'string',
                description: 'Search query text',
              },
              filter: {
                type: 'string',
                enum: ['page', 'database'],
                description: 'Filter by object type',
              },
              sort_direction: {
                type: 'string',
                enum: ['ascending', 'descending'],
                description: 'Sort direction for last_edited_time',
              },
            },
          },
        },
      ];

      return { tools };
    });

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

      try {
        if (!args) {
          throw new Error('Missing required arguments');
        }

        switch (name) {
          // Pages
          case 'notion_get_page': {
            const page = await this.client.getPage(args.page_id as PageId);
            return {
              content: [{ type: 'text', text: JSON.stringify(page, null, 2) }],
            };
          }

          case 'notion_create_page': {
            const parent =
              args.parent_type === 'database_id'
                ? { database_id: args.parent_id as DatabaseId }
                : { page_id: args.parent_id as PageId };
            const page = await this.client.createPage({
              parent,
              properties: args.properties as Record<string, Partial<PageProperty>>,
              children: args.children as Block[] | undefined,
            });
            return {
              content: [{ type: 'text', text: JSON.stringify(page, null, 2) }],
            };
          }

          case 'notion_update_page': {
            const page = await this.client.updatePage(args.page_id as PageId, {
              properties: args.properties as Record<string, Partial<PageProperty>> | undefined,
              archived: args.archived as boolean | undefined,
            });
            return {
              content: [{ type: 'text', text: JSON.stringify(page, null, 2) }],
            };
          }

          // Databases
          case 'notion_get_database': {
            const db = await this.client.getDatabase(args.database_id as DatabaseId);
            return {
              content: [{ type: 'text', text: JSON.stringify(db, null, 2) }],
            };
          }

          case 'notion_query_database': {
            const results = await this.client.queryDatabase({
              database_id: args.database_id as DatabaseId,
              filter: args.filter as Filter | undefined,
              sorts: args.sorts as Sort[] | undefined,
              page_size: args.page_size as number | undefined,
            });
            return {
              content: [{ type: 'text', text: JSON.stringify(results, null, 2) }],
            };
          }

          case 'notion_create_database': {
            const db = await this.client.createDatabase({
              parent: { page_id: args.parent_page_id as PageId },
              title: [{ type: 'text', text: { content: args.title as string } }],
              properties: args.properties as Record<string, DatabaseProperty>,
            });
            return {
              content: [{ type: 'text', text: JSON.stringify(db, null, 2) }],
            };
          }

          // Blocks
          case 'notion_get_block': {
            const block = await this.client.getBlock(args.block_id as BlockId);
            return {
              content: [{ type: 'text', text: JSON.stringify(block, null, 2) }],
            };
          }

          case 'notion_get_block_children': {
            const children = await this.client.getBlockChildren(
              args.block_id as BlockId,
              undefined,
              args.page_size as number | undefined
            );
            return {
              content: [{ type: 'text', text: JSON.stringify(children, null, 2) }],
            };
          }

          case 'notion_append_block_children': {
            const result = await this.client.appendBlockChildren(
              args.block_id as BlockId,
              args.children as Block[]
            );
            return {
              content: [{ type: 'text', text: JSON.stringify(result, null, 2) }],
            };
          }

          case 'notion_delete_block': {
            const block = await this.client.deleteBlock(args.block_id as BlockId);
            return {
              content: [{ type: 'text', text: JSON.stringify(block, null, 2) }],
            };
          }

          // Users
          case 'notion_get_user': {
            const user = await this.client.getUser(args.user_id as UserId);
            return {
              content: [{ type: 'text', text: JSON.stringify(user, null, 2) }],
            };
          }

          case 'notion_list_users': {
            const users = await this.client.listUsers(undefined, args.page_size as number | undefined);
            return {
              content: [{ type: 'text', text: JSON.stringify(users, null, 2) }],
            };
          }

          // Comments
          case 'notion_create_comment': {
            const parent =
              args.parent_type === 'page_id'
                ? { page_id: args.parent_id as PageId }
                : { block_id: args.parent_id as BlockId };
            const comment = await this.client.createComment({
              parent,
              rich_text: [{ type: 'text', text: { content: args.text as string } }],
            });
            return {
              content: [{ type: 'text', text: JSON.stringify(comment, null, 2) }],
            };
          }

          case 'notion_list_comments': {
            const comments = await this.client.listComments(args.block_id as BlockId);
            return {
              content: [{ type: 'text', text: JSON.stringify(comments, null, 2) }],
            };
          }

          // Search
          case 'notion_search': {
            const params: any = {};
            if (args.query) params.query = args.query;
            if (args.filter) {
              params.filter = { value: args.filter, property: 'object' };
            }
            if (args.sort_direction) {
              params.sort = {
                direction: args.sort_direction,
                timestamp: 'last_edited_time',
              };
            }
            const results = await this.client.search(params);
            return {
              content: [{ type: 'text', text: JSON.stringify(results, null, 2) }],
            };
          }

          default:
            throw new Error(`Unknown tool: ${name}`);
        }
      } catch (error) {
        const errorMessage =
          error instanceof Error ? error.message : 'Unknown error occurred';
        return {
          content: [{ type: 'text', text: `Error: ${errorMessage}` }],
          isError: true,
        };
      }
    });
  }

  async connect(transport: StdioServerTransport): Promise<void> {
    await this.server.connect(transport);
  }

  getServer(): Server {
    return this.server;
  }
}
