import { z } from 'zod';
import { Tool } from '@modelcontextprotocol/sdk/types.js';
import type { NotionClient } from '../clients/notion.js';
import type { PageId, BlockId } from '../types/index.js';

// Zod schemas
const CreateCommentSchema = z.object({
  parent_type: z.enum(['page_id', 'block_id']).describe('Type of parent (page_id or block_id)'),
  parent_id: z.string().describe('ID of the parent page or block'),
  text: z.string().describe('Comment text content'),
});

const ListCommentsSchema = z.object({
  block_id: z.string().describe('The ID of the block'),
  page_size: z.number().min(1).max(100).optional().describe('Number of results to return'),
  start_cursor: z.string().optional().describe('Pagination cursor'),
});

const ListCommentsAllSchema = z.object({
  block_id: z.string().describe('The ID of the block'),
});

export function getTools(client: NotionClient): Tool[] {
  return [
    {
      name: 'notion_create_comment',
      description:
        'Add a comment to a page or block. Comments are visible in the Notion UI and can be used for collaboration.',
      inputSchema: {
        type: 'object',
        properties: {
          parent_type: {
            type: 'string',
            enum: ['page_id', 'block_id'],
            description: 'Type of parent (page_id or block_id)',
          },
          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. Returns paginated list of comments in a discussion thread.',
      inputSchema: {
        type: 'object',
        properties: {
          block_id: {
            type: 'string',
            description: 'The ID of the block',
          },
          page_size: {
            type: 'number',
            description: 'Number of results to return (max 100)',
          },
          start_cursor: {
            type: 'string',
            description: 'Pagination cursor',
          },
        },
        required: ['block_id'],
      },
    },
    {
      name: 'notion_list_comments_all',
      description:
        'Retrieve ALL comments for a block (auto-paginating). Returns complete comment thread.',
      inputSchema: {
        type: 'object',
        properties: {
          block_id: {
            type: 'string',
            description: 'The ID of the block',
          },
        },
        required: ['block_id'],
      },
    },
  ];
}

export async function handleCommentTool(
  toolName: string,
  args: any,
  client: NotionClient
): Promise<any> {
  switch (toolName) {
    case 'notion_create_comment': {
      const validated = CreateCommentSchema.parse(args);
      const parent =
        validated.parent_type === 'page_id'
          ? { page_id: validated.parent_id as PageId }
          : { block_id: validated.parent_id as BlockId };

      const comment = await client.createComment({
        parent,
        rich_text: [{ type: 'text', text: { content: validated.text } }],
      });
      return comment;
    }

    case 'notion_list_comments': {
      const validated = ListCommentsSchema.parse(args);
      const comments = await client.listComments(
        validated.block_id as BlockId,
        validated.start_cursor,
        validated.page_size
      );
      return comments;
    }

    case 'notion_list_comments_all': {
      const validated = ListCommentsAllSchema.parse(args);
      const allComments = [];
      for await (const comment of client.listCommentsAll(validated.block_id as BlockId)) {
        allComments.push(comment);
      }
      return { object: 'list', results: allComments, has_more: false, next_cursor: null };
    }

    default:
      throw new Error(`Unknown comment tool: ${toolName}`);
  }
}
