package openapi

import (
	"context"
	"encoding/json"
	"fmt"
)

// ToolHandler is a function that implements a tool's logic.
type ToolHandler func(ctx context.Context, s *Session, args json.RawMessage) (any, error)

// ToolDefinition pairs the OpenAI-compatible tool schema with a Go handler.
type ToolDefinition struct {
	Tool    Tool
	Handler ToolHandler
}

// ToolRegistry defines the interface for tool discovery and execution.
type ToolRegistry interface {
	Tools(ctx context.Context) []Tool
	Execute(ctx context.Context, s *Session, name string, args json.RawMessage) (any, error)
}

// MapRegistry is a simple map-based implementation of ToolRegistry.
type MapRegistry struct {
	tools map[string]ToolDefinition
}

// NewMapRegistry creates a new MapRegistry.
func NewMapRegistry() *MapRegistry {
	return &MapRegistry{tools: make(map[string]ToolDefinition)}
}

// Register adds a tool to the registry.
func (r *MapRegistry) Register(name, description string, parameters any, handler ToolHandler) error {
	paramsJSON, err := json.Marshal(parameters)
	if err != nil {
		return fmt.Errorf("marshal parameters: %w", err)
	}

	r.tools[name] = ToolDefinition{
		Tool: Tool{
			Type: "function",
			Function: FunctionDefinition{
				Name:        name,
				Description: description,
				Parameters:  paramsJSON,
			},
		},
		Handler: handler,
	}
	return nil
}

// Tools returns the list of tools in the registry.
func (r *MapRegistry) Tools(ctx context.Context) []Tool {
	tools := make([]Tool, 0, len(r.tools))
	for _, td := range r.tools {
		tools = append(tools, td.Tool)
	}
	return tools
}

// Execute looks up and runs a tool by name.
func (r *MapRegistry) Execute(ctx context.Context, s *Session, name string, args json.RawMessage) (any, error) {
	td, ok := r.tools[name]
	if !ok {
		return nil, fmt.Errorf("tool %s not found", name)
	}
	return td.Handler(ctx, s, args)
}

// CompositeRegistry allows combining multiple tool registries.
type CompositeRegistry struct {
	registries []ToolRegistry
}

// Add adds a registry to the composite.
func (r *CompositeRegistry) Add(registry ToolRegistry) {
	r.registries = append(r.registries, registry)
}

// Tools returns the union of all tools from all registries.
func (r *CompositeRegistry) Tools(ctx context.Context) []Tool {
	var all []Tool
	for _, reg := range r.registries {
		all = append(all, reg.Tools(ctx)...)
	}
	return all
}

// Execute attempts to execute the tool by trying each registry in order.
func (r *CompositeRegistry) Execute(ctx context.Context, s *Session, name string, args json.RawMessage) (any, error) {
	for _, reg := range r.registries {
		res, err := reg.Execute(ctx, s, name, args)
		if err == nil {
			return res, nil
		}
	}
	return nil, fmt.Errorf("tool %s not found in any registry", name)
}