package com.cameleer.server.core.agent;

import com.cameleer.server.core.runtime.CreateGuard;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;


/**
 * In-memory agent registry managing agent lifecycle, heartbeats, and commands.
 * <p>
 * Plain class (no Spring annotations) -- wired as a bean by the app module.
 * Uses {@link ConcurrentHashMap} for thread-safe agent storage with atomic
 * record-swapping via {@code compute} operations.
 */
public class AgentRegistryService {

    private static final Logger log = LoggerFactory.getLogger(AgentRegistryService.class);

    private final long staleThresholdMs;
    private final long deadThresholdMs;
    private final long commandExpiryMs;
    private final CreateGuard registerGuard;

    private final ConcurrentHashMap<String, AgentInfo> agents = new ConcurrentHashMap<>();
    private final ConcurrentHashMap<String, ConcurrentLinkedQueue<AgentCommand>> commands = new ConcurrentHashMap<>();
    private final ConcurrentHashMap<String, CompletableFuture<CommandReply>> pendingReplies = new ConcurrentHashMap<>();

    private volatile AgentEventListener eventListener;

    /** Backwards-compatible 3-arg ctor (no enforcement). Used by tests. */
    public AgentRegistryService(long staleThresholdMs, long deadThresholdMs, long commandExpiryMs) {
        this(staleThresholdMs, deadThresholdMs, commandExpiryMs, CreateGuard.NOOP);
    }

    /** Production ctor with license-cap enforcement on new registrations. */
    public AgentRegistryService(long staleThresholdMs, long deadThresholdMs, long commandExpiryMs,
                                CreateGuard registerGuard) {
        this.staleThresholdMs = staleThresholdMs;
        this.deadThresholdMs = deadThresholdMs;
        this.commandExpiryMs = commandExpiryMs;
        this.registerGuard = registerGuard;
    }

    /**
     * Register a new agent or re-register an existing one.
     * Re-registration updates metadata, transitions state to LIVE, and resets timestamps.
     */
    public AgentInfo register(String id, String name, String application, String environmentId,
                               String version, List<String> routeIds, Map<String, Object> capabilities) {
        Instant now = Instant.now();
        AgentInfo newAgent = new AgentInfo(id, name, application, environmentId, version,
                List.copyOf(routeIds), Map.copyOf(capabilities),
                AgentState.LIVE, now, now, null);

        return agents.compute(id, (key, existing) -> {
            if (existing != null) {
                // Re-registration: update metadata, reset to LIVE.
                // Re-registers bypass the license cap check — they don't grow the registry,
                // and rejecting them would orphan an agent that already counts against the cap.
                log.info("Agent {} re-registering (was {})", id, existing.state());
                return existing
                        .withMetadata(name, application, environmentId, version, List.copyOf(routeIds), Map.copyOf(capabilities))
                        .withState(AgentState.LIVE)
                        .withLastHeartbeat(now)
                        .withRegisteredAt(now)
                        .withStaleTransitionTime(null);
            }
            // NEW registration — consult the cap. ConcurrentHashMap.compute propagates the
            // exception thrown here, so the controller advice (LicenseExceptionAdvice) maps
            // LicenseCapExceededException to 403 with the structured envelope.
            registerGuard.check(liveAgentCount());
            log.info("Agent {} registered (name={}, application={}, env={})", id, name, application, environmentId);
            return newAgent;
        });
    }

    /**
     * Live-only agent count for license-cap enforcement.
     * STALE/DEAD/SHUTDOWN agents are excluded so the cap reflects the working fleet, not
     * historical residue. Re-registers of an existing agent revive it via the
     * {@code existing != null} branch in {@link #register}, so this count never double-counts.
     */
    private long liveAgentCount() {
        return agents.values().stream().filter(a -> a.state() == AgentState.LIVE).count();
    }

