package tech.easyflow.agent.runtime.agui;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.data.redis.core.HashOperations;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;
import com.easyagents.agui.AguiExtendedEvent;
import tech.easyflow.agent.config.AgentRuntimeProperties;
import tech.easyflow.common.entity.LoginAccount;
import tech.easyflow.common.web.exceptions.BusinessException;

import java.time.Duration;
import java.nio.charset.StandardCharsets;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.regex.Pattern;

/**
 * 基于 Redis 的 AG-UI 运行元数据与顺序事件日志。
 *
 * <p>运行 ID 只承担幂等查找作用，所有读取和取消仍必须校验登录用户与租户归属。</p>
 */
@Component
public class AgentAguiRunStore {

    private static final String RUN_PREFIX = "easyflow:agent:agui:run:";
    private static final String STATUS_FIELD = "status";
    private static final int REPLAY_BATCH_SIZE = 256;
    private static final Pattern RUN_ID_PATTERN = Pattern.compile("[A-Za-z0-9_-]{1,128}");
    private static final long MAX_RUN_EVENT_BYTES = 16L * 1024L * 1024L;
    private static final long TERMINAL_EVENT_RESERVE_BYTES = 64L * 1024L;
    private static final DefaultRedisScript<Long> APPEND_SCRIPT = new DefaultRedisScript<>("""
            local status = redis.call('HGET', KEYS[1], 'status')
            if not status then
                return -3
            end
            local lastCursor = tonumber(redis.call('HGET', KEYS[1], 'lastCursor') or '0')
            if redis.call('HGET', KEYS[1], 'lastBatchId') == ARGV[1] then
                return lastCursor
            end
            if status ~= 'RUNNING' then
                return lastCursor
            end
            local currentBytes = tonumber(redis.call('HGET', KEYS[1], 'eventBytes') or '0')
            local incomingBytes = tonumber(ARGV[5])
            local byteLimit = tonumber(ARGV[6])
            if ARGV[4] ~= '' then
                byteLimit = byteLimit + tonumber(ARGV[7])
            end
            if currentBytes + incomingBytes > byteLimit then
                return -1
            end
            local eventCount = tonumber(ARGV[8])
            for index = 1, eventCount do
                redis.call('RPUSH', KEYS[2], ARGV[8 + index])
            end
            local cursor = redis.call('LLEN', KEYS[2])
            redis.call('HSET', KEYS[1],
                'updatedAt', ARGV[2],
                'lastBatchId', ARGV[1],
                'eventBytes', currentBytes + incomingBytes,
                'lastCursor', cursor)
            if ARGV[4] ~= '' then
                redis.call('HSET', KEYS[1], 'status', ARGV[4])
            end
            redis.call('PEXPIRE', KEYS[1], ARGV[3])
            redis.call('PEXPIRE', KEYS[2], ARGV[3])
            redis.call('PEXPIRE', KEYS[3], ARGV[3])
            return cursor
            """, Long.class);
    private static final DefaultRedisScript<Long> REQUEST_CANCEL_SCRIPT = new DefaultRedisScript<>("""
            if redis.call('HGET', KEYS[1], 'status') ~= 'RUNNING' then
                return 0
            end
            redis.call('HSET', KEYS[1], 'cancelRequested', 'true', 'updatedAt', ARGV[1])
            redis.call('PEXPIRE', KEYS[1], ARGV[2])
            return 1
            """, Long.class);

    private final StringRedisTemplate redisTemplate;
    private final ObjectMapper objectMapper;
    private final Duration retention;

    /**
     * 创建 AG-UI 运行存储。
     *
     * @param redisTemplate Redis 字符串模板
     * @param objectMapper JSON 映射器
     * @param properties Agent 运行配置
     */
    public AgentAguiRunStore(StringRedisTemplate redisTemplate,
                             ObjectMapper objectMapper,
                             AgentRuntimeProperties properties) {
        this.redisTemplate = redisTemplate;
        this.objectMapper = objectMapper;
        this.retention = properties.getAguiRunRetention();
    }

