package agents

import (
	"context"
	"encoding/json"
	"os"
	"screenmark/screenmark/.gen/haystack/haystack/model"
	"screenmark/screenmark/agents/client"
	"screenmark/screenmark/models"
	"time"

	"github.com/charmbracelet/log"
	"github.com/google/uuid"
)

const locationPrompt = `
You are an agent.

The user will send you images and you have to identify if they have any location or a place. This could a picture of a real place, an address, or it's name.

There are various tools you can use to perform this task.

listLocations
Lists the users already existing locations, you should do this before using createLocation to avoid creating duplicates. 

createLocation
Use this to create a new location. Avoid making duplicates and only create a new location if listLocations doesnt contain the location on the image.

linkLocation
Links an image to a location.

finish
Call when there is nothing else to do.
`

const locationTools = `
[
    {
        "type": "function",
        "function": {
            "name": "listLocations",
            "description": "List the locations the user already has.",
            "parameters": {
                "type": "object",
                "properties": {},
                "required": []
            }
        }
    },
	{
        "type": "function",
        "function": {
            "name": "createLocation",
            "description": "Use to create a new location",
            "parameters": {
                "type": "object",
                "properties": {
                    "name": {
                        "type": "string"
                    },
                    "address": {
                        "type": "string"
                    }
				},
                "required": ["name"]
            }
        }
    },
	{
        "type": "function",
        "function": {
            "name": "linkLocation",
            "description": "Use to link an already existing location to the image you were sent",
            "parameters": {
                "type": "object",
                "properties": {
					"locationId": {
						"type": "string"
					}
				},
                "required": ["locationId"]
            }
        }
    },
	{
        "type": "function",
        "function": {
            "name": "finish",
            "description": "Call this when there is nothing left to do.",
            "parameters": {
                "type": "object",
                "properties": {},
                "required": []
            }
        }
    }
]`

type LocationAgent struct {
	client client.AgentClient

	locationModel models.LocationModel
}

type listLocationArguments struct{}
type createLocationArguments struct {
	Name    string  `json:"name"`
	Address *string `json:"address"`
}
type linkLocationArguments struct {
	LocationID string `json:"locationId"`
}

func NewLocationAgent(locationModel models.LocationModel) (LocationAgent, error) {
	agentClient, err := client.CreateAgentClient(log.NewWithOptions(os.Stdout, log.Options{
		ReportTimestamp: true,
		TimeFormat:      time.Kitchen,
		Prefix:          "Locations 📍",
	}))

	if err != nil {
		return LocationAgent{}, err
	}

	agent := LocationAgent{
		client:        agentClient,
		locationModel: locationModel,
	}

	agentClient.ToolHandler.AddTool("listLocations", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
		return agent.locationModel.List(context.Background(), info.UserId)
	})

	agentClient.ToolHandler.AddTool("createLocation", func(info client.ToolHandlerInfo, _args string, call client.ToolCall) (any, error) {
		args := createLocationArguments{}
		err := json.Unmarshal([]byte(_args), &args)
		if err != nil {
			return model.Locations{}, err
		}

		ctx := context.Background()

		location, err := agent.locationModel.Save(ctx, info.UserId, model.Locations{
			Name:    args.Name,
			Address: args.Address,
		})

		if err != nil {
			return model.Locations{}, err
		}

		_, err = agent.locationModel.SaveToImage(ctx, info.ImageId, location.ID)
		if err != nil {
			return model.Locations{}, err
		}

		return location, nil
	})

	agentClient.ToolHandler.AddTool("linkLocation", func(info client.ToolHandlerInfo, _args string, call client.ToolCall) (any, error) {
		args := linkLocationArguments{}
		err := json.Unmarshal([]byte(_args), &args)
		if err != nil {
			return "", err
		}

		ctx := context.Background()

		contactUuid, err := uuid.Parse(args.LocationID)
		if err != nil {
			return "", err
		}

		agent.locationModel.SaveToImage(ctx, info.ImageId, contactUuid)
		return "Saved", nil
	})

	return agent, nil
}