    /**
     * Process a heartbeat from an agent.
     * Updates lastHeartbeat, routeIds (if provided), capabilities (if provided),
     * and revives STALE or DEAD agents back to LIVE. A DEAD revival happens
     * when an agent resumes (e.g. docker unpause) after a pause long enough
     * to cross both the stale and dead thresholds — the process is the same,
     * no re-registration is required.
     *
     * @return true if the agent is known, false otherwise
     */
    public boolean heartbeat(String id, List<String> routeIds, Map<String, Object> capabilities) {
        AgentInfo updated = agents.computeIfPresent(id, (key, existing) -> {
            Instant now = Instant.now();
            AgentInfo result = existing.withLastHeartbeat(now);
            if (routeIds != null && !routeIds.isEmpty()) {
                result = result.withRouteIds(List.copyOf(routeIds));
            }
            if (capabilities != null && !capabilities.isEmpty()) {
                result = result.withCapabilities(Map.copyOf(capabilities));
            }
            if (existing.state() == AgentState.STALE || existing.state() == AgentState.DEAD) {
                result = result.withState(AgentState.LIVE).withStaleTransitionTime(null);
                log.info("Agent {} revived from {} to LIVE via heartbeat", id, existing.state());
            }
            return result;
        });
        return updated != null;
    }

    /** Overload for callers without routeIds or capabilities (backward compatibility). */
    public boolean heartbeat(String id) {
        return heartbeat(id, null, null);
    }

    /**
     * Manually transition an agent to a new state.
     * Sets staleTransitionTime when transitioning to STALE.
     */
    public void transitionState(String id, AgentState newState) {
        agents.computeIfPresent(id, (key, existing) -> {
            AgentInfo result = existing.withState(newState);
            if (newState == AgentState.STALE) {
                result = result.withStaleTransitionTime(Instant.now());
            } else if (newState == AgentState.LIVE) {
                result = result.withStaleTransitionTime(null);
            }
            log.debug("Agent {} transitioned {} -> {}", id, existing.state(), newState);
            return result;
        });
    }

    /**
     * Gracefully shut down an agent. Transitions to SHUTDOWN state,
     * which is excluded from the LIVE -> STALE -> DEAD lifecycle.
     *
     * @return true if the agent was found and transitioned
     */
    public boolean shutdown(String id) {
        AgentInfo updated = agents.computeIfPresent(id, (key, existing) -> {
            log.info("Agent {} graceful shutdown ({} -> SHUTDOWN)", id, existing.state());
            return existing.withState(AgentState.SHUTDOWN);
        });
        return updated != null;
    }

    /**
     * Remove an agent from the registry entirely.
     *
     * @return true if the agent was found and removed
     */
    public boolean deregister(String id) {
        AgentInfo removed = agents.remove(id);
        if (removed != null) {
            commands.remove(id);
            log.info("Agent {} deregistered (was {})", id, removed.state());
        }
        return removed != null;
    }

    /**
     * Check all agents and apply lifecycle transitions:
     * LIVE -> STALE when lastHeartbeat exceeds staleThresholdMs,
     * STALE -> DEAD when staleTransitionTime exceeds deadThresholdMs.
     * SHUTDOWN and DEAD agents are skipped.
     */
    public void checkLifecycle() {
        Instant now = Instant.now();
        for (Map.Entry<String, AgentInfo> entry : agents.entrySet()) {
            AgentInfo agent = entry.getValue();

            if (agent.state() == AgentState.LIVE) {
                Duration sinceHeartbeat = Duration.between(agent.lastHeartbeat(), now);
                if (sinceHeartbeat.toMillis() > staleThresholdMs) {
                    agents.computeIfPresent(entry.getKey(), (k, current) -> {
                        if (current.state() == AgentState.LIVE) {
                            log.info("Agent {} LIVE -> STALE (no heartbeat for {}ms)", k, sinceHeartbeat.toMillis());
                            return current.withState(AgentState.STALE)
                                    .withStaleTransitionTime(now);
                        }
                        return current;
                    });
                }
            } else if (agent.state() == AgentState.STALE && agent.staleTransitionTime() != null) {
                Duration sinceStale = Duration.between(agent.staleTransitionTime(), now);
                if (sinceStale.toMillis() > deadThresholdMs) {
                    agents.computeIfPresent(entry.getKey(), (k, current) -> {
                        if (current.state() == AgentState.STALE) {
                            log.info("Agent {} STALE -> DEAD (stale for {}ms)", k, sinceStale.toMillis());
                            return current.withState(AgentState.DEAD);
                        }
                        return current;
                    });
                }
            }
            // DEAD agents remain DEAD (no auto-purge)
        }
    }

