package core

import (
	"context"
	"encoding/xml"
	"fmt"
	"nudgebee/llm/agents/prompts_repo"
	"nudgebee/llm/common"
	"nudgebee/llm/security"
	"strconv"
	"strings"
	"time"

	"github.com/tmc/langchaingo/llms"
	"github.com/tmc/langchaingo/prompts"
)

const (
	planCritiquerDecisionAccept = "accept"
	planCritiquerDecisionReject = "reject"
)

// Response from the plan critiquer LLM
type rewooPlanCritiqueResponse struct {
	XMLName  xml.Name `xml:"plan_critique"`
	Thought  string   `xml:"thought"`
	Decision string   `xml:"decision"` // "accept" or "reject"
	Feedback string   `xml:"feedback"`
}

// ReWooPlanCritiquer critiques a plan generated by the ReWoo planner.
type ReWooPlanCritiquer struct {
	ctx     *security.RequestContext
	prompt  prompts.FormatPrompter
	request NBAgentRequest
}

// NewReWooPlanCritiquer creates a new plan critiquer.
func NewReWooPlanCritiquer(ctx *security.RequestContext, request NBAgentRequest, nbAgent NBAgent) (*ReWooPlanCritiquer, error) {
	promptTemplate := prompts.NewPromptTemplate(
		prompts_repo.GetPrompt(prompts_repo.PromptPlannerRewooPlanCritiquer),
		[]string{"input", "tool_names", "tool_descriptions", "plan"},
	)

	return &ReWooPlanCritiquer{
		ctx:     ctx,
		prompt:  promptTemplate,
		request: request,
	}, nil
}

// CritiquePlan evaluates a plan and decides if it needs refinement.
func (c *ReWooPlanCritiquer) CritiquePlan(
	_ context.Context, // ctx is unused, but kept for interface compatibility
	input string,
	planToCritique string,
	toolNames string,
	toolDescriptions string,
	maxStepsForPlan int,
) (bool, string, error) {
	questionType := "query"
	if IsInvestigationRequestTask(input) {
		questionType = "investigation"
	}

	fullInputs := map[string]any{
		"input":             input,
		"tool_names":        toolNames,
		"tool_descriptions": toolDescriptions,
		"plan":              planToCritique,
		"today":             time.Now().Format("January 02, 2006"),
		"max_plan_steps":    strconv.Itoa(maxStepsForPlan),
		"question_type":     questionType,
	}

	prompt, err := c.prompt.FormatPrompt(fullInputs)
	if err != nil {
		return false, "", fmt.Errorf("failed to format plan critique prompt: %w", err)
	}

	mcList := []llms.MessageContent{llms.TextParts(llms.ChatMessageTypeHuman, prompt.String())}

	result, err := GenerateAndTrackLLMContent(c.ctx, c.request.UserId, c.request.AccountId, c.request.ConversationId, c.request.MessageId, "rewoo_planner", false, mcList, true,
		llms.WithTemperature(0.0))

	if err != nil {
		return false, "", fmt.Errorf("plan critiquer LLM call failed: %w", err)
	}
	if len(result.Choices) == 0 {
		return false, "", fmt.Errorf("plan critiquer LLM returned no choices")
	}

	return c.parseOutput(result.Choices[0].Content)
}

// parseOutput parses the XML response from the critiquer LLM.
func (c *ReWooPlanCritiquer) parseOutput(output string) (bool, string, error) {
	output = common.XmlSanitize(output, "plan_critique")

	var critiqueResp rewooPlanCritiqueResponse
	err := xml.Unmarshal([]byte(output), &critiqueResp)
	if err != nil {
		return false, "", fmt.Errorf("failed to unmarshal plan critique response XML: %w, data: %s", err, output)
	}

	// Decision is to reject, refinement is needed.
	if strings.ToLower(critiqueResp.Decision) == planCritiquerDecisionReject {
		return true, critiqueResp.Feedback, nil
	}

	// Decision is to accept, no refinement needed.
	if strings.ToLower(critiqueResp.Decision) == planCritiquerDecisionAccept {
		return false, "", nil
	}

	return false, "", fmt.Errorf("could not parse decision from critiquer output: %s", output)
}