    /**
     * 在模型运行启动前注册唯一运行元数据。
     *
     * @param descriptor 初始运行描述
     * @throws BusinessException runId 已存在时抛出
     */
    public void create(AgentAguiRunDescriptor descriptor) {
        String runId = requireRunId(descriptor == null ? null : descriptor.runId());
        Boolean reserved = redisTemplate.opsForValue().setIfAbsent(
                reservationKey(runId), descriptor.userId(), retention);
        if (!Boolean.TRUE.equals(reserved)) {
            throw new BusinessException("当前 Agent 运行标识已存在，请勿重复提交");
        }
        try {
            Map<String, String> values = new LinkedHashMap<>();
            values.put("runId", runId);
            values.put("requestId", value(descriptor.requestId()));
            values.put("threadId", value(descriptor.threadId()));
            values.put("agentId", value(descriptor.agentId()));
            values.put("sessionId", value(descriptor.sessionId()));
            values.put("userId", value(descriptor.userId()));
            values.put("tenantId", value(descriptor.tenantId()));
            values.put("draft", Boolean.toString(descriptor.draft()));
            values.put(STATUS_FIELD, AgentAguiRunStatus.RUNNING.name());
            values.put("cancelRequested", Boolean.FALSE.toString());
            values.put("eventBytes", "0");
            values.put("lastBatchId", "");
            values.put("lastCursor", "0");
            values.put("createdAt", Long.toString(descriptor.createdAt()));
            values.put("updatedAt", Long.toString(descriptor.updatedAt()));
            meta().putAll(metaKey(runId), values);
            redisTemplate.expire(metaKey(runId), retention);
            redisTemplate.expire(eventsKey(runId), retention);
        } catch (RuntimeException exception) {
            redisTemplate.delete(reservationKey(runId));
            throw exception;
        }
    }

    /**
     * 批量追加严格有序的协议事件。
     *
     * @param runId 运行 ID
     * @param events 已编码 AG-UI 事件
     * @return 追加后的最后事件游标
     */
    public long append(String runId, String batchId, List<String> events) {
        requireRunId(runId);
        if (events == null || events.isEmpty()) {
            AgentAguiRunDescriptor descriptor = find(runId);
            return descriptor == null ? 0L : descriptor.lastCursor();
        }
        if (batchId == null || batchId.isBlank()) {
            throw new IllegalArgumentException("AG-UI 事件批次 ID 不能为空");
        }
        AgentAguiRunStatus terminalStatus = terminalStatus(events);
        long eventBytes = events.stream()
                .mapToLong(event -> event.getBytes(StandardCharsets.UTF_8).length)
                .sum();
        List<String> arguments = new java.util.ArrayList<>(8 + events.size());
        arguments.add(batchId);
        arguments.add(Long.toString(System.currentTimeMillis()));
        arguments.add(Long.toString(retention.toMillis()));
        arguments.add(terminalStatus == null ? "" : terminalStatus.name());
        arguments.add(Long.toString(eventBytes));
        arguments.add(Long.toString(MAX_RUN_EVENT_BYTES));
        arguments.add(Long.toString(TERMINAL_EVENT_RESERVE_BYTES));
        arguments.add(Integer.toString(events.size()));
        arguments.addAll(events);
        Long cursor = redisTemplate.execute(
                APPEND_SCRIPT,
                List.of(metaKey(runId), eventsKey(runId), reservationKey(runId)),
                arguments.toArray());
        if (cursor == null || cursor == -3L) {
            throw new IllegalStateException("AG-UI 运行记录不存在或已过期");
        }
        if (cursor == -1L) {
            throw new IllegalStateException("AG-UI 运行事件超过 16 MiB 安全上限");
        }
        return cursor;
    }

    /**
     * 读取游标之后的一批事件。
     *
     * @param runId 运行 ID
     * @param afterCursor 已消费的最后游标，零表示从头读取
     * @return 按游标升序排列的事件
     */
    public List<String> readAfter(String runId, long afterCursor) {
        requireRunId(runId);
        long start = Math.max(0L, afterCursor);
        List<String> values = redisTemplate.opsForList().range(
                eventsKey(runId), start, start + REPLAY_BATCH_SIZE - 1L);
        return values == null ? List.of() : values;
    }

    /**
     * 查询运行描述。
     *
     * @param runId 运行 ID
     * @return 运行描述；不存在时为 null
     */
    public AgentAguiRunDescriptor find(String runId) {
        requireRunId(runId);
        Map<Object, Object> values = meta().entries(metaKey(runId));
        if (values == null || values.isEmpty()) {
            return null;
        }
        return new AgentAguiRunDescriptor(
                text(values, "runId"),
                text(values, "requestId"),
                text(values, "threadId"),
                text(values, "agentId"),
                text(values, "sessionId"),
                text(values, "userId"),
                text(values, "tenantId"),
                Boolean.parseBoolean(text(values, "draft")),
                parseStatus(text(values, STATUS_FIELD)),
                number(values, "lastCursor"),
                number(values, "createdAt"),
                number(values, "updatedAt"));
    }

