package agents

import (
	"encoding/json"
	"fmt"
	"nudgebee/llm/agents/core"
	"nudgebee/llm/security"
	"nudgebee/llm/tools"
	toolcore "nudgebee/llm/tools/core"
	"strings"
)

func init() {
}

const LokiAgentName = "loki"

// newLokiAgent is used by integration tests (build tag `e2e`).
// nolint:unused // referenced from agent_*_e2e_test.go files.
func newLokiAgent(accountId string) LokiAgent {
	return LokiAgent{ // This agent is typically called by the generic 'logs' agent.
		accountId: accountId,
	}
}

type LokiAgent struct {
	accountId string
}

func (l LokiAgent) GetName() string {
	return LokiAgentName
}

func (l LokiAgent) GetNameAliases() []string {
	return []string{"Loki"}
}

func (l LokiAgent) GetDescription() string {
	return `Retrieves logs from Loki. Input should be a natural language question about logs.`
}

func (l LokiAgent) GetSystemPrompt(ctx *security.RequestContext, query core.NBAgentRequest) core.NBAgentPrompt {

	instructions := []string{
		"**Understand the Question:** Carefully analyze the user's question to identify the log information they need.",
		"**Generate LogQL Query:** Use the `logql_query_generator` tool to generate a correct LogQL query.",
		"**Execute Query:** Always use the `loki_execute` tool to execute the generated LogQL query.",
		"**Present Results:** Examine the results returned by the `loki_execute` tool and provide a concise, markdown-readable summary.",
		"**Error Handling:** Handle errors from both `logql_query_generator` and `loki_execute` gracefully and provide informative error messages to the user.",
		"**Context:** Pay attention to specifics in the query given by users such as time range filters, workloads, namespaces, and any keywords or error messages",
		"**Process:**  First use 'logql_query_generator' tool to generate query && then use 'loki_execute' tool to execute the query",
		"**Query Validation:**  To validate query generated from 'logql_query_generator' always execute generated query using 'loki_execute' tool to decide if query is valid or not. DO NOT ASSUME ANYTHING without executing the query",
	}

	constraints := []string{
		"You MUST use the `logql_query_generator` tool to generate the LogQL query.",
		"You MUST use the `loki_execute` tool to execute the LogQL query generated by `logql_query_generator` tool.",
		"You MUST NOT answer the question without using the tools.",
		"Do not ask for any clarification from the user, try to resolve using the available tools",
	}

	toolUsage := map[string][]string{
		LogqlAgentName: {
			"Use this tool to generate a valid LogQL query.",
			"Input: Natural language request for logs",
			"Output: Valid LogQL query",
		},
		tools.ToolGetLokiLogs: {
			"Use this tool to execute a LogQL query against Loki.",
			"Input: A valid LogQL query.",
			"Output: JSON response from Loki.",
		},
	}

	return core.NBAgentPrompt{
		Role:         "an SRE expert specializing in Loki log analysis",
		Instructions: instructions,
		Constraints:  constraints,
		ToolUsage:    toolUsage,
		Examples: []core.NBAgentPromptExample{
			{
				Question: "Show me logs for the application 'my-app'.",
				Answer: `				
					tool - logql_query_generator
					input - generate query for - Show me logs for the application 'my-app' in namespace abc
					output - {app="my-app", namespace="abc"}
					}
					thought - I will use logql_query_generator tool to generate query for loki to execute

					tool - loki_execute
					input - {app="my-app", namespace="abc"}
					output -  Loki Logs returned by executing the query				
					thought - i got query from logql_query_generator tool, now i will use loki_execute tool to execute the query 
				`,
				Explanation: "First use `logql_query_generator` tool to generate query and then use `loki_execute` tool to execute the generated query to get logs",
			},
		},
	}
}

func (p LokiAgent) GetSupportedTools(ctx *security.RequestContext) []toolcore.NBTool {
	toolsList := []toolcore.NBTool{}
	if tool, ok := toolcore.GetNBTool(p.accountId, LogqlAgentName); ok {
		toolsList = append(toolsList, tool)
	}
	toolsList = append(toolsList, tools.LokiExecuteTool{})
	return toolsList
}

func (l LokiAgent) GetPlannerType() core.AgentPlannerType {
	return core.AgentPlannerTypeReAct
}

func (l LokiAgent) PostProcessResponse(ctx *security.RequestContext, request core.NBAgentRequest, resp core.NBAgentResponse) core.NBAgentResponse {
	responseStr := ""
	responseLogQL := ""
	for _, stepRes := range resp.AgentStepResponse {
		if stepRes.Call.FunctionCall.Name == "loki_execute" {
			responseStr = stepRes.Response.Content
			responseLogQL = stepRes.Call.FunctionCall.Arguments
		}
	}
	if responseStr != "" && responseLogQL != "" {
		result := map[string]string{
			"logql_query": responseLogQL,
			"logs":        responseStr,
		}
		jsonBytes, err := json.Marshal(result)
		if err == nil {
			resp.Response = []string{string(jsonBytes)}
		}
	}
	return resp
}

func (l LokiAgent) UpdateToolResponseForPlanner(toolRequest core.NBAgentPlannerToolAction, toolResponse string) string {
	if strings.EqualFold(toolRequest.Tool, tools.ToolGetLokiLogs) {
		originalCount := tools.CountLogsInObservabilityResponse(toolResponse)

		newToolResponse, err := tools.GetErrorLinesFromObservabilityLogString(toolResponse)
		if err != nil {
			return toolResponse
		}

		// If filtering reduced the count, prepend a summary so the LLM knows data was filtered
		filteredCount := tools.CountLogsInObservabilityResponse(newToolResponse)
		if originalCount > 0 && filteredCount < originalCount {
			summary := fmt.Sprintf(
				"<note>Showing %d error/warning lines out of %d total log lines. "+
					"Non-error lines were filtered for brevity. If the user asked for all logs (not just errors), "+
					"mention that %d additional non-error log lines exist.</note>\n",
				filteredCount, originalCount, originalCount-filteredCount)
			return summary + newToolResponse
		}
		return newToolResponse
	} else if strings.EqualFold(toolRequest.Tool, LogqlAgentName) {
		return fmt.Sprintf(`<output>%s</output>`, toolResponse)
	}
	return toolResponse
}
