package tech.easyflow.agent.runtime.session;

import com.easyagents.agent.runtime.persistence.session.AgentSessionStore;
import com.mybatisflex.core.query.QueryWrapper;
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 tech.easyflow.agent.entity.AgentSession;
import tech.easyflow.agent.mapper.AgentSessionMapper;

import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.*;
import java.util.concurrent.TimeUnit;

/**
 * session store 持久化实现
 */
@Service
public class EasyFlowAgentSessionStore implements AgentSessionStore {

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

    private final AgentSessionMapper agentSessionMapper;
    private final StringRedisTemplate stringRedisTemplate;
    private final AgentRuntimeProperties properties;

    /**
     * 创建 EasyFlow Agent session store。
     *
     * @param agentSessionMapper session Mapper
     * @param stringRedisTemplate Redis 模板
     * @param properties Agent 运行态配置
     */
    public EasyFlowAgentSessionStore(AgentSessionMapper agentSessionMapper,
                                     StringRedisTemplate stringRedisTemplate,
                                     AgentRuntimeProperties properties) {
        this.agentSessionMapper = agentSessionMapper;
        this.stringRedisTemplate = stringRedisTemplate;
        this.properties = properties;
    }

    /**
     * 绑定业务会话元信息。
     *
     * <p>AgentScope Session API 不会把 EasyFlow 的 {@code agentId/chatSessionId/tenantId} 传入
     * {@code save(...)}，因此运行入口必须先调用本方法建立或刷新元信息。后续 state 写入会复用
     * 这些字段，避免表里只剩裸 sessionKey。</p>
     *
     * @param metadata 业务会话元信息
     */
    public void bindSession(AgentSessionMetadata metadata) {
        if (metadata == null || !StringUtils.hasText(metadata.sessionKey())) {
            return;
        }
        AgentSession session = findBySessionKey(metadata.sessionKey());
        Date now = new Date();
        if (session == null) {
            session = new AgentSession();
            session.setTenantId(metadata.tenantId());
            session.setAgentId(metadata.agentId());
            session.setChatSessionId(metadata.chatSessionId());
            session.setRuntimeSessionId(metadata.runtimeSessionId());
            session.setSessionKey(metadata.sessionKey());
            session.setStateJson(emptyEnvelope());
            session.setVersion(0L);
            session.setCacheVersion(0L);
            session.setCreated(now);
            session.setCreatedBy(metadata.operatorId());
            session.setIsDeleted(0);
        }
        session.setTenantId(firstNonNull(metadata.tenantId(), session.getTenantId()));
        session.setAgentId(firstNonNull(metadata.agentId(), session.getAgentId()));
        session.setChatSessionId(firstNonNull(metadata.chatSessionId(), session.getChatSessionId()));
        session.setRuntimeSessionId(firstText(metadata.runtimeSessionId(), session.getRuntimeSessionId()));
        session.setLastAccessAt(now);
        session.setModified(now);
        session.setModifiedBy(metadata.operatorId());
        agentSessionMapper.insertOrUpdate(session);
        writeCache(session.getSessionKey(), session.getStateJson());
    }

    /**
     * 按聊天会话清理 AgentScope session。
     *
     * @param chatSessionId 聊天会话 ID
     */
    public void deleteByChatSessionId(BigInteger chatSessionId) {
        if (chatSessionId == null) {
            return;
        }
        List<AgentSession> sessions = agentSessionMapper.selectListByQuery(QueryWrapper.create()
                .eq("chat_session_id", chatSessionId)
                .eq("is_deleted", 0));
        Date now = new Date();
        for (AgentSession session : sessions) {
            session.setIsDeleted(1);
            session.setModified(now);
            agentSessionMapper.update(session);
            deleteCache(session.getSessionKey());
        }
    }

