package agent

import (
	"sync"

	"git.bebrik.xyz/Dan4ick/ZovOS_AI/internal/ai"
)

// Memory stores conversation history per session.
type Memory struct {
	mu       sync.RWMutex
	sessions map[string][]ai.Message
}

// NewMemory creates a new Memory store.
func NewMemory() *Memory {
	return &Memory{
		sessions: make(map[string][]ai.Message),
	}
}

// Get returns the message history for a session.
func (m *Memory) Get(sessionID string) []ai.Message {
	m.mu.RLock()
	defer m.mu.RUnlock()
	msgs, ok := m.sessions[sessionID]
	if !ok {
		return nil
	}
	// Return a copy
	result := make([]ai.Message, len(msgs))
	copy(result, msgs)
	return result
}

// Set replaces the message history for a session.
func (m *Memory) Set(sessionID string, messages []ai.Message) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.sessions[sessionID] = messages
}

// Append adds messages to a session's history.
func (m *Memory) Append(sessionID string, msgs ...ai.Message) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.sessions[sessionID] = append(m.sessions[sessionID], msgs...)
}

// Clear removes all messages for a session.
func (m *Memory) Clear(sessionID string) {
	m.mu.Lock()
	defer m.mu.Unlock()
	delete(m.sessions, sessionID)
}

// ListSessions returns all active session IDs.
func (m *Memory) ListSessions() []string {
	m.mu.RLock()
	defer m.mu.RUnlock()
	ids := make([]string, 0, len(m.sessions))
	for id := range m.sessions {
		ids = append(ids, id)
	}
	return ids
}