// /internal/agents/agent_manager.go

package agents

import (
	"fmt"
	"sync"

	"github.com/blog/conversational-agent/internal/logger"
	"github.com/tmc/langchaingo/chains"
	"github.com/tmc/langchaingo/llms/openai"
	"github.com/tmc/langchaingo/memory"
	"github.com/tmc/langchaingo/schema"
	"github.com/tmc/langchaingo/vectorstores/weaviate"
)

type AgentManager struct {
	LLM                 *openai.LLM
	VectorStore         weaviate.Store
	AgentMemory         map[string]*memory.ConversationBuffer
	ConversationalChain *chains.ConversationalRetrievalQA
	WeaviateIndex       string
	LLMChain            *chains.LLMChain
	messageBuffer       []schema.Document
	bufferMutex         sync.Mutex
	maxBufferMessages   int
	memoryMutex         sync.Mutex
}

func NewAgentManager(
	openAIApiKey, openAIModel, weaviateHost, weaviateApiKey, weaviateIndex string,
	maxBufferMessages int,
) (*AgentManager, error) {
	log := logger.GetLogger()

	log.Info().Msgf("Initializing OpenAI LLM with Model: %s", openAIModel)
	llm, err := InitializeLLM(openAIApiKey, openAIModel)
	if err != nil {
		return nil, fmt.Errorf("failed to initialize OpenAI LLM: %w", err)
	}
	log.Info().Msg("OpenAI LLM initialized successfully")

	log.Info().Msg("Initializing OpenAI Embedder...")
	embedder, err := InitializeEmbedder(llm)
	if err != nil {
		return nil, fmt.Errorf("failed to initialize OpenAI embedder: %w", err)
	}
	log.Info().Msg("OpenAI Embedder initialized successfully")

	log.Info().Msg("Initializing Weaviate vector store...")
	vectorStore, err := InitializeVectorStore(weaviateHost, weaviateApiKey, weaviateIndex, embedder)
	if err != nil {
		return nil, fmt.Errorf("failed to initialize Weaviate vector store: %w", err)
	}

	log.Info().Msg("Weaviate vector store initialized successfully")

	log.Info().Msg("Initializing ConversationBuffer memory...")
	agentMemory := InitializeMemory()
	log.Info().Msg("ConversationBuffer memory initialized successfully")

	log.Info().Msg("Initializing Chain...")
	chain, err := InitializeChain(llm, agentMemory)
	if err != nil {
		return nil, fmt.Errorf("failed to initialize chain: %w", err)
	}
	log.Info().Msg("Chain initialized successfully")

	return &AgentManager{
		LLM:               llm,
		VectorStore:       vectorStore,
		AgentMemory:       make(map[string]*memory.ConversationBuffer),
		LLMChain:          chain,
		WeaviateIndex:     weaviateIndex,
		messageBuffer:     []schema.Document{},
		maxBufferMessages: maxBufferMessages,
	}, nil
}

// /internal/agents/agent_utils.go

package agents

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/blog/conversational-agent/internal/logger"
	"github.com/tmc/langchaingo/llms"
	"github.com/tmc/langchaingo/memory"
	"github.com/tmc/langchaingo/schema"
	"github.com/tmc/langchaingo/vectorstores"
	"github.com/weaviate/weaviate-go-client/v4/weaviate/filters"
)

// GetThreadMemory retrieves or initializes a ConversationBuffer for the given thread ID.
func (am *AgentManager) GetThreadMemory(threadID string) *memory.ConversationBuffer {
	am.memoryMutex.Lock()
	defer am.memoryMutex.Unlock()

	// Check if the memory for the thread already exists
	if mem, exists := am.AgentMemory[threadID]; exists {
		return mem
	}

	// Initialize a new memory buffer
	newMemory := memory.NewConversationBuffer(
		memory.WithMemoryKey(fmt.Sprintf("thread:%s", threadID)),
		memory.WithReturnMessages(false),
	)
	am.AgentMemory[threadID] = newMemory
	return newMemory
}

