package work.slhaf.partner.core.memory;

import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import work.slhaf.partner.core.memory.pojo.MemorySlice;
import work.slhaf.partner.core.memory.pojo.MemorySliceSnapshot;
import work.slhaf.partner.core.memory.pojo.MemoryUnit;
import work.slhaf.partner.core.memory.pojo.MemoryUnitSnapshot;
import work.slhaf.partner.framework.agent.factory.capability.annotation.CapabilityCore;
import work.slhaf.partner.framework.agent.factory.capability.annotation.CapabilityMethod;
import work.slhaf.partner.framework.agent.model.pojo.Message;
import work.slhaf.partner.framework.agent.state.State;
import work.slhaf.partner.framework.agent.state.StateSerializable;
import work.slhaf.partner.framework.agent.state.StateValue;
import work.slhaf.partner.framework.agent.support.Result;
import work.slhaf.partner.module.memory.runtime.exception.MemoryLookupException;

import java.nio.file.Path;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

@CapabilityCore(value = "memory")
@Slf4j
public class MemoryCore implements StateSerializable {

    private final Lock memoryLock = new ReentrantLock();
    private final ConcurrentHashMap<String, MemoryUnit> memoryUnits = new ConcurrentHashMap<>();

    // 默认值一般只存在于智能体初次启动时
    private String memorySessionId = UUID.randomUUID().toString();

    public MemoryCore() {
        register();
    }

    @CapabilityMethod
    public MemoryUnitSnapshot updateMemoryUnit(List<Message> chatMessages, String summary) {
        memoryLock.lock();
        try {
            MemoryUnit unit = getOrLoadMemoryUnit(memorySessionId);
            unit.updateTimestamp();

            List<Message> conversationMessages = unit.getConversationMessages();
            int startIndex = conversationMessages.size();

            MemorySlice memorySlice = new MemorySlice(
                    startIndex,
                    startIndex + chatMessages.size(),
                    summary
            );

            conversationMessages.addAll(chatMessages);

            unit.getSlices().add(memorySlice);
            normalizeMemoryUnit(unit);
            return unit.snapshot();
        } finally {
            memoryLock.unlock();
        }
    }

    @CapabilityMethod
    public MemoryUnitSnapshot getMemoryUnit(String unitId) {
        memoryLock.lock();
        try {
            MemoryUnit unit = getOrLoadMemoryUnit(unitId);
            normalizeMemoryUnit(unit);
            return unit.snapshot();
        } finally {
            memoryLock.unlock();
        }
    }

    @CapabilityMethod
    public Result<MemorySliceSnapshot> getMemorySlice(String unitId, String sliceId) {
        memoryLock.lock();
        try {
            MemoryUnit memoryUnit = memoryUnits.get(unitId);
            if (memoryUnit == null) {
                return memorySliceNotFound(unitId, sliceId);
            }
            memoryUnit.load();
            normalizeMemoryUnit(memoryUnit);
            for (MemorySlice slice : memoryUnit.getSlices()) {
                if (sliceId.equals(slice.getId())) {
                    return Result.success(slice.snapshot());
                }
            }
            return memorySliceNotFound(unitId, sliceId);
        } finally {
            memoryLock.unlock();
        }
    }

    @CapabilityMethod
    public Collection<MemoryUnitSnapshot> listMemoryUnits() {
        memoryLock.lock();
        try {
            return memoryUnits.values().stream()
                    .peek(MemoryUnit::load)
                    .peek(this::normalizeMemoryUnit)
                    .map(MemoryUnit::snapshot)
                    .toList();
        } finally {
            memoryLock.unlock();
        }
    }

    private MemoryUnit getOrLoadMemoryUnit(String unitId) {
        MemoryUnit unit = memoryUnits.computeIfAbsent(unitId, id -> {
            MemoryUnit newUnit = new MemoryUnit(id);
            newUnit.register();
            return newUnit;
        });
        unit.load();
        return unit;
    }

    private Result<MemorySliceSnapshot> memorySliceNotFound(String unitId, String sliceId) {
        return Result.failure(new MemoryLookupException(
                "Memory slice not found: " + unitId + ":" + sliceId,
                unitId + ":" + sliceId,
                "MEMORY_SLICE"
        ));
    }

    @CapabilityMethod
    public void refreshMemorySession() {
        memorySessionId = UUID.randomUUID().toString();
    }

    @CapabilityMethod
    public String getMemorySessionId() {
        return memorySessionId;
    }

    private void normalizeMemoryUnit(MemoryUnit memoryUnit) {
        if (memoryUnit.getTimestamp() == null || memoryUnit.getTimestamp() <= 0) {
            memoryUnit.updateTimestamp();
        }
        int maxEndExclusive = memoryUnit.getConversationMessages().size();
        List<MemorySlice> normalizedSlices = new ArrayList<>(memoryUnit.getSlices().size());
        for (MemorySlice slice : memoryUnit.getSlices()) {
            if (slice == null) {
                continue;
            }
            String sliceId = slice.getId();
            if (sliceId == null || sliceId.isBlank()) {
                sliceId = UUID.randomUUID().toString();
            }
            long sliceTimestamp = slice.getTimestamp() == null || slice.getTimestamp() <= 0
                    ? memoryUnit.getTimestamp()
                    : slice.getTimestamp();
            int startIndex = slice.getStartIndex() == null || slice.getStartIndex() < 0
                    ? 0
                    : Math.min(slice.getStartIndex(), maxEndExclusive);
            int endIndex = slice.getEndIndex() == null || slice.getEndIndex() < startIndex
                    ? maxEndExclusive
                    : Math.min(slice.getEndIndex(), maxEndExclusive);
            normalizedSlices.add(MemorySlice.restore(
                    sliceId,
                    startIndex,
                    endIndex,
                    slice.getSummary(),
                    sliceTimestamp
            ));
        }
        memoryUnit.getSlices().clear();
        memoryUnit.getSlices().addAll(normalizedSlices);
        memoryUnit.getSlices().sort(Comparator.naturalOrder());
    }

    @Override
    public @NotNull Path statePath() {
        return Path.of("core", "memory.json");
    }

    @Override
    public void load(@NotNull JSONObject state) {
        String memorySessionId = state.getString("memory_session_id");
        if (memorySessionId == null) {
            throw new IllegalStateException("Memory session id is missing");
        }
        JSONArray array = state.getJSONArray("memory_unit_uuid_set");
        if (array == null) {
            throw new IllegalStateException("Memory unit uuid set is missing");
        }
        for (int i = 0; i < array.size(); i++) {
            String unitUuid = array.getString(i);
            if (unitUuid == null) {
                throw new IllegalStateException("memory_unit_uuid_set is not a uuid array, index: " + i);
            }
            MemoryUnit memoryUnit = new MemoryUnit(unitUuid);
            memoryUnits.put(unitUuid, memoryUnit);
        }
    }

    @Override
    public @NotNull State convert() {
        State state = new State();
        state.append("memory_session_id", StateValue.str(memorySessionId));

        List<String> unitOverview = memoryUnits.keySet().stream()
                .toList();
        state.append("memory_unit_uuid_set", StateValue.arr(unitOverview));
        return state;
    }
}
