package com.gzzm.lobster.tool.mcp;

import com.gzzm.lobster.config.LobsterConfig;
import com.gzzm.platform.commons.Tools;
import net.cyan.nest.annotation.Inject;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Runtime snapshot cache for MCP tool execution.
 *
 * <p>ToolRegistry keeps executors by tool name, so the executor must not hold a
 * long-lived McpServerConfig. This cache gives every call a recent server/tool
 * snapshot while keeping the hot path mostly in-memory.
 */
public class McpRuntimeToolCache {

    private static final long DEFAULT_TTL_MS = 600_000L;
    private static final Map<String, Entry> CACHE = new ConcurrentHashMap<>();

    @Inject private McpServerConfigDao serverConfigDao;
    @Inject private McpToolCacheDao toolCacheDao;

    public Snapshot get(String localToolName) {
        if (localToolName == null || localToolName.trim().isEmpty()) return null;
        long now = System.currentTimeMillis();
        Entry cached = CACHE.get(localToolName);
        if (cached != null && cached.expiresAtMs > now) return cached.snapshot;
        return refreshLocal(localToolName);
    }

    public Snapshot refreshLocal(String localToolName) {
        if (localToolName == null || localToolName.trim().isEmpty()) return null;
        try {
            McpToolCache tool = toolCacheDao().getByLocalToolName(localToolName);
            if (tool == null || Boolean.FALSE.equals(tool.getEnabled())
                    || tool.getExposureMode() != McpToolExposureMode.DIRECT) {
                CACHE.remove(localToolName);
                return null;
            }
            McpServerConfig server = serverConfigDao().getConfig(tool.getServerId());
            if (server == null || Boolean.FALSE.equals(server.getEnabled())) {
                CACHE.remove(localToolName);
                return null;
            }
            Snapshot snapshot = new Snapshot(server, tool);
            CACHE.put(localToolName, new Entry(snapshot, System.currentTimeMillis() + ttlMs()));
            return snapshot;
        } catch (Throwable t) {
            try { Tools.log("[McpRuntimeToolCache] refresh local failed: " + localToolName, t); }
            catch (Throwable ignore) { /* ignore */ }
            Entry cached = CACHE.get(localToolName);
            if (cached != null) {
                CACHE.put(localToolName, new Entry(cached.snapshot, System.currentTimeMillis() + ttlMs()));
                return cached.snapshot;
            }
            return null;
        }
    }

    public List<Snapshot> refreshServer(String serverId) {
        if (serverId == null || serverId.trim().isEmpty()) return Collections.emptyList();
        try {
            McpServerConfig server = serverConfigDao().getConfig(serverId);
            if (server == null || Boolean.FALSE.equals(server.getEnabled())) {
                invalidateServer(serverId);
                return Collections.emptyList();
            }
            List<McpToolCache> tools = toolCacheDao().listByServer(serverId);
            List<Snapshot> out = new ArrayList<>();
            long expires = System.currentTimeMillis() + ttlMs();
            for (McpToolCache tool : tools) {
                if (tool == null || tool.getLocalToolName() == null) continue;
                if (Boolean.FALSE.equals(tool.getEnabled())
                        || tool.getExposureMode() != McpToolExposureMode.DIRECT) {
                    CACHE.remove(tool.getLocalToolName());
                    continue;
                }
                Snapshot snapshot = new Snapshot(server, tool);
                CACHE.put(tool.getLocalToolName(), new Entry(snapshot, expires));
                out.add(snapshot);
            }
            return out;
        } catch (Throwable t) {
            try { Tools.log("[McpRuntimeToolCache] refresh server failed: " + serverId, t); }
            catch (Throwable ignore) { /* ignore */ }
            return null;
        }
    }

    public void invalidateLocal(String localToolName) {
        if (localToolName != null) CACHE.remove(localToolName);
    }

    public void invalidateServer(String serverId) {
        if (serverId == null || serverId.trim().isEmpty()) return;
        for (Map.Entry<String, Entry> e : CACHE.entrySet()) {
            Snapshot snapshot = e.getValue() == null ? null : e.getValue().snapshot;
            if (snapshot != null && snapshot.server != null
                    && serverId.equals(snapshot.server.getServerId())) {
                CACHE.remove(e.getKey(), e.getValue());
            }
        }
    }

    private McpServerConfigDao serverConfigDao() {
        try {
            McpServerConfigDao d = Tools.getBean(McpServerConfigDao.class);
            if (d != null) return d;
        } catch (Throwable ignore) { /* fallback */ }
        return serverConfigDao;
    }

    private long ttlMs() {
        try {
            long configured = LobsterConfig.getMcpRuntimeToolCacheTtlMs();
            return configured > 0 ? configured : DEFAULT_TTL_MS;
        } catch (Throwable ignore) {
            return DEFAULT_TTL_MS;
        }
    }

    private McpToolCacheDao toolCacheDao() {
        try {
            McpToolCacheDao d = Tools.getBean(McpToolCacheDao.class);
            if (d != null) return d;
        } catch (Throwable ignore) { /* fallback */ }
        return toolCacheDao;
    }

    private static final class Entry {
        final Snapshot snapshot;
        final long expiresAtMs;

        Entry(Snapshot snapshot, long expiresAtMs) {
            this.snapshot = snapshot;
            this.expiresAtMs = expiresAtMs;
        }
    }

    public static final class Snapshot {
        private final McpServerConfig server;
        private final McpToolCache tool;

        Snapshot(McpServerConfig server, McpToolCache tool) {
            this.server = server;
            this.tool = tool;
        }

        public McpServerConfig getServer() { return server; }
        public McpToolCache getTool() { return tool; }
    }
}