package model

import "sync"

// InMemoryAgentRegistry 基于内存的 Agent 注册表
type InMemoryAgentRegistry struct {
	mu     sync.RWMutex
	agents map[string]RegisteredAgent
}

func NewInMemoryAgentRegistry() *InMemoryAgentRegistry {
	return &InMemoryAgentRegistry{agents: make(map[string]RegisteredAgent)}
}

func (r *InMemoryAgentRegistry) Register(agent RegisteredAgent) error {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.agents[agent.AgentID] = agent
	return nil
}

func (r *InMemoryAgentRegistry) Get(agentID string) (RegisteredAgent, bool) {
	r.mu.RLock()
	defer r.mu.RUnlock()
	agent, ok := r.agents[agentID]
	return agent, ok
}

func (r *InMemoryAgentRegistry) List() []RegisteredAgent {
	r.mu.RLock()
	defer r.mu.RUnlock()
	agents := make([]RegisteredAgent, 0, len(r.agents))
	for _, agent := range r.agents {
		agents = append(agents, agent)
	}
	return agents
}

// InMemorySessionStore 基于内存的会话存储
type InMemorySessionStore struct {
	mu       sync.RWMutex
	sessions map[string]string
}

func NewInMemorySessionStore() *InMemorySessionStore {
	return &InMemorySessionStore{sessions: make(map[string]string)}
}

func (s *InMemorySessionStore) Get(userID, agentID string) (string, bool) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	sessionID, ok := s.sessions[userID+":"+agentID]
	return sessionID, ok
}

func (s *InMemorySessionStore) Set(userID, agentID, sessionID string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.sessions[userID+":"+agentID] = sessionID
	return nil
}