package agents

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

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

const contactPrompt = `
**Role:** AI Contact Processor from Images.

**Goal:** Extract contacts from an image, check against existing list using listContacts, add *only* new contacts using createContact, and call stopAgent when finished. Avoid duplicates.

**Input:** Image potentially containing contact info (Name, Phone, Email, Address).

**Workflow:**
1.  **Scan Image:** Extract all contact details. If none, call stopAgent.
2.  **Think:** Using the think tool, you must layout your thoughts about the contacts on the image. If they are duplicates or not, and what your next action should be,
3.  **Check Duplicates:** If contacts found, *first* call listContacts. Compare extracted info to list. If all found contacts already exist, use createExistingContact.
4.  **Add New:** If you detect a new contact on the image, call createContact to create a new contact.
5.  **Finish:** Call stopAgent once all new contacts are created OR if steps 1 or 2 determined no action/creation was needed.

**Tools:**
* listContacts: Check existing contacts (Use first if contacts found).
* createContact: Add a NEW contact (Name required).
* createExistingContact: Adds this image to an existing contact, if one is found in listContacts.
* stopAgent: Signal task completion.
`

const contactTools = `
[
	{
        "type": "function",
        "function": {
            "name": "think",
            "description": "Use this tool to think through the image, evaluating the contact and whether or not it exists in the users listContacts.",
            "parameters": {
                "type": "object",
                "properties": {
					"thought": {
						"type": "string",
						"description": "A singular thought about the image"
					}
				},
                "required": ["thought"]
            }
        }
  },
  {
    "type": "function",
    "function": {
      "name": "listContacts",
      "description": "Retrieves the complete list of the user's currently saved contacts (e.g., names, phone numbers, emails if available in the stored data). This tool is essential and **must** be called *before* attempting to create a new contact if potential contact info is found in the image, to check if the person already exists and prevent duplicate entries.",
      "parameters": {
        "type": "object",
        "properties": {},
        "required": []
      }
    }
  },
  {
    "type": "function",
    "function": {
      "name": "createContact",
      "description": "Saves a new contact to the user's contact list. Only use this function **after** confirming the contact does not already exist by checking the output of listContacts. Provide all available extracted information for the new contact. Process one new contact per call.",
      "parameters": {
        "type": "object",
        "properties": {
			"contactId": {
				"type": "string",
				"description": "The UUID of the contact. You should only provide this IF you believe the contact already exists, from listContacts."
			},
          "name": {
            "type": "string",
            "description": "The full name of the person being added as a contact. This field is mandatory."
          },
          "phoneNumber": {
            "type": "string",
            "description": "The contact's primary phone number, including area or country code if available. Provide this if extracted from the image. Only include this if you see a phone number, do not make it up."
          },
          "address": {
            "type": "string",
            "description": "The complete physical mailing address of the contact (e.g., street number, street name, city, state/province, postal code, country). Provide this if extracted from the image. Only include this if you see an address. No not make it up."
          },
          "email": {
            "type": "string",
            "description": "The contact's primary email address. Provide this if extracted from the image. Only include this if you see an email, do not make it up."
          }
        },
        "required": ["name"]
      }
    }
  },
	{
        "type": "function",
        "function": {
            "name": "createExistingContact",
            "description": "Called when a contact already exists in the users list, from listContas. Only call this to indicate this image contains a duplicate.",
            "parameters": {
                "type": "object",
                "properties": {
                    "contactId": {
                        "type": "string",
                        "description": "The UUID of the contact"
                    }
                },
                "required": ["contactId"]
            }
        }
    },
  {
    "type": "function",
    "function": {
      "name": "stopAgent",
      "description": "Use this tool to signal that the contact processing for the current image is complete. Call this *only* when: 1) No contact info was found initially, OR 2) All found contacts were confirmed to already exist after calling listContacts, OR 3) All necessary createContact calls for new individuals have been completed.",
      "parameters": {
        "type": "object",
        "properties": {},
        "required": []
      }
    }
  }
]
`

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

func NewContactAgent(log *log.Logger, contactModel models.ContactModel) client.AgentClient {
	agentClient := client.CreateAgentClient(client.CreateAgentClientOptions{
		SystemPrompt: contactPrompt,
		JsonTools:    contactTools,
		Log:          log,
		EndToolCall:  "stopAgent",
	})

	agentClient.ToolHandler.AddTool("think", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
		return "Thought", nil
	})

	agentClient.ToolHandler.AddTool("listContacts", func(info client.ToolHandlerInfo, args string, call client.ToolCall) (any, error) {
		return 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()

		contactId := uuid.Nil
		if args.ContactID != nil {
			contactUuid, err := uuid.Parse(*args.ContactID)
			if err != nil {
				return model.Contacts{}, err
			}

			contactId = contactUuid
		}

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

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

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

		return contact, nil
	})

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

		ctx := context.Background()

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

		_, err = contactModel.SaveToImage(ctx, info.ImageId, contactId)
		if err != nil {
			return "", err
		}

		return "", nil
	})

	return agentClient
}
