package tech.easyflow.agent.runtime.agui;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.easyagents.agui.AguiExtendedEvent;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.ReentrantLock;
import java.util.UUID;

/**
 * 将高频 AG-UI delta 在内存中短暂合并并批量写入 Redis。
 *
 * <p>该组件以每个 runId 的同步缓冲保证事件顺序；Redis 故障时保留原批次等待下一次重试，
 * 避免每个模型 token 都产生一次网络往返。</p>
 */
@Component
public class AgentAguiRunJournal {

    private static final Logger LOG = LoggerFactory.getLogger(AgentAguiRunJournal.class);
    private static final int MAX_BUFFERED_EVENTS_PER_RUN = 4096;
    private static final long MAX_EVENT_BYTES = 512L * 1024L;
    private static final long MAX_BUFFERED_BYTES_PER_RUN = 4L * 1024L * 1024L;
    private static final long MAX_BUFFERED_BYTES_GLOBAL = 64L * 1024L * 1024L;
    private static final long COMPLETED_RETRY_WINDOW_MILLIS = 5L * 60L * 1000L;

    private final AgentAguiRunStore store;
    private final ObjectMapper objectMapper;
    private final Map<String, RunBuffer> buffers = new ConcurrentHashMap<>();
    private final AtomicLong globalBufferedBytes = new AtomicLong();

    /**
     * 创建 AG-UI 事件日志协调器。
     *
     * @param store Redis 运行存储
     * @param objectMapper JSON 映射器
     */
    public AgentAguiRunJournal(AgentAguiRunStore store, ObjectMapper objectMapper) {
        this.store = store;
        this.objectMapper = objectMapper;
    }

    /**
     * 追加一条已编码协议事件。
     *
     * @param runId 运行 ID
     * @param data 协议事件 JSON
     */
    public void append(String runId, String data) {
        if (data == null || data.isBlank()) {
            throw new IllegalArgumentException("AG-UI 事件不能为空");
        }
        long dataBytes = utf8Bytes(data);
        if (dataBytes > MAX_EVENT_BYTES) {
            throw new IllegalStateException("单条 AG-UI 事件超过 512 KiB 安全上限");
        }
        RunBuffer buffer = buffers.computeIfAbsent(runId, ignored -> new RunBuffer());
        try {
            synchronized (buffer) {
                if (buffer.completed) {
                    throw new IllegalStateException("AG-UI 运行事件日志已经收口");
                }
                int pendingEventCount = buffer.pendingBatch == null ? 0 : buffer.pendingBatch.events().size();
                if (pendingEventCount + buffer.tailEvents.size() >= MAX_BUFFERED_EVENTS_PER_RUN) {
                    throw new IllegalStateException("AG-UI 事件日志暂时不可用，请稍后重试");
                }
                String merged = buffer.tailEvents.isEmpty()
                        ? null
                        : mergeDelta(buffer.tailEvents.get(buffer.tailEvents.size() - 1), data);
                long previousBytes = 0L;
                long nextBytes = dataBytes;
                if (merged != null) {
                    String previous = buffer.tailEvents.get(buffer.tailEvents.size() - 1);
                    previousBytes = utf8Bytes(previous);
                    nextBytes = utf8Bytes(merged);
                    if (nextBytes > MAX_EVENT_BYTES) {
                        throw new IllegalStateException("合并后的 AG-UI 事件超过 512 KiB 安全上限");
                    }
                }
                long deltaBytes = nextBytes - previousBytes;
                if (buffer.totalBytes() + deltaBytes > MAX_BUFFERED_BYTES_PER_RUN) {
                    throw new IllegalStateException("单次 AG-UI 运行待写日志超过 4 MiB 安全上限");
                }
                long globalBytes = globalBufferedBytes.addAndGet(deltaBytes);
                if (globalBytes > MAX_BUFFERED_BYTES_GLOBAL) {
                    globalBufferedBytes.addAndGet(-deltaBytes);
                    throw new IllegalStateException("AG-UI 待写日志超过 64 MiB 全局安全上限");
                }
                if (merged == null) {
                    buffer.tailEvents.add(data);
                } else {
                    buffer.tailEvents.set(buffer.tailEvents.size() - 1, merged);
                }
                buffer.tailBytes += deltaBytes;
            }
        } catch (RuntimeException exception) {
            removeEmptyBuffer(runId, buffer);
            throw exception;
        }
    }