// RetrieveMemory retrieves the chat history for a specific thread
func (am *AgentManager) RetrieveMemory(ctx context.Context, threadID string) ([]map[string]string, error) {
	log := logger.GetLogger()
	log.Debug().Msgf("Starting RetrieveMemory function for thread ID: %s", threadID)

	// Get the specific memory for the thread
	threadMemory := am.GetThreadMemory(threadID)
	if threadMemory == nil {
		log.Debug().Msgf("AgentMemory for thread ID '%s' is not initialized.", threadID)
		return nil, fmt.Errorf("memory is not initialized for thread ID: %s", threadID)
	}

	// Retrieve messages from the chat history
	messages, err := threadMemory.ChatHistory.Messages(ctx)
	if err != nil {
		log.Error().Err(err).Msg("Failed to retrieve messages from chat history.")
		return nil, fmt.Errorf("failed to retrieve messages for thread_id %s: %w", threadID, err)
	}

	// Log retrieved messages
	log.Debug().Msgf("Retrieved messages for thread ID %s: %+v", threadID, messages)

	// Format the messages for output
	formattedMessages := formatMessages(messages)
	return formattedMessages, nil
}

// formatMessages formats the messages for output
func formatMessages(messages []llms.ChatMessage) []map[string]string {
	var formattedMessages []map[string]string
	for _, msg := range messages {
		role := "user"
		if msg.GetType() == llms.ChatMessageTypeAI {
			role = "ai"
		}
		formattedMessages = append(formattedMessages, map[string]string{
			"role":    role,
			"content": msg.GetContent(),
		})
	}
	return formattedMessages
}

// addToBuffer adds the input and response to the message buffer
func (am *AgentManager) addToBuffer(threadID, input, response, userID, orgID string) {
	am.bufferMutex.Lock()
	defer am.bufferMutex.Unlock()
        // Chunk the input and response for fine-grained storage
	chunks := ChunkContent(input+"\n"+response, 300) // Adjust chunk size
	for _, chunk := range chunks {
		doc := schema.Document{
			PageContent: chunk,
			Metadata: map[string]any{
				"thread_id": threadID,
				"user_id":   userID,
				"org_id":    orgID,
				"timestamp": time.Now().Format(time.RFC3339),
				"source":    "conversation",
			},
		}

		// Ensure uniqueness before adding
		am.messageBuffer = append(am.messageBuffer, doc)
	}

	// Flush the buffer if it exceeds max size
	if len(am.messageBuffer) >= am.maxBufferMessages {
		go am.flushBufferToWeaviate(context.Background())
	}
}

// flushBufferToWeaviate flushes the message buffer to the vector store
func (am *AgentManager) flushBufferToWeaviate(ctx context.Context) {
	am.bufferMutex.Lock()
	buffer := am.messageBuffer
	am.messageBuffer = nil
	am.bufferMutex.Unlock()

	if len(buffer) == 0 {
		return
	}

	// Deduplicate buffer entries before inserting into Weaviate
	uniqueDocs := deduplicateDocuments(buffer)

	log := logger.GetLogger()
	log.Debug().Msgf("Batch inserting %d unique messages into Weaviate...", len(uniqueDocs))

	// Assuming each document has org_id in its metadata
	if len(uniqueDocs) > 0 {
		orgID, ok := uniqueDocs[0].Metadata["org_id"].(string)
		if !ok {
			log.Error().Msg("Failed to retrieve org_id from document metadata.")
			return
		}

		namespace := orgID
		log.Debug().Msgf("Using namespace: %s for flushing buffer to Weaviate", 
namespace)

		_, err := am.VectorStore.AddDocuments(ctx, uniqueDocs, vectorstores.WithNameSpace(namespace))
		if err != nil {
			log.Error().Err(err).Msg("Failed to batch insert messages into Weaviate.")
			return
		}
	}

	log.Debug().Msg("Batch insertion to Weaviate completed successfully.")
}

// deduplicateDocuments removes duplicate documents from a slice based on PageContent.
func deduplicateDocuments(docs []schema.Document) []schema.Document {
	seen := make(map[string]bool)
	uniqueDocs := []schema.Document{}

	for _, doc := range docs {
		if _, exists := seen[doc.PageContent]; !exists {
			seen[doc.PageContent] = true
			uniqueDocs = append(uniqueDocs, doc)
		}
	}

	return uniqueDocs
}

// SyncMemory syncs the message buffer to the vector store
func (am *AgentManager) SyncMemory(ctx context.Context) error {
	log := logger.GetLogger()

	am.bufferMutex.Lock()
	defer am.bufferMutex.Unlock()

	if len(am.messageBuffer) == 0 {
		log.Debug().Msg("No messages to sync from buffer.")
		return nil
	}

	log.Debug().Msgf("Syncing %d messages from buffer to Weaviate...", len(am.messageBuffer))
	_, err := am.VectorStore.AddDocuments(ctx, am.messageBuffer)
	if err != nil {
		log.Error().Err(err).Msg("Failed to sync messages to Weaviate.")
		return err
	}

	am.messageBuffer = nil
	log.Debug().Msg("Memory synced successfully to Weaviate.")
	return nil
}

