package agents

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

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

const locationPrompt = `
Role: Location AI Assistant

Objective: Identify locations from images/text, manage a saved list (create, update), and answer user queries about saved locations using the provided tools.

Core Logic:

**Extract Location Details:** Attempt to extract location details (like InputName, InputAddress) from the user's input (image or text).
	* If no details can be extracted, inform the user and use stopAgent.

**Check for Existing Location:** If details *were* extracted:
    * Use listLocations with the extracted InputName and/or InputAddress to search for potentially matching locations already saved in the list.

**Decide Action based on Search Results:**
    * **If listLocations returns one or more likely matches:**
        * Identify the *best* match (based on name, address similarity).
        * **Crucially:** Call upsertLocation, providing the locationId of that best match. Include the newly extracted InputName (required) and any other extracted details (InputAddress, etc.) to potentially *update* the existing record or simply link the current input to it.
    * **If listLocations returns no matches OR no returned location is a confident match:**
        * Call upsertLocation providing *only* the newly extracted InputName (required) and any other extracted details (InputAddress, etc.). **Do NOT provide a locationId in this case.** This will create a *new* location entry.

4.  **Finalize:** After successfully calling upsertLocation (or determining no action could be taken), use stopAgent.

Tool Usage:

* **listLocations**: Searches the saved locations list based on provided criteria (like name or address). Used specifically to check if a location potentially already exists before using upsertLocation. Returns a list of matching locations, *each including its locationId*.
* **upsertLocation**: Creates or updates a location in the saved list. Requires name. Can include address, etc.
    * **To UPDATE:** If you identified an existing location using listLocations, provide its locationId along with any new/updated details (name, address, etc.).
    * **To CREATE:** If no existing location was found (or you are creating intentionally), provide the location details (name, address, etc.) but **omit the locationId**.
* **stopAgent**: Signals the end of the agent's processing for the current turn. Call this *after* completing the location task (create/update/failed extraction).
`

const replyTool = `
{
        "type": "function",
        "function": {
            "name": "reply",
            "description": "Signals intent to provide information about a specific known location in response to a user's query. Use only if the user asked a question and the location's ID was found via listLocations.",
            "parameters": {
                "type": "object",
                "properties": {
                    "locationId": {
                        "type": "string",
                        "description": "The unique identifier of the saved location that the user is asking about."
                    }
                },
                "required": ["locationId"]
            }
        }
},`

const locationTools = `
[
    {
        "type": "function",
        "function": {
            "name": "listLocations",
            "description": "Retrieves the list of the user's currently saved locations (names, addresses, IDs). Use this first to check if a location from an image already exists, or to find the ID of a location the user is asking about.",
            "parameters": {
                "type": "object",
                "properties": {},
                "required": []
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "upsertLocation",
            "description": "Upserts a location. This is used for both creating new locations, and updating existing ones. Providing locationId from an existing ID from listLocations, will make this an update function. Not providing one will create a new location. You must provide a locationId if you think the input is a location that already exists.",
            "parameters": {
                "type": "object",
                "properties": {
                    "name": {
                        "type": "string",
                        "description": "The primary name of the location (e.g., 'Eiffel Tower', 'Mom's House', 'Acme Corp HQ'). This field is mandatory."
                    },
					"locationId": {
						"type": "string",
						"description": "The UUID of the location. You should only provide this IF you believe the location already exists, from listLocation."
					},
                    "address": {
                        "type": "string",
                        "description": "The full street address of the location, if available (e.g., 'Champ de Mars, 5 Av. Anatole France, 75007 Paris, France'). Include if extracted."
                    }
                },
                "required": ["name"]
            }
        }
    },
	%s
	{
		"type": "function",
		"function": {
		  "name": "stopAgent",
		  "description": "Use this tool to signal that the contact processing for the current image is complete.",
		  "parameters": {
			"type": "object",
			"properties": {},
			"required": []
		  }
		}
	  }
]`

func getLocationAgentTools(allowReply bool) string {
	if allowReply {
		return fmt.Sprintf(locationTools, replyTool)
	} else {
		return fmt.Sprintf(locationTools, "")
	}
}

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

func NewLocationAgentWithComm(log *log.Logger, locationModel models.LocationModel) client.AgentClient {
	client := NewLocationAgent(log, locationModel)

	client.Options.JsonTools = getLocationAgentTools(true)

	return client
}

func NewLocationAgent(log *log.Logger, locationModel models.LocationModel) client.AgentClient {
	agentClient := client.CreateAgentClient(client.CreateAgentClientOptions{
		SystemPrompt: locationPrompt,
		JsonTools:    getLocationAgentTools(false),
		Log:          log,
		EndToolCall:  "stopAgent",
	})

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

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

		ctx := context.Background()

		locationId := uuid.Nil
		if args.LocationID != nil {
			locationUuid, err := uuid.Parse(*args.LocationID)
			if err != nil {
				return model.Locations{}, err
			}

			locationId = locationUuid
		}

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

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

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

		return location, nil
	})

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

	return agentClient
}