/**
 * @author: dn-jinmin/dn-jinmin
 * @doc:
 */

package memoryx

import (
	"context"
	"github.com/tmc/langchaingo/schema"
	"imooc.com/aiworkc/pkg/langchain"
	"sync"
)

type Memoryx struct {
	sync.Mutex
	getMemory     func() schema.Memory
	memorys       map[string]schema.Memory
	defaultMemory schema.Memory
}

func NewMemoryx(handle func() schema.Memory) *Memoryx {
	return &Memoryx{
		getMemory:     handle,
		memorys:       make(map[string]schema.Memory),
		defaultMemory: handle(),
	}
}

// GetMemoryKey getter for memory key.
func (s *Memoryx) GetMemoryKey(ctx context.Context) string {
	return s.memory(ctx).GetMemoryKey(ctx)
}

// MemoryVariables Input keys this memory class will load dynamically.
func (s *Memoryx) MemoryVariables(ctx context.Context) []string {
	return s.memory(ctx).MemoryVariables(ctx)
}

// LoadMemoryVariables Return key-value pairs given the text input to the chain.
// If None, return all memories
func (s *Memoryx) LoadMemoryVariables(ctx context.Context, inputs map[string]any) (map[string]any, error) {
	return s.memory(ctx).LoadMemoryVariables(ctx, inputs)
}

// SaveContext Save the context of this model run to memory.
func (s *Memoryx) SaveContext(ctx context.Context, inputs map[string]any, outputs map[string]any) error {
	return s.memory(ctx).SaveContext(ctx, inputs, outputs)
}

// Clear memory contents.
func (s *Memoryx) Clear(ctx context.Context) error {
	return s.memory(ctx).Clear(ctx)
}

func (s *Memoryx) memory(ctx context.Context) schema.Memory {
	s.Lock()
	defer s.Unlock()

	var chatId string
	v := ctx.Value(langchain.ChatId)
	if v == nil {
		return s.defaultMemory
	}

	chatId = v.(string)
	memory, ok := s.memorys[chatId]
	if !ok {
		memory = s.getMemory()
		s.memorys[chatId] = memory
	}

	return memory
}