// AddDocuments adds documents to the vector store
func (am *AgentManager) AddDocuments(ctx context.Context, docs []schema.Document, userID, orgID string) error {
	log := logger.GetLogger()

	// Add user_id to document metadata
	for i := range docs {
		docs[i].Metadata["user_id"] = userID
	}

	// Use org_id as namespace
	namespace := orgID
	log.Debug().Msgf("Using namespace: %s for adding documents", namespace)

	// Retry mechanism for adding documents
	maxRetries := 3
	var err error

	for attempt := 1; attempt <= maxRetries; attempt++ {
		_, err = am.VectorStore.AddDocuments(ctx, docs, vectorstores.WithNameSpace(namespace))
		if err != nil {
			log.Warn().Err(err).Msgf("Attempt %d to add documents to vector store failed.", attempt)
			if attempt < maxRetries {
				continue
			}
		} else {
			log.Info().Msgf("Documents added to vector store successfully on attempt %d.", attempt)
			return nil
		}
	}

	// Log final failure after retries
	log.Error().Err(err).Msg("Failed to add documents to vector store after multiple attempts.")
	return fmt.Errorf("failed to add documents to vector store: %w", err)
}

// AddDataset adds a dataset to the vector store
func (am *AgentManager) AddDataset(ctx context.Context, threadID, filePath string) error {
	file, err := os.Open(filePath)
	if err != nil {
		return fmt.Errorf("failed to open file: %w", err)
	}
	defer file.Close()

	// Parse the dataset
	var documents []schema.Document
	if err := json.NewDecoder(file).Decode(&documents); err != nil {
		return fmt.Errorf("failed to parse JSON dataset: %w", err)
	}

	// Retrieve the thread-specific memory
	threadMemory := am.GetThreadMemory(threadID)
	if threadMemory == nil {
		return fmt.Errorf("failed to retrieve memory for thread ID: %s", threadID)
	}

	// Add each document to the memory
	for _, doc := range documents {
		err := threadMemory.ChatHistory.AddAIMessage(ctx, doc.PageContent)
		if err != nil {
			return fmt.Errorf("failed to add document to memory for thread ID %s: %w", threadID, err)
		}
	}

	return nil
}

// Chunk content into smaller parts based on word count
func ChunkContent(content string, maxWords int) []string {
	words := strings.Fields(content)
	var chunks []string

	for i := 0; i < len(words); i += maxWords {
		end := i + maxWords
		if end > len(words) {
			end = len(words)
		}
		chunks = append(chunks, strings.Join(words[i:end], " "))
	}

	return chunks
}

// QueryMemoryByUserAndOrgID queries the memory for a specific user and organization
func (am *AgentManager) QueryMemoryByUserAndOrgID(ctx context.Context, userID, orgID string) ([]schema.Document, error) {
	log := logger.GetLogger()
	log.Debug().Msgf("Querying memory for user_id: %s and org_id: %s", userID, orgID)

	// Build metadata filters using filters.Where
	whereBuilder := filters.Where().
		WithOperator(filters.And).
		WithOperands([]*filters.WhereBuilder{
			filters.Where().WithPath([]string{"user_id"}).WithOperator(filters.Equal).WithValueString(userID),
	
filters.Where().WithPath([]string{"org_id"}).WithOperator(filters.Equal).WithValueString(orgID),
		})

	// Execute MetadataSearch with the constructed filter
	docs, err := am.VectorStore.MetadataSearch(ctx, 10, vectorstores.WithFilters(whereBuilder))
	if err != nil {
		log.Error().Err(err).Msg("Failed to query memory by user_id and org_id.")
		return nil, fmt.Errorf("failed to query memory for user_id %s and org_id %s: %w", userID, orgID, err)
	}

	log.Debug().Msgf("Retrieved %d documents for user_id: %s and org_id: %s", len(docs), userID, orgID)
	return docs, nil
}

// ConvertContentToString converts content to a string
func ConvertContentToString(content interface{}) (string, error) {
	switch v := content.(type) {
	case string:
		return v, nil
	case map[string]interface{}, []interface{}:
		// Serialize JSON content to a string
		jsonContent, err := json.Marshal(v)
		if err != nil {
			return "", fmt.Errorf("failed to serialize JSON content: %w", err)
		}
		return string(jsonContent), nil
	default:
		return "", fmt.Errorf("unsupported content type: %T", content)
	}
}

