import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { z } from "zod";
import type { Logger } from "pino";

import {
  backlinkDocumentId,
  collectionId,
  direction,
  limit,
  offset,
  parentDocumentId,
  sort,
  userId,
} from "./common";
import { handleSuccess, handleError } from "./utils";
import type { ToolsFactory } from "./toolsFactory";
import type { DocumentsApi } from "../gen/api/outline";

export const registerDocumentTools: ToolsFactory<DocumentsApi> = (
  server: McpServer,
  client: DocumentsApi,
  logger: Logger
): void => {
  server.registerTool(
    "documents_info",
    {
      title: "Get Document Info",
      description:
        "Retrieve a document by its UUID, urlId, or shareId. At least one of these parameters must be provided.",
      inputSchema: {
        id: z
          .string()
          .optional()
          .describe(
            "Unique identifier for the document. Either the UUID or the urlId is acceptable."
          ),
        shareId: z
          .string()
          .optional()
          .describe(
            "Unique identifier for a document share, a shareId may be used in place of a document UUID"
          ),
      },
    },
    async (args) => {
      try {
        const response = await client.documentsInfo({
          documentsInfoRequest: {
            id: args?.id,
            shareId: args?.shareId,
          },
        });
        return handleSuccess(
          response,
          logger.child({ tool: "documents_info" })
        );
      } catch (error) {
        return handleError(error, logger.child({ tool: "documents_info" }));
      }
    }
  );

  const documentsListSchema = z.object({
    offset,
    limit,
    sort,
    direction,
    collectionId,
    userId,
    backlinkDocumentId,
    parentDocumentId,
    template: z.boolean().optional(),
  });

  server.registerTool(
    "documents_list",
    {
      title: "List Documents",
      description:
        "This method will list all published documents and draft documents belonging to the current user.",
      inputSchema: documentsListSchema.shape,
    },
    async (args) => {
      try {
        const validatedArgs = documentsListSchema.parse(args);
        const response = await client.documentsList({
          documentsListRequest: validatedArgs,
        });
        return handleSuccess(
          response,
          logger.child({ tool: "documents_list" })
        );
      } catch (error) {
        return handleError(error, logger.child({ tool: "documents_list" }));
      }
    }
  );
};