package com.easyagents.agent.runtime.persistence.session.memory;

import com.easyagents.agent.runtime.persistence.session.AgentSessionStore;
import io.agentscope.core.state.State;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 用于测试和单节点 MVP 的内存会话存储。
 */
public class InMemoryAgentSessionStore implements AgentSessionStore {

    private final Map<String, Map<String, List<State>>> states = new ConcurrentHashMap<>();

    @Override
    public void save(String sessionKey, String name, State state) {
        if (sessionKey == null || name == null || state == null) {
            return;
        }
        saveList(sessionKey, name, List.of(state));
    }

    @Override
    public void saveList(String sessionKey, String name, List<? extends State> states) {
        if (sessionKey == null || name == null) {
            return;
        }
        this.states.computeIfAbsent(sessionKey, key -> new ConcurrentHashMap<>())
                .put(name, states == null ? new ArrayList<>() : new ArrayList<>(states));
    }

    @Override
    public <T extends State> Optional<T> get(String sessionKey, String name, Class<T> type) {
        List<T> list = getList(sessionKey, name, type);
        return list.isEmpty() ? Optional.empty() : Optional.of(list.get(0));
    }

    @Override
    public <T extends State> List<T> getList(String sessionKey, String name, Class<T> itemType) {
        Map<String, List<State>> byName = states.getOrDefault(sessionKey, new LinkedHashMap<>());
        List<State> values = byName.getOrDefault(name, new ArrayList<>());
        if (itemType == null) {
            return new ArrayList<>();
        }
        List<T> typedValues = new ArrayList<>();
        for (State value : values) {
            if (itemType.isInstance(value)) {
                typedValues.add(itemType.cast(value));
            }
        }
        return typedValues;
    }

    @Override
    public boolean exists(String sessionKey) {
        return states.containsKey(sessionKey);
    }

    @Override
    public void delete(String sessionKey) {
        states.remove(sessionKey);
    }

    @Override
    public Set<String> listSessionKeys() {
        return new LinkedHashSet<>(states.keySet());
    }
}
