package mcpserver

import (
	"context"
	"fmt"
	"reflect"

	"github.com/google/jsonschema-go/jsonschema"
	"github.com/google/uuid"
	"github.com/modelcontextprotocol/go-sdk/mcp"
)

var toolSchemaOptions = &jsonschema.ForOptions{
	TypeSchemas: map[reflect.Type]*jsonschema.Schema{
		reflect.TypeFor[uuid.UUID](): {
			Type:   "string",
			Format: "uuid",
		},
	},
}

func addTool[In any, Out any](server *mcp.Server, tool *mcp.Tool, handler func(context.Context, *mcp.CallToolRequest, In) (*mcp.CallToolResult, Out, error)) {
	if err := setToolSchemas[In, Out](tool); err != nil {
		panic(fmt.Sprintf("configure MCP tool %q schemas: %v", tool.Name, err))
	}
	mcp.AddTool(server, tool, handler)
}

func setToolSchemas[In any, Out any](tool *mcp.Tool) error {
	if tool.InputSchema == nil {
		inputSchema, err := jsonschema.For[In](toolSchemaOptions)
		if err != nil {
			return fmt.Errorf("infer input schema: %w", err)
		}
		tool.InputSchema = inputSchema
	}

	if tool.OutputSchema == nil {
		outputSchema, err := jsonschema.For[Out](toolSchemaOptions)
		if err != nil {
			return fmt.Errorf("infer output schema: %w", err)
		}
		tool.OutputSchema = outputSchema
	}

	return nil
}