import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"
import z from "zod"
import { readDictionary } from "@/services/dictionaries"
import { generateMadeUpWords } from "@/services/makeUpWords"

export const addMakeUpWordsTool = (server: McpServer) => {
  server.registerTool(
    "generate_made_up_words_from_dict",
    {
      title: "Generate Made-up Words from Dictionary",
      description: "Generate new words using Markov chains based on a dictionary.",
      inputSchema: z.object({
        dictionary_name: z.string(),
        count: z.number().optional().default(1),
        minLength: z.number().optional(),
        maxLength: z.number().optional(),
        order: z.number().optional().default(2),
        startWith: z.string().optional(),
        maxAttempts: z.number().optional().default(1000)
      })
    },
    async (input) => {
      try {
        if (input.count <= 0) {
          return {
            isError: true,
            content: [{ type: "text", text: "Count must be a positive number" }]
          }
        }
        if (input.minLength !== undefined && input.maxLength !== undefined && input.minLength > input.maxLength) {
          return {
            isError: true,
            content: [{ type: "text", text: "minLength cannot be greater than maxLength" }]
          }
        }

        const words: string[] = (await readDictionary(input.dictionary_name)).split("\n").filter((w) => w.trim())

        const generatedWords = generateMadeUpWords(words, input.count, {
          minLength: input.minLength,
          maxLength: input.maxLength,
          order: input.order,
          startWith: input.startWith,
          maxAttempts: input.maxAttempts
        })

        return { content: [{ type: "text", text: generatedWords.join("\n") }] }
      } catch (error) {
        if (error instanceof Error) {
          return {
            isError: true,
            content: [{ type: "text", text: `Error generating words: ${error.message}` }]
          }
        }
        return {
          isError: true,
          content: [{ type: "text", text: "Unknown error generating words" }]
        }
      }
    }
  )

  server.registerTool(
    "generate_made_up_words_from_list",
    {
      title: "Generate Made-up Words from List",
      description: "Generate new words using Markov chains based on a list of words.",
      inputSchema: z.object({
        words: z.array(z.string()),
        count: z.number().optional().default(1),
        minLength: z.number().optional(),
        maxLength: z.number().optional(),
        order: z.number().optional().default(2),
        startWith: z.string().optional(),
        maxAttempts: z.number().optional().default(1000)
      })
    },
    async (input) => {
      if (input.count <= 0) {
        return {
          isError: true,
          content: [{ type: "text", text: "Count must be a positive number" }]
        }
      }
      if (input.minLength !== undefined && input.maxLength !== undefined && input.minLength > input.maxLength) {
        return {
          isError: true,
          content: [{ type: "text", text: "minLength cannot be greater than maxLength" }]
        }
      }

      const generatedWords = generateMadeUpWords(input.words, input.count, {
        minLength: input.minLength,
        maxLength: input.maxLength,
        order: input.order,
        startWith: input.startWith,
        maxAttempts: input.maxAttempts
      })

      return { content: [{ type: "text", text: generatedWords.join("\n") }] }
    }
  )
}