    @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));
        persistEnvelope(sessionKey, envelope);
    }

    @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);
        persistEnvelope(sessionKey, envelope);
    }

    @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));
    }

    @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;
    }

    @Override
    public boolean exists(String sessionKey) {
        if (!StringUtils.hasText(sessionKey)) {
            return false;
        }
        if (readCache(sessionKey) != null) {
            return true;
        }
        return findBySessionKey(sessionKey) != null;
    }

    @Override
    public void delete(String sessionKey) {
        if (!StringUtils.hasText(sessionKey)) {
            return;
        }
        AgentSession session = findBySessionKey(sessionKey);
        if (session != null) {
            session.setIsDeleted(1);
            session.setModified(new Date());
            agentSessionMapper.update(session);
        }
        deleteCache(sessionKey);
    }

    @Override
    public Set<String> listSessionKeys() {
        List<AgentSession> sessions = agentSessionMapper.selectListByQuery(QueryWrapper.create()
                .eq("is_deleted", 0)
                .select("session_key"));
        Set<String> keys = new LinkedHashSet<>();
        for (AgentSession session : sessions) {
            if (StringUtils.hasText(session.getSessionKey())) {
                keys.add(session.getSessionKey());
            }
        }
        return keys;
    }

    private void persistEnvelope(String sessionKey, Map<String, Object> envelope) {
        AgentSession session = findBySessionKey(sessionKey);
        Date now = new Date();
        if (session == null) {
            session = new AgentSession();
            session.setRuntimeSessionId(sessionKey);
            session.setSessionKey(sessionKey);
            session.setCreated(now);
            session.setVersion(0L);
            session.setCacheVersion(0L);
            session.setIsDeleted(0);
        }
        long nextVersion = session.getVersion() == null ? 1L : session.getVersion() + 1L;
        session.setStateJson(envelope);
        session.setVersion(nextVersion);
        session.setCacheVersion(nextVersion);
        session.setLastAccessAt(now);
        session.setModified(now);
        session.setIsDeleted(0);
        // 同步写会话状态
        agentSessionMapper.insertOrUpdate(session);
        // 同步写缓存
        writeCache(sessionKey, envelope);
    }

    private Map<String, Object> loadEnvelope(String sessionKey) {
        Map<String, Object> cached = readCache(sessionKey);
        if (cached != null) {
            return cached;
        }
        AgentSession session = findBySessionKey(sessionKey);
        if (session == null || session.getStateJson() == null || session.getStateJson().isEmpty()) {
            return emptyEnvelope();
        }
        writeCache(sessionKey, session.getStateJson());
        return deepCopy(session.getStateJson());
    }

    private AgentSession findBySessionKey(String sessionKey) {
        return agentSessionMapper.selectOneByQuery(QueryWrapper.create()
                .eq("session_key", sessionKey)
                .eq("is_deleted", 0)
                .limit(1));
    }

    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) {
        try {
            long seconds = Math.max(1L, properties.getSessionCacheTtl().toSeconds());
            stringRedisTemplate.opsForValue().set(cacheKey(sessionKey), JsonUtils.getJsonCodec().toJson(envelope),
                    seconds, TimeUnit.SECONDS);
        } catch (RuntimeException e) {
            // MySQL 是 AgentScope session 的真相源。Redis 只做热缓存，写缓存失败不能掩盖已完成的持久化写入。
        }
    }

    private void deleteCache(String sessionKey) {
        try {
            stringRedisTemplate.delete(cacheKey(sessionKey));
        } catch (RuntimeException ignored) {
            // 清理缓存失败不应掩盖 MySQL 删除结果，后续 TTL 会自然回收。
        }
    }

    @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(':', '_');
        }
    }

    private <T> T firstNonNull(T left, T right) {
        return left == null ? right : left;
    }

    private String firstText(String left, String right) {
        return StringUtils.hasText(left) ? left : right;
    }

    /**
     * AgentScope session 绑定的 EasyFlow 业务元信息。
     *
     * @param sessionKey AgentScope session key
     * @param runtimeSessionId runtime session ID
     * @param agentId Agent ID
     * @param chatSessionId chatlog 会话 ID
     * @param tenantId 租户 ID
     * @param operatorId 操作人 ID
     */
    public record AgentSessionMetadata(String sessionKey,
                                       String runtimeSessionId,
                                       BigInteger agentId,
                                       BigInteger chatSessionId,
                                       BigInteger tenantId,
                                       BigInteger operatorId) {
    }
}
