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 contactPrompt = `
You are an agent that performs actions on contacts and people you find on an image.

You can use tools to achieve your task.

You should use listContacts to make sure that you don't create duplicate contacts.

Call createContact when you see there is a new contact on this image.

Call finish if you dont think theres anything else to do.
`

const contactTools = `
[
    {
        "type": "function",
        "function": {
            "name": "listContacts",
            "description": "List the users existing contacts",
            "parameters": {
                "type": "object",
                "properties": {},
                "required": []
            }
        }
    },
	{
        "type": "function",
        "function": {
            "name": "createContact",
            "description": "Creates a new contact",
            "parameters": {
                "type": "object",
                "properties": {
                    "name": {
                        "type": "string",
						"description": "the name of the person"
                    },
                    "phoneNumber": {
                        "type": "string"
                    },
                    "address": {
                        "type": "string",
						"description": "their physical address"
                    },
                    "email": {
                        "type": "string"
                    }
                },
                "required": ["name"]
            }
        }
    },
	{
        "type": "function",
        "function": {
            "name": "finish",
            "description": "Call when you dont think theres anything to do",
            "parameters": {
                "type": "object",
                "properties": {},
                "required": []
            }
        }
    }
]
`

type ContactAgent struct {
	client client.AgentClient

	contactModel models.ContactModel
}

type listContactsArguments struct{}
type createContactsArguments struct {
	Name        string  `json:"name"`
	PhoneNumber *string `json:"phoneNumber"`
	Address     *string `json:"address"`
	Email       *string `json:"email"`
}

// Yeah this is just a copy of the other one.
func (agent ContactAgent) GetContacts(userId uuid.UUID, imageId uuid.UUID, imageName string, imageData []byte) error {
	var tools any
	err := json.Unmarshal([]byte(contactTools), &tools)

	toolChoice := "any"

	request := client.AgentRequestBody{
		Tools:       &tools,
		ToolChoice:  &toolChoice,
		Model:       "pixtral-12b-2409",
		Temperature: 0.3,
		EndToolCall: "finish",
		ResponseFormat: client.ResponseFormat{
			Type: "text",
		},
		Chat: &client.Chat{
			Messages: make([]client.ChatMessage, 0),
		},
	}

	request.Chat.AddSystem(eventLocationPrompt)
	request.Chat.AddImage(imageName, imageData)

	_, err = agent.client.Request(&request)
	if err != nil {
		return err
	}

	toolHandlerInfo := client.ToolHandlerInfo{
		ImageId: imageId,
		UserId:  userId,
	}

	return agent.client.ToolLoop(toolHandlerInfo, &request)
}

func NewContactAgent(contactModel models.ContactModel) (ContactAgent, error) {
	agentClient, err := client.CreateAgentClient(log.NewWithOptions(os.Stdout, log.Options{
		ReportTimestamp: true,
		TimeFormat:      time.Kitchen,
		Prefix:          "Contacts 👥",
	}))
	if err != nil {
		return ContactAgent{}, err
	}

	agent := ContactAgent{
		client:       agentClient,
		contactModel: contactModel,
	}

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

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

		ctx := context.Background()

		contact, err := agent.contactModel.Save(ctx, info.UserId, model.Contacts{
			Name:        args.Name,
			PhoneNumber: args.PhoneNumber,
			Email:       args.Email,
		})

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

		_, err = agent.contactModel.SaveToImage(ctx, info.ImageId, contact.ID)
		if err != nil {
			return model.Contacts{}, err
		}

		return contact, nil
	})

	return agent, nil
}