package client

import (
	"encoding/json"
	"errors"
	"log"

	"github.com/google/uuid"
)

type ToolHandlerInfo struct {
	UserId  uuid.UUID
	ImageId uuid.UUID
}

type ToolHandler struct {
	Fn func(info ToolHandlerInfo, args string, call ToolCall) (string, error)
}

type ToolsHandlers struct {
	handlers *map[string]ToolHandler
}

func (handler ToolsHandlers) Handle(info ToolHandlerInfo, request *AgentRequestBody) (string, error) {
	agentMessage := request.Messages[len(request.Messages)-1]

	toolCall, ok := agentMessage.(AgentAssistantToolCall)
	if !ok {
		return "", errors.New("Latest message was not a tool call.")
	}

	fnName := toolCall.ToolCalls[0].Function.Name
	arguments := toolCall.ToolCalls[0].Function.Arguments

	fnHandler, exists := (*handler.handlers)[fnName]
	if !exists {
		return "", errors.New("Could not find tool with this name.")
	}

	log.Printf("Calling: %s\n", fnName)
	res, err := fnHandler.Fn(info, arguments, toolCall.ToolCalls[0])
	if err != nil {
		return "", err
	}

	request.AddText(AgentTextMessage{
		Role:       "tool",
		Name:       "createLocation",
		Content:    res,
		ToolCallId: toolCall.ToolCalls[0].Id,
	})

	return res, nil
}

func (handler ToolsHandlers) AddTool(name string, getArgs func() any, fn func(info ToolHandlerInfo, args any, call ToolCall) (any, error)) {
	(*handler.handlers)["createLocation"] = ToolHandler{
		Fn: func(info ToolHandlerInfo, args string, call ToolCall) (string, error) {
			argsStruct := getArgs()

			err := json.Unmarshal([]byte(args), &argsStruct)
			if err != nil {
				return "", err
			}

			res, err := fn(info, argsStruct, call)
			if err != nil {
				return "", err
			}

			marshalledRes, err := json.Marshal(res)
			if err != nil {
				return "", err
			}

			return string(marshalledRes), nil
		},
	}
}