package tools

import (
	"context"
	"encoding/json"
	"fmt"
	"maps"
	"strings"

	"pentagi/pkg/database"
	"pentagi/pkg/database/knowledge/limits"
	"pentagi/pkg/graph/model"
	obs "pentagi/pkg/observability"
	"pentagi/pkg/observability/langfuse"
	"pentagi/pkg/providers/embeddings"

	"github.com/sirupsen/logrus"
	"github.com/vxcontrol/cloud/anonymizer"
	"github.com/vxcontrol/langchaingo/documentloaders"
	"github.com/vxcontrol/langchaingo/schema"
	"github.com/vxcontrol/langchaingo/vectorstores"
	"github.com/vxcontrol/langchaingo/vectorstores/pgvector"
)

const (
	guideVectorStoreThreshold   = 0.2
	guideVectorStoreResultLimit = 3
	guideVectorStoreDefaultType = "guide"
	guideNotFoundMessage        = "nothing found in guide store and you need to store it after figure out this case"
)

type guide struct {
	userID            int64
	flowID            int64
	taskID            *int64
	subtaskID         *int64
	replacer          anonymizer.Replacer
	store             *pgvector.Store
	embedder          embeddings.Embedder
	db                database.Querier
	maxEmbeddingBytes int
	vslp              VectorStoreLogProvider
	knp               KnowledgeProvider
}

func NewGuideTool(
	userID int64, flowID int64, taskID, subtaskID *int64,
	replacer anonymizer.Replacer,
	store *pgvector.Store,
	embedder embeddings.Embedder,
	db database.Querier,
	maxEmbeddingBytes int,
	vslp VectorStoreLogProvider,
	knp KnowledgeProvider,
) Tool {
	return &guide{
		userID:            userID,
		flowID:            flowID,
		taskID:            taskID,
		subtaskID:         subtaskID,
		replacer:          replacer,
		store:             store,
		embedder:          embedder,
		db:                db,
		maxEmbeddingBytes: maxEmbeddingBytes,
		vslp:              vslp,
		knp:               knp,
	}
}

