package tech.easyflow.agent.runtime.session;

import com.easyagents.agent.runtime.persistence.session.AgentSessionStore;
import io.agentscope.core.state.State;
import io.agentscope.core.util.JsonUtils;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Service;
import org.springframework.util.StringUtils;
import tech.easyflow.agent.config.AgentRuntimeProperties;

import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.ArrayList;
import java.util.Base64;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.TimeUnit;

/**
 * Agent 草稿试运行 Redis-only session store。
 */
@Service
public class DraftAgentSessionStore implements AgentSessionStore {

    private static final String REDIS_PREFIX = "easyflow:agent:draft-session:";
    private static final String ENVELOPE_VERSION = "1";
    private static final String SINGLE_STATES = "singleStates";
    private static final String LIST_STATES = "listStates";

    private final StringRedisTemplate stringRedisTemplate;
    private final AgentRuntimeProperties properties;

    /**
     * 创建草稿试运行 session store。
     *
     * @param stringRedisTemplate Redis 模板
     * @param properties Agent 运行态配置
     */
    public DraftAgentSessionStore(StringRedisTemplate stringRedisTemplate,
                                  AgentRuntimeProperties properties) {
        this.stringRedisTemplate = stringRedisTemplate;
        this.properties = properties;
    }

    /**
     * 保存单个状态项。
     *
     * @param sessionKey 会话键
     * @param name 状态名称
     * @param state 状态值
     */
    @Override
    public void save(String sessionKey, String name, State state) {
        if (!StringUtils.hasText(sessionKey) || !StringUtils.hasText(name) || state == null) {
            return;
        }
        Map<String, Object> envelope = loadEnvelope(sessionKey);
        singleStates(envelope).put(name, JsonUtils.getJsonCodec().toJson(state));
        writeCache(sessionKey, envelope);
    }

    /**
     * 保存状态列表。
     *
     * @param sessionKey 会话键
     * @param name 状态名称
     * @param states 状态列表
     */
    @Override
    public void saveList(String sessionKey, String name, List<? extends State> states) {
        if (!StringUtils.hasText(sessionKey) || !StringUtils.hasText(name)) {
            return;
        }
        List<String> values = new ArrayList<>();
        if (states != null) {
            for (State state : states) {
                values.add(JsonUtils.getJsonCodec().toJson(state));
            }
        }
        Map<String, Object> envelope = loadEnvelope(sessionKey);
        listStates(envelope).put(name, values);
        writeCache(sessionKey, envelope);
    }

    /**
     * 获取单个状态项。
     *
     * @param sessionKey 会话键
     * @param name 状态名称
     * @param type 状态类型
     * @param <T> 状态类型
     * @return 可选状态
     */
    @Override
    public <T extends State> Optional<T> get(String sessionKey, String name, Class<T> type) {
        if (!StringUtils.hasText(sessionKey) || !StringUtils.hasText(name) || type == null) {
            return Optional.empty();
        }
        Object json = singleStates(loadEnvelope(sessionKey)).get(name);
        if (!(json instanceof String text) || text.isBlank()) {
            return Optional.empty();
        }
        return Optional.of(JsonUtils.getJsonCodec().fromJson(text, type));
    }

    /**
     * 获取状态列表。
     *
     * @param sessionKey 会话键
     * @param name 状态名称
     * @param itemType 状态元素类型
     * @param <T> 状态元素类型
     * @return 状态列表
     */
    @Override
    public <T extends State> List<T> getList(String sessionKey, String name, Class<T> itemType) {
        if (!StringUtils.hasText(sessionKey) || !StringUtils.hasText(name) || itemType == null) {
            return List.of();
        }
        Object raw = listStates(loadEnvelope(sessionKey)).get(name);
        if (!(raw instanceof List<?> values) || values.isEmpty()) {
            return List.of();
        }
        List<T> result = new ArrayList<>();
        for (Object value : values) {
            if (value instanceof String text && !text.isBlank()) {
                result.add(JsonUtils.getJsonCodec().fromJson(text, itemType));
            }
        }
        return result;
    }

    /**
     * 判断会话键是否存在。
     *
     * @param sessionKey 会话键
     * @return 存在时为 true
     */
    @Override
    public boolean exists(String sessionKey) {
        return StringUtils.hasText(sessionKey) && readCache(sessionKey) != null;
    }

    /**
     * 删除指定会话键下的全部状态。
     *
     * @param sessionKey 会话键
     */
    @Override
    public void delete(String sessionKey) {
        if (!StringUtils.hasText(sessionKey)) {
            return;
        }
        deleteCache(sessionKey);
    }

    /**
     * 列出当前存储中的会话键。
     *
     * <p>草稿 session 使用哈希 Redis key，不维护反向索引，避免为试运行引入额外持久化状态。</p>
     *
     * @return 空集合
     */
    @Override
    public Set<String> listSessionKeys() {
        return new LinkedHashSet<>();
    }

    private Map<String, Object> loadEnvelope(String sessionKey) {
        Map<String, Object> cached = readCache(sessionKey);
        return cached == null ? emptyEnvelope() : deepCopy(cached);
    }

    private Map<String, Object> emptyEnvelope() {
        Map<String, Object> envelope = new LinkedHashMap<>();
        envelope.put("version", ENVELOPE_VERSION);
        envelope.put(SINGLE_STATES, new LinkedHashMap<String, Object>());
        envelope.put(LIST_STATES, new LinkedHashMap<String, Object>());
        return envelope;
    }

    @SuppressWarnings("unchecked")
    private Map<String, Object> singleStates(Map<String, Object> envelope) {
        return (Map<String, Object>) envelope.computeIfAbsent(SINGLE_STATES, key -> new LinkedHashMap<String, Object>());
    }

    @SuppressWarnings("unchecked")
    private Map<String, Object> listStates(Map<String, Object> envelope) {
        return (Map<String, Object>) envelope.computeIfAbsent(LIST_STATES, key -> new LinkedHashMap<String, Object>());
    }

    @SuppressWarnings("unchecked")
    private Map<String, Object> readCache(String sessionKey) {
        try {
            String value = stringRedisTemplate.opsForValue().get(cacheKey(sessionKey));
            if (!StringUtils.hasText(value)) {
                return null;
            }
            return JsonUtils.getJsonCodec().fromJson(value, Map.class);
        } catch (RuntimeException e) {
            return null;
        }
    }

    private void writeCache(String sessionKey, Map<String, Object> envelope) {
        long seconds = Math.max(1L, properties.getSessionCacheTtl().toSeconds());
        stringRedisTemplate.opsForValue().set(cacheKey(sessionKey), JsonUtils.getJsonCodec().toJson(envelope),
                seconds, TimeUnit.SECONDS);
    }

    private void deleteCache(String sessionKey) {
        stringRedisTemplate.delete(cacheKey(sessionKey));
    }

    @SuppressWarnings("unchecked")
    private Map<String, Object> deepCopy(Map<String, Object> source) {
        if (source == null || source.isEmpty()) {
            return emptyEnvelope();
        }
        return JsonUtils.getJsonCodec().fromJson(JsonUtils.getJsonCodec().toJson(source), Map.class);
    }

    private String cacheKey(String sessionKey) {
        return REDIS_PREFIX + hash(sessionKey);
    }

    private String hash(String value) {
        try {
            MessageDigest digest = MessageDigest.getInstance("SHA-256");
            byte[] bytes = digest.digest(value.getBytes(StandardCharsets.UTF_8));
            return Base64.getUrlEncoder().withoutPadding().encodeToString(bytes);
        } catch (NoSuchAlgorithmException e) {
            return value.replace(':', '_');
        }
    }
}