    /**
     * 丢弃尚未确认持久化的普通事件，并以独立小批次记录日志故障终态。
     *
     * <p>该路径保留已经成功写入 Redis 的历史。即使前一批次处于结果不确定状态，
     * 新终态也使用新的批次 ID，避免幂等判重吞掉故障事件。</p>
     *
     * @param threadId AG-UI 线程 ID
     * @param runId 运行 ID
     * @param code 稳定错误码
     * @param message 用户可理解的错误信息
     */
    public void fail(String threadId, String runId, String code, String message) {
        String terminalEvent;
        try {
            terminalEvent = objectMapper.writeValueAsString(new AguiExtendedEvent.RunError(
                    threadId, runId, message, code));
        } catch (Exception exception) {
            throw new IllegalStateException("AG-UI 日志故障终态编码失败", exception);
        }
        long terminalBytes = utf8Bytes(terminalEvent);
        RunBuffer buffer = buffers.computeIfAbsent(runId, ignored -> new RunBuffer());
        synchronized (buffer) {
            long discardedBytes = buffer.totalBytes();
            buffer.pendingBatch = null;
            buffer.tailEvents.clear();
            buffer.tailBytes = 0L;
            if (discardedBytes > 0L) {
                globalBufferedBytes.addAndGet(-discardedBytes);
            }
            buffer.pendingBatch = new PendingBatch(
                    UUID.randomUUID().toString(), List.of(terminalEvent), terminalBytes);
            globalBufferedBytes.addAndGet(terminalBytes);
            buffer.completed = true;
            buffer.completedAt = System.currentTimeMillis();
        }
        flush(runId, buffer, true);
    }

    /**
     * 在运行收口前同步尝试刷入该运行的剩余事件。
     *
     * @param runId 运行 ID
     */
    public void flush(String runId) {
        RunBuffer buffer = buffers.get(runId);
        if (buffer != null) {
            flush(runId, buffer, true);
        }
    }

    /**
     * 同步刷入终态事件，并在确认缓冲为空后释放本地运行缓冲。
     *
     * @param runId 运行 ID
     */
    public void complete(String runId) {
        RunBuffer buffer = buffers.get(runId);
        if (buffer == null) {
            return;
        }
        synchronized (buffer) {
            buffer.completed = true;
            buffer.completedAt = System.currentTimeMillis();
        }
        flush(runId, buffer, true);
    }

    /**
     * 定期批量写入所有活动运行的待持久化事件。
     */
    @Scheduled(fixedDelayString = "${easyflow.agent.runtime.agui-journal-flush-interval:50ms}")
    public void flushPending() {
        for (Map.Entry<String, RunBuffer> entry : buffers.entrySet()) {
            flush(entry.getKey(), entry.getValue(), false);
        }
    }