    /**
     * Find an agent by its ID.
     *
     * @return the agent info, or null if not found
     */
    public AgentInfo findById(String id) {
        return agents.get(id);
    }

    /**
     * Return all registered agents regardless of state.
     */
    public List<AgentInfo> findAll() {
        return new ArrayList<>(agents.values());
    }

    /**
     * Return all agents in the specified state.
     */
    public List<AgentInfo> findByState(AgentState state) {
        return agents.values().stream()
                .filter(a -> a.state() == state)
                .toList();
    }

    /**
     * Return all agents belonging to the given application.
     */
    public List<AgentInfo> findByApplication(String application) {
        return agents.values().stream()
                .filter(a -> application.equals(a.applicationId()))
                .toList();
    }

    /**
     * Return all agents belonging to the given application and environment.
     */
    public List<AgentInfo> findByApplicationAndEnvironment(String application, String environment) {
        return agents.values().stream()
                .filter(a -> application.equals(a.applicationId()))
                .filter(a -> environment.equals(a.environmentId()))
                .toList();
    }

    /**
     * Add a command to an agent's pending queue.
     * Notifies the event listener if one is set.
     *
     * @return the created command
     */
    public AgentCommand addCommand(String agentId, CommandType type, String payload) {
        AgentCommand command = new AgentCommand(
                UUID.randomUUID().toString(),
                type,
                payload,
                agentId,
                Instant.now(),
                CommandStatus.PENDING
        );

        commands.computeIfAbsent(agentId, k -> new ConcurrentLinkedQueue<>()).add(command);
        log.debug("Command {} ({}) added for agent {}", command.id(), type, agentId);

        AgentEventListener listener = this.eventListener;
        if (listener != null) {
            listener.onCommandReady(agentId, command);
        }

        return command;
    }

    /**
     * Acknowledge a command, transitioning it from PENDING/DELIVERED to ACKNOWLEDGED.
     *
     * @return true if the command was found and acknowledged
     */
    public boolean acknowledgeCommand(String agentId, String commandId) {
        ConcurrentLinkedQueue<AgentCommand> queue = commands.get(agentId);
        if (queue == null) {
            return false;
        }

        // Replace the command in the queue with an acknowledged version
        boolean[] found = {false};
        ConcurrentLinkedQueue<AgentCommand> newQueue = new ConcurrentLinkedQueue<>();
        for (AgentCommand cmd : queue) {
            if (cmd.id().equals(commandId) && !found[0]) {
                newQueue.add(cmd.withStatus(CommandStatus.ACKNOWLEDGED));
                found[0] = true;
            } else {
                newQueue.add(cmd);
            }
        }

        if (found[0]) {
            commands.put(agentId, newQueue);
            log.debug("Command {} acknowledged by agent {}", commandId, agentId);
        }

        return found[0];
    }

    /**
     * Return all PENDING commands for the specified agent.
     */
    public List<AgentCommand> findPendingCommands(String agentId) {
        ConcurrentLinkedQueue<AgentCommand> queue = commands.get(agentId);
        if (queue == null) {
            return List.of();
        }
        return queue.stream()
                .filter(cmd -> cmd.status() == CommandStatus.PENDING)
                .toList();
    }