func (g *guide) Handle(ctx context.Context, name string, args json.RawMessage) (string, error) {
	if !g.IsAvailable() {
		return "", fmt.Errorf("guide is not available")
	}

	ctx, observation := obs.Observer.NewObservation(ctx)
	logger := logrus.WithContext(ctx).WithFields(enrichLogrusFields(g.flowID, g.taskID, g.subtaskID, logrus.Fields{
		"tool": name,
		"args": string(args),
	}))

	if g.store == nil {
		logger.Error("pgvector store is not initialized")
		return "", fmt.Errorf("pgvector store is not initialized")
	}

	switch name {
	case SearchGuideToolName:
		var action SearchGuideAction
		if err := json.Unmarshal(args, &action); err != nil {
			logger.WithError(err).Error("failed to unmarshal search guide action")
			return "", fmt.Errorf("failed to unmarshal %s search guide action arguments: %w", name, err)
		}

		filters := map[string]any{
			"doc_type":   guideVectorStoreDefaultType,
			"guide_type": action.Type,
		}

		metadata := langfuse.Metadata{
			"tool_name":     name,
			"message":       action.Message,
			"limit":         guideVectorStoreResultLimit,
			"threshold":     guideVectorStoreThreshold,
			"doc_type":      guideVectorStoreDefaultType,
			"guide_type":    action.Type,
			"queries_count": len(action.Questions),
		}

		retriever := observation.Retriever(
			langfuse.WithRetrieverName("retrieve guide from vector store"),
			langfuse.WithRetrieverInput(map[string]any{
				"queries":     action.Questions,
				"threshold":   guideVectorStoreThreshold,
				"max_results": guideVectorStoreResultLimit,
				"filters":     filters,
			}),
			langfuse.WithRetrieverMetadata(metadata),
		)
		ctx, observation = retriever.Observation(ctx)

		logger = logger.WithFields(logrus.Fields{
			"queries_count": len(action.Questions),
			"type":          action.Type,
		})

		var allDocs []schema.Document
		for i, query := range action.Questions {
			queryLogger := logger.WithFields(logrus.Fields{
				"query_index": i + 1,
				"query":       query[:min(len(query), 1000)],
			})

			docs, err := g.store.SimilaritySearch(
				ctx,
				query,
				guideVectorStoreResultLimit,
				vectorstores.WithScoreThreshold(guideVectorStoreThreshold),
				vectorstores.WithFilters(filters),
			)
			if err != nil {
				obs.LogErrorOrCancel(queryLogger, err, "failed to search for similar documents")
				continue // Continue with other queries even if one fails
			}

			queryLogger.WithField("docs_found", len(docs)).Debug("query executed")
			allDocs = append(allDocs, docs...)
		}

		logger.WithFields(logrus.Fields{
			"total_docs_before_dedup": len(allDocs),
		}).Debug("all queries completed")

		docs := MergeAndDeduplicateDocs(allDocs, guideVectorStoreResultLimit)

		logger.WithFields(logrus.Fields{
			"docs_after_dedup": len(docs),
		}).Debug("documents deduplicated and sorted")

		if len(docs) == 0 {
			retriever.End(
				langfuse.WithRetrieverStatus("no guide found"),
				langfuse.WithRetrieverLevel(langfuse.ObservationLevelWarning),
				langfuse.WithRetrieverOutput([]any{}),
			)
			observation.Score(
				langfuse.WithScoreComment("no guide found"),
				langfuse.WithScoreName("guide_search_result"),
				langfuse.WithScoreStringValue("not_found"),
			)
			return guideNotFoundMessage, nil
		}

		retriever.End(
			langfuse.WithRetrieverStatus("success"),
			langfuse.WithRetrieverLevel(langfuse.ObservationLevelDebug),
			langfuse.WithRetrieverOutput(docs),
		)

		buffer := strings.Builder{}
		for i, doc := range docs {
			observation.Score(
				langfuse.WithScoreComment("guide vector store result"),
				langfuse.WithScoreName("guide_search_result"),
				langfuse.WithScoreFloatValue(float64(doc.Score)),
			)
			fmt.Fprintf(&buffer, "# Document %d Match score: %f\n\n", i+1, doc.Score)
			fmt.Fprintf(&buffer, "## Original Guide Type: %s\n\n", doc.Metadata["guide_type"])
			fmt.Fprintf(&buffer, "## Original Guide Question\n\n%s\n\n", doc.Metadata["question"])
			buffer.WriteString("## Content\n\n")
			buffer.WriteString(doc.PageContent)
			buffer.WriteString("\n\n")
		}

		if agentCtx, ok := GetAgentContext(ctx); ok {
			filtersData, err := json.Marshal(filters)
			if err != nil {
				logger.WithError(err).Error("failed to marshal filters")
				return "", fmt.Errorf("failed to marshal filters: %w", err)
			}
			queriesText := g.replacer.ReplaceString(strings.Join(action.Questions, "\n--------------------------------\n"))
			_, _ = g.vslp.PutLog(
				ctx,
				agentCtx.ParentAgentType,
				agentCtx.CurrentAgentType,
				filtersData,
				queriesText,
				database.VecstoreActionTypeRetrieve,
				buffer.String(),
				g.taskID,
				g.subtaskID,
			)
		}

		return buffer.String(), nil

	case StoreGuideToolName:
		var action StoreGuideAction
		if err := json.Unmarshal(args, &action); err != nil {
			logger.WithError(err).Error("failed to unmarshal store guide action")
			return "", fmt.Errorf("failed to unmarshal %s store guide action arguments: %w", name, err)
		}

		action.Question = limits.FillBlankQuestion(action.Question, action.Guide, "Untitled guide")
		if strings.TrimSpace(string(action.Type)) == "" {
			action.Type = GuideType(model.KnowledgeGuideTypeOther)
		}

		guide := fmt.Sprintf("Question:\n%s\n\nGuide:\n%s", action.Question, action.Guide)

		// Anonymize before anything else so all downstream paths (including error
		// branches that emit langfuse events) only ever expose the anonymized form.
		var (
			anonymizedGuide     = boundedContent(g.replacer, guide)
			anonymizedQuestion  = limits.TruncateToLimit(g.replacer.ReplaceString(action.Question), limits.MaxQuestionLen)
			anonymizedGuideOnly = g.replacer.ReplaceString(action.Guide) // used in slow-path embedding text
		)

		eventMetadata := map[string]any{
			"tool_name":  name,
			"message":    action.Message,
			"doc_type":   guideVectorStoreDefaultType,
			"guide_type": action.Type,
		}
		opts := []langfuse.EventOption{
			langfuse.WithEventName("store guide to vector store"),
			langfuse.WithEventInput(action.Question),
			langfuse.WithEventOutput(anonymizedGuide),
			langfuse.WithEventMetadata(eventMetadata),
		}

		logger = logger.WithFields(logrus.Fields{
			"query": action.Question[:min(len(action.Question), 1000)],
			"type":  action.Type,
			"guide": action.Guide[:min(len(action.Guide), 1000)],
		})

		metadata := map[string]any{
			"user_id":    g.userID,
			"flow_id":    g.flowID,
			"doc_type":   guideVectorStoreDefaultType,
			"guide_type": action.Type,
			"question":   anonymizedQuestion,
			"part_size":  len(anonymizedGuide),
			"total_size": len(anonymizedGuide),
		}
		if g.taskID != nil {
			metadata["task_id"] = *g.taskID
		}
		if g.subtaskID != nil {
			metadata["subtask_id"] = *g.subtaskID
		}

		var (
			docs []schema.Document
			ids  []string
			err  error
		)

		if len(anonymizedGuide) <= g.maxEmbeddingBytes || g.embedder == nil {
			// Fast path: document fits within the embedding limit.
			docs, err = documentloaders.NewText(strings.NewReader(anonymizedGuide)).Load(ctx)
			if err != nil {
				observation.Event(append(opts,
					langfuse.WithEventStatus(err.Error()),
					langfuse.WithEventLevel(langfuse.ObservationLevelError),
				)...)
				logger.WithError(err).Error("failed to load document")
				return "", fmt.Errorf("failed to load document: %w", err)
			}
			for i := range docs {
				if docs[i].Metadata == nil {
					docs[i].Metadata = map[string]any{}
				}
				maps.Copy(docs[i].Metadata, metadata)
				docs[i].Metadata["part_size"] = len(docs[i].PageContent)
			}
			ids, err = g.store.AddDocuments(ctx, docs)
			eventMetadata["ids"] = ids
			if err != nil {
				observation.Event(append(opts,
					langfuse.WithEventStatus(err.Error()),
					langfuse.WithEventLevel(langfuse.ObservationLevelError),
				)...)
				logger.WithError(err).Error("failed to store guide")
				return "", fmt.Errorf("failed to store guide: %w", err)
			}
		} else {
			// Slow path: Guide field exceeds embedding limit.
			// Template: "Question:\n{question}\n\nGuide:\n{guide}"
			prefix := "Question:\n" + anonymizedQuestion + "\n\nGuide:\n"
			available := max(g.maxEmbeddingBytes-len(prefix), 0)
			embeddingText := prefix + truncateForEmbedding(anonymizedGuideOnly, available)

			id, err := storeDocumentWithEmbeddingLimit(ctx, g.db, g.embedder,
				embeddingText, anonymizedGuide, metadata)
			if err != nil {
				observation.Event(append(opts,
					langfuse.WithEventStatus(err.Error()),
					langfuse.WithEventLevel(langfuse.ObservationLevelError),
				)...)
				logger.WithError(err).Error("failed to store guide with embedding limit")
				return "", fmt.Errorf("failed to store guide: %w", err)
			}
			ids = []string{id}
			docs = []schema.Document{
				{
					PageContent: anonymizedGuide,
					Metadata:    metadata,
				},
			}
			eventMetadata["ids"] = ids
		}

		observation.Event(append(opts,
			langfuse.WithEventStatus("success"),
			langfuse.WithEventLevel(langfuse.ObservationLevelDebug),
			langfuse.WithEventOutput(docs),
		)...)

		if g.knp != nil {
			guideType := model.KnowledgeGuideType(action.Type)
			for _, id := range ids {
				knDoc := &model.KnowledgeDocument{
					ID:        id,
					UserID:    g.userID,
					DocType:   model.KnowledgeDocTypeGuide,
					Content:   anonymizedGuide,
					Question:  anonymizedQuestion,
					GuideType: &guideType,
					PartSize:  len(anonymizedGuide),
					TotalSize: len(anonymizedGuide),
					Manual:    false,
				}
				if g.flowID != 0 {
					knDoc.FlowID = &g.flowID
				}
				knDoc.TaskID = g.taskID
				knDoc.SubtaskID = g.subtaskID
				g.knp.KnowledgeDocumentCreated(ctx, knDoc)
			}
		}

		if agentCtx, ok := GetAgentContext(ctx); ok {
			data := map[string]any{
				"doc_type":   guideVectorStoreDefaultType,
				"guide_type": action.Type,
			}
			if g.taskID != nil {
				data["task_id"] = *g.taskID
			}
			if g.subtaskID != nil {
				data["subtask_id"] = *g.subtaskID
			}
			filtersData, err := json.Marshal(data)
			if err != nil {
				logger.WithError(err).Error("failed to marshal filters")
				return "", fmt.Errorf("failed to marshal filters: %w", err)
			}
			_, _ = g.vslp.PutLog(
				ctx,
				agentCtx.ParentAgentType,
				agentCtx.CurrentAgentType,
				filtersData,
				anonymizedQuestion,
				database.VecstoreActionTypeStore,
				anonymizedGuide,
				g.taskID,
				g.subtaskID,
			)
		}

		return "guide stored successfully", nil

	default:
		logger.Error("unknown tool")
		return "", fmt.Errorf("unknown tool: %s", name)
	}
}

func (g *guide) IsAvailable() bool {
	return g.store != nil
}