// Convert metadata values
func ConvertMetadata(metadata map[string]string) map[string]any {
	converted := make(map[string]any)
	for k, v := range metadata {
		converted[k] = v
	}
	return converted
}		

// internal/handlers/memory_handler.go

package handlers

import (
	"context"
	"encoding/json"
	"net/http"
	"os"
	"github.com/blog/conversational-agent/internal/agents"
	"github.com/blog/conversational-agent/internal/logger"
	"github.com/labstack/echo/v4"
	"github.com/tmc/langchaingo/schema"
)

type ImportDatasetRequest struct {
	FilePath string `json:"file_path"`
}

// GetMemoryHandler retrieves the chat history for a specific thread.
func (h *AgentHandler) GetMemoryHandler(c echo.Context) error {
	ctx := c.Request().Context()

	// Get threadID from the request (e.g., query parameter or header)
	threadID := c.Param("thread_id")
	if threadID == "" {
		return c.JSON(http.StatusBadRequest, map[string]string{"error": "Missing thread_id"})
	}

	log := logger.GetLogger()

	// Retrieve memory from AgentManager
	memory, err := h.AgentManager.RetrieveMemory(ctx, threadID)
	if err != nil {
		return c.JSON(http.StatusInternalServerError, map[string]string{"error": err.Error()})
	}

	if len(memory) == 0 {
		log.Debug().Msgf("Handler: No memory found for thread ID: %s", threadID)
		return c.JSON(http.StatusOK, map[string]string{"memory": "Memory is empty"})
	}

	return c.JSON(http.StatusOK, map[string]any{"memory": memory})
}

// AddDocumentHandler processes and stores documents with chunking and metadata.
func (h *AgentHandler) AddDocumentHandler(c echo.Context) error {
	type AddDocumentRequest struct {
		Content  interface{}       `json:"page_content"` // Accepts string or JSON
		Metadata map[string]string `json:"metadata"`
	}

	var req AddDocumentRequest
	if err := c.Bind(&req); err != nil {
		return c.JSON(http.StatusBadRequest, map[string]string{"error": "Invalid request payload"})
	}

	// Convert Content to a string
	content, err := agents.ConvertContentToString(req.Content)
	if err != nil {
		return c.JSON(http.StatusBadRequest, map[string]string{"error": err.Error()})
	}

	// Create a document to add
	doc := schema.Document{
		PageContent: content,
		Metadata:    agents.ConvertMetadata(req.Metadata),
	}

	// Add document to Weaviate vector store
	ctx := c.Request().Context()
	docID, err := h.AgentManager.VectorStore.AddDocuments(ctx, []schema.Document{doc})
	if err != nil {
		return c.JSON(http.StatusInternalServerError, map[string]string{"error": "Failed to add document"})
	}

	return c.JSON(http.StatusOK, map[string]string{"doc_id": docID[0]})
}

// ImportMemoryHandler imports a dataset into the vector store.
func (h *AgentHandler) ImportMemoryHandler(c echo.Context) error {
	type ImportRequest struct {
		FilePath string `json:"file_path"`
		UserID   string `json:"user_id"`
		OrgID    string `json:"org_id"`
	}

	var req ImportRequest
	if err := c.Bind(&req); err != nil || req.FilePath == "" || req.UserID == "" || req.OrgID == "" {
		return c.JSON(http.StatusBadRequest, map[string]string{
			"error": "Invalid request payload. 'file_path', 'user_id', and 'org_id' are required.",
		})
	}

	// Open the JSON file
	file, err := os.Open(req.FilePath)
	if err != nil {
		return c.JSON(http.StatusInternalServerError, map[string]string{"error": "Failed to open the file."})
	}
	defer file.Close()

	// Parse the JSON file
	var dataset []map[string]interface{}
	if err := json.NewDecoder(file).Decode(&dataset); err != nil {
		return c.JSON(http.StatusBadRequest, map[string]string{"error": "Invalid JSON file format."})
	}

	// Convert dataset into schema.Document format
	docs := []schema.Document{}
	for _, item := range dataset {
		if summary, ok := item["summary"].(string); ok {
			doc := schema.Document{
				PageContent: summary,
				Metadata:    item,
			}
			docs = append(docs, doc)
		}
	}

	// Add documents to vector store with userID and orgID
	ctx := context.Background()
	if err := h.AgentManager.AddDocuments(ctx, docs, req.UserID, req.OrgID); err != nil {
		return c.JSON(http.StatusInternalServerError, map[string]string{"error": "Failed to add documents to vector store."})
	}

	return c.JSON(http.StatusOK, map[string]string{"message": "Dataset imported successfully."})
}