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:

**Handle Image/Text Location (if no query was handled in Step 2):**
    * If location details (InputName, InputAddress, etc.) were successfully extracted from the input in Step 1:
        * Use listLocations to check if a location matching InputName or InputAddress already exists in the saved list.
        * **If *no match*** is found:
            * Use createLocation, providing the extracted InputName (required) and any other details like InputAddress.
            * *(Proceed to step 5)*
        * **If a *match*** is found (meaning it's a potential duplicate/update candidate - let the matching saved location have ExistingLocationId):
            * Use updateLocation. Provide the locationId = ExistingLocationId.
            * Also provide any *new or potentially refined* details extracted from the current input (e.g., name = InputName, address = InputAddress). The updateLocation tool should handle updating the record with these details and/or linking the new input context (like the image) to this existing location.
            * *(Proceed to step 5)*
    * If no location details could be extracted from the input in Step 1.
        * *(Proceed to step 5)*

**Summarize & Stop:** Always finish by writing a concise message explaining the outcome of the turn. Examples:
    * "Okay, I've answered your query about [Location Name]." (after calling reply)
    * "I couldn't find [Queried Location Name] in my saved list." (after failing to find a match for a query)
    * "I've saved '[InputName]' as a new location." (after calling createLocation)
    * "I found that '[InputName]' was already saved, so I've updated its information/context based on your latest input." (after calling updateLocation)
    * "I couldn't identify a specific location from your input." (if Step 1 failed or no action was taken)
    * After providing the summary message, call stopAgent to signal the end of processing for this turn.

Tool Usage:

* listLocations: Check saved locations. Used to find matches for user queries or to detect existing entries before creating/updating. Returns matching location(s) including their locationId.
* createLocation: Save a *new* location. Requires name, can include address, etc.
* updateLocation: Update an *existing* location. Requires locationId. Can include name, address, etc., to update specific fields or simply to associate the new input context with the existing location.
* stopAgent: Signals the end of the agent's processing for the current turn. Call this *after* providing the summary message.
`

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": "createLocation",
            "description": "Creates a new location entry in the user's saved list. Use only after listLocations confirms the location does not already exist.",
            "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."
                    },
                    "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"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "updateLocation",
            "description": "Updates an existing saved location identified by its locationId. Use when input matches a pre-existing location. Pass locationId and any new details (name, address) to update.",
            "parameters": {
                "type": "object",
                "properties": {
                    "locationId": {
                        "type": "string",
                        "description": "The UUID of the location you are trying to update"
                    }
                },
                "required": ["locationId"]
            }
        }
    },
	%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 createLocationArguments struct {
	Name    string  `json:"name"`
	Address *string `json:"address"`
}
type updateLocationArguments struct {
	LocationID string `json:"locationId"`
}

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("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 := locationModel.Save(ctx, info.UserId, model.Locations{
			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("updateLocation", func(info client.ToolHandlerInfo, _args string, call client.ToolCall) (any, error) {
		args := updateLocationArguments{}
		err := json.Unmarshal([]byte(_args), &args)
		if err != nil {
			return "", err
		}

		ctx := context.Background()

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

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

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

	return agentClient
}