    private void flush(String runId, RunBuffer buffer, boolean propagateFailure) {
        boolean acquired;
        if (propagateFailure) {
            buffer.flushLock.lock();
            acquired = true;
        } else {
            acquired = buffer.flushLock.tryLock();
        }
        if (!acquired) {
            return;
        }
        try {
            while (true) {
                PendingBatch batch;
                synchronized (buffer) {
                    if (buffer.pendingBatch == null && !buffer.tailEvents.isEmpty()) {
                        buffer.pendingBatch = new PendingBatch(
                                UUID.randomUUID().toString(),
                                List.copyOf(buffer.tailEvents),
                                buffer.tailBytes);
                        buffer.tailEvents.clear();
                        buffer.tailBytes = 0L;
                    }
                    batch = buffer.pendingBatch;
                    if (batch == null) {
                        if (buffer.completed) {
                            removeBuffer(runId, buffer);
                        }
                        return;
                    }
                }
                try {
                    // Redis I/O 必须位于 RunBuffer 锁外，避免阻塞模型事件线程。
                    store.append(runId, batch.id(), batch.events());
                    synchronized (buffer) {
                        if (buffer.pendingBatch == batch) {
                            buffer.pendingBatch = null;
                            globalBufferedBytes.addAndGet(-batch.bytes());
                        }
                    }
                } catch (RuntimeException exception) {
                    LOG.error("AG-UI 事件日志写入失败: runId={}, eventCount={}",
                            runId, batch.events().size(), exception);
                    boolean batchWasReplaced;
                    synchronized (buffer) {
                        batchWasReplaced = buffer.pendingBatch != batch;
                        if (!batchWasReplaced && buffer.completed && buffer.completedAt > 0L
                                && System.currentTimeMillis() - buffer.completedAt
                                >= COMPLETED_RETRY_WINDOW_MILLIS) {
                            LOG.error("丢弃超过重试窗口的 AG-UI 终态日志缓冲: runId={}, eventCount={}, bytes={}",
                                    runId, buffer.eventCount(), buffer.totalBytes());
                            removeBuffer(runId, buffer);
                        }
                    }
                    if (batchWasReplaced) {
                        continue;
                    }
                    if (propagateFailure) {
                        throw exception;
                    }
                    return;
                }
            }
        } finally {
            buffer.flushLock.unlock();
        }
    }

    private void removeEmptyBuffer(String runId, RunBuffer buffer) {
        synchronized (buffer) {
            if (buffer.pendingBatch == null && buffer.tailEvents.isEmpty() && !buffer.flushLock.isLocked()) {
                buffers.remove(runId, buffer);
            }
        }
    }

    private void removeBuffer(String runId, RunBuffer buffer) {
        if (!buffers.remove(runId, buffer)) {
            return;
        }
        long remainingBytes = buffer.totalBytes();
        buffer.pendingBatch = null;
        buffer.tailEvents.clear();
        buffer.tailBytes = 0L;
        if (remainingBytes > 0L) {
            globalBufferedBytes.addAndGet(-remainingBytes);
        }
    }

    private long utf8Bytes(String value) {
        return value.getBytes(StandardCharsets.UTF_8).length;
    }

    private String mergeDelta(String previous, String current) {
        try {
            JsonNode previousNode = objectMapper.readTree(previous);
            JsonNode currentNode = objectMapper.readTree(current);
            String type = currentNode.path("type").asText();
            if (!type.equals(previousNode.path("type").asText()) || !isMergeable(type)) {
                return null;
            }
            String identityField = "TOOL_CALL_ARGS".equals(type) ? "toolCallId" : "messageId";
            if (!currentNode.path(identityField).asText().equals(previousNode.path(identityField).asText())) {
                return null;
            }
            if (!(previousNode instanceof ObjectNode previousObject)) {
                return null;
            }
            previousObject.put("delta",
                    previousNode.path("delta").asText() + currentNode.path("delta").asText());
            return objectMapper.writeValueAsString(previousObject);
        } catch (Exception ignored) {
            return null;
        }
    }

    private boolean isMergeable(String type) {
        return "TEXT_MESSAGE_CONTENT".equals(type)
                || "REASONING_MESSAGE_CONTENT".equals(type)
                || "TOOL_CALL_ARGS".equals(type);
    }

    private record PendingBatch(String id, List<String> events, long bytes) {
    }

    private static final class RunBuffer {
        private final List<String> tailEvents = new ArrayList<>();
        private final ReentrantLock flushLock = new ReentrantLock();
        private PendingBatch pendingBatch;
        private long tailBytes;
        private boolean completed;
        private long completedAt;

        private int eventCount() {
            return tailEvents.size() + (pendingBatch == null ? 0 : pendingBatch.events().size());
        }

        private long totalBytes() {
            return tailBytes + (pendingBatch == null ? 0L : pendingBatch.bytes());
        }
    }
}