    /**
     * Mark a command as DELIVERED (sent via SSE but not yet acknowledged).
     */
    public void markDelivered(String agentId, String commandId) {
        ConcurrentLinkedQueue<AgentCommand> queue = commands.get(agentId);
        if (queue == null) {
            return;
        }

        ConcurrentLinkedQueue<AgentCommand> newQueue = new ConcurrentLinkedQueue<>();
        for (AgentCommand cmd : queue) {
            if (cmd.id().equals(commandId) && cmd.status() == CommandStatus.PENDING) {
                newQueue.add(cmd.withStatus(CommandStatus.DELIVERED));
            } else {
                newQueue.add(cmd);
            }
        }
        commands.put(agentId, newQueue);
    }

    /**
     * Expire PENDING commands older than commandExpiryMs.
     */
    public void expireOldCommands() {
        Instant cutoff = Instant.now().minusMillis(commandExpiryMs);
        for (Map.Entry<String, ConcurrentLinkedQueue<AgentCommand>> entry : commands.entrySet()) {
            ConcurrentLinkedQueue<AgentCommand> queue = entry.getValue();
            queue.removeIf(cmd ->
                    cmd.status() == CommandStatus.PENDING && cmd.createdAt().isBefore(cutoff));
        }
    }

    /**
     * Register a command that expects a synchronous reply from the agent.
     * Returns a CompletableFuture that will be completed when the agent ACKs the command.
     * Auto-cleans up from the pending map on completion or timeout.
     */
    public CompletableFuture<CommandReply> addCommandWithReply(String agentId, CommandType type, String payload) {
        AgentCommand command = addCommand(agentId, type, payload);
        CompletableFuture<CommandReply> future = new CompletableFuture<>();
        pendingReplies.put(command.id(), future);
        future.whenComplete((result, ex) -> pendingReplies.remove(command.id()));
        return future;
    }

    /**
     * Send a command to all LIVE agents in a group and return futures for collecting replies.
     * Returns a map of agentId -> CompletableFuture&lt;CommandReply&gt;.
     */
    public Map<String, CompletableFuture<CommandReply>> addGroupCommandWithReplies(
            String group, CommandType type, String payload) {
        return addGroupCommandWithReplies(group, null, type, payload);
    }

    /**
     * Send a command to all LIVE agents in a group, optionally filtered by environment.
     * When environment is null, targets all agents for the application.
     * Returns a map of agentId -> CompletableFuture&lt;CommandReply&gt;.
     */
    public Map<String, CompletableFuture<CommandReply>> addGroupCommandWithReplies(
            String group, String environment, CommandType type, String payload) {
        Map<String, CompletableFuture<CommandReply>> results = new LinkedHashMap<>();
        List<AgentInfo> candidates = environment != null
                ? findByApplicationAndEnvironment(group, environment)
                : findByApplication(group);
        List<AgentInfo> liveAgents = candidates.stream()
                .filter(a -> a.state() == AgentState.LIVE)
                .toList();

        for (AgentInfo agent : liveAgents) {
            AgentCommand cmd = addCommand(agent.instanceId(), type, payload);
            CompletableFuture<CommandReply> future = new CompletableFuture<>();
            pendingReplies.put(cmd.id(), future);
            future.whenComplete((r, ex) -> pendingReplies.remove(cmd.id()));
            results.put(agent.instanceId(), future);
        }

        return results;
    }

    /**
     * Complete a pending reply future for a command.
     * Called when an agent ACKs a command that was registered via {@link #addCommandWithReply}.
     * No-op if no pending future exists for the given command ID.
     */
    public void completeReply(String commandId, String status, String message, String data) {
        CompletableFuture<CommandReply> future = pendingReplies.remove(commandId);
        if (future != null) {
            future.complete(new CommandReply(status, message, data));
        }
    }

    /**
     * Set the event listener for command notifications.
     * The SSE layer in the app module implements this interface.
     */
    public void setEventListener(AgentEventListener listener) {
        this.eventListener = listener;
    }
}