    /**
     * 查询并校验当前登录账号对运行的访问权。
     *
     * @param runId 运行 ID
     * @param account 当前登录账号
     * @return 已授权运行描述
     * @throws BusinessException 运行不存在或无访问权时抛出
     */
    public AgentAguiRunDescriptor requireOwned(String runId, LoginAccount account) {
        AgentAguiRunDescriptor descriptor = find(runId);
        if (descriptor == null) {
            throw new BusinessException("Agent 运行记录不存在或已过期");
        }
        String userId = account == null || account.getId() == null ? null : account.getId().toString();
        String tenantId = account == null || account.getTenantId() == null ? null : account.getTenantId().toString();
        if (!descriptor.userId().equals(userId) || !descriptor.tenantId().equals(tenantId)) {
            throw new BusinessException("无权访问该 Agent 运行");
        }
        return descriptor;
    }

    /**
     * 将无法继续的运行收口为失败状态。
     *
     * @param runId 运行 ID
     */
    public void failOwnerLost(AgentAguiRunDescriptor descriptor) {
        if (descriptor == null || descriptor.status().isTerminal()) {
            return;
        }
        try {
            String event = objectMapper.writeValueAsString(new AguiExtendedEvent.RunError(
                    descriptor.threadId(),
                    descriptor.runId(),
                    "Agent 运行节点已不可用",
                    "AGENT_RUN_OWNER_LOST"));
            append(descriptor.runId(), "owner-lost-" + UUID.randomUUID(), List.of(event));
        } catch (RuntimeException exception) {
            throw exception;
        } catch (Exception exception) {
            throw new IllegalStateException("AG-UI owner-lost 终态编码失败", exception);
        }
    }

    /**
     * 记录显式取消意图，供 Runtime 注册前后的启动窗口检查。
     *
     * @param runId 运行 ID
     * @return 本次是否首次接受取消意图
     */
    public boolean requestCancellation(String runId) {
        requireRunId(runId);
        Long result = redisTemplate.execute(
                REQUEST_CANCEL_SCRIPT,
                List.of(metaKey(runId)),
                Long.toString(System.currentTimeMillis()),
                Long.toString(retention.toMillis()));
        return Long.valueOf(1L).equals(result);
    }

    /**
     * 判断运行是否收到显式取消意图。
     *
     * @param runId 运行 ID
     * @return 已请求取消时为 true
     */
    public boolean isCancellationRequested(String runId) {
        Object value = meta().get(metaKey(requireRunId(runId)), "cancelRequested");
        return Boolean.parseBoolean(value == null ? "false" : value.toString());
    }

    private HashOperations<String, Object, Object> meta() {
        return redisTemplate.opsForHash();
    }

    private AgentAguiRunStatus terminalStatus(List<String> events) {
        AgentAguiRunStatus result = null;
        for (String event : events) {
            try {
                JsonNode value = objectMapper.readTree(event);
                String type = value.path("type").asText();
                if ("RUN_FINISHED".equals(type)) {
                    result = AgentAguiRunStatus.COMPLETED;
                } else if ("RUN_ERROR".equals(type)) {
                    result = "RUN_CANCELLED".equals(value.path("code").asText())
                            ? AgentAguiRunStatus.CANCELLED
                            : AgentAguiRunStatus.FAILED;
                }
            } catch (Exception exception) {
                throw new IllegalStateException("AG-UI 运行事件格式不合法", exception);
            }
        }
        return result;
    }

    private AgentAguiRunStatus parseStatus(String value) {
        try {
            return AgentAguiRunStatus.valueOf(value);
        } catch (Exception ignored) {
            return AgentAguiRunStatus.FAILED;
        }
    }

    private String requireRunId(String runId) {
        if (runId == null || !RUN_ID_PATTERN.matcher(runId).matches()) {
            throw new BusinessException("Agent 运行 ID 不合法");
        }
        return runId;
    }

    private String reservationKey(String runId) {
        return runKey(runId, "reserved");
    }

    private String metaKey(String runId) {
        return runKey(runId, "meta");
    }

    private String eventsKey(String runId) {
        return runKey(runId, "events");
    }

    private String runKey(String runId, String suffix) {
        return RUN_PREFIX + "{" + runId + "}:" + suffix;
    }

    private String value(String value) {
        return value == null ? "" : value;
    }

    private String text(Map<Object, Object> values, String key) {
        Object value = values.get(key);
        return value == null ? "" : value.toString();
    }

    private long number(Map<Object, Object> values, String key) {
        try {
            return Long.parseLong(text(values, key));
        } catch (NumberFormatException ignored) {
            return 0L;
        }
    }
}
