package one.edee.mcp.jdwp;

import com.sun.jdi.VirtualMachine;
import org.jspecify.annotations.Nullable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Service;

import java.time.Duration;
import java.time.Instant;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicReference;

/**
 * Passive JDI traffic observer plus active probe used to answer "is the JDI connection wedged?"
 * without parking on a user-facing tool. Three concerns:
 * <ul>
 *   <li><b>Passive observation</b> — {@link #notifyTraffic()} is called on the two signals the
 *       service actually has: every {@code EventSet} the JDI event listener drains from the
 *       queue (inbound side), and {@link JDIConnectionService} on a successful attach handshake.
 *       Outbound JDWP commands issued by short-lived tools do <b>not</b> currently refresh the
 *       timestamp; if the target VM is genuinely responsive but idle (no breakpoint / step /
 *       exception events firing), the active probe is what re-classifies the connection as
 *       responsive. No extra socket traffic is generated by the passive side.</li>
 *   <li><b>Active probe</b> — a daemon thread wakes every {@link #PROBE_TICK} and, when the
 *       silence interval exceeds {@link #SILENCE_THRESHOLD}, issues a single {@code vm.allThreads()}
 *       JDWP command wrapped in {@link Future#get(long, TimeUnit)} so the watchdog itself cannot
 *       deadlock. A successful probe re-classifies the connection as responsive; a hung probe
 *       conclusively classifies it as wedged.</li>
 *   <li><b>No automatic recovery</b> — the monitor never tears down the connection or restarts
 *       anything. It publishes a {@link JdiHealthSnapshot} consumed by {@code jdwp_diagnose} and
 *       the soft-wait envelope, leaving the recovery decision to the agent.</li>
 * </ul>
 *
 * <p>Thread-safety: snapshots are published via a single {@link AtomicReference} so callers never
 * observe a torn triple; the {@code vm} reference is also volatile-via-{@link AtomicReference} so
 * a {@link #stop()} concurrent with a probe tick simply makes the next tick a no-op rather than
 * raising on a stale value.
 */
@Service
public class JdiHealthMonitor {
    private static final Logger log = LoggerFactory.getLogger(JdiHealthMonitor.class);

    /**
     * How often the watchdog thread wakes up to evaluate the silence interval. Short enough that
     * the agent's diagnose poll never has to wait for a full silence window, long enough that the
     * thread does no useful work most of the time.
     */
    private static final Duration PROBE_TICK = Duration.ofSeconds(5);

    /**
     * Silence threshold before the watchdog escalates to an active probe. Aligned with the
     * soft-wait envelope window so the agent sees {@code "unresponsive_for_30s"} right as its
     * first soft-wait window also expires.
     */
    private static final Duration SILENCE_THRESHOLD = Duration.ofSeconds(30);

    /**
     * Maximum time the active probe is allowed to take before it is considered hung. Short enough
     * that a wedged probe does not itself become the wedge.
     */
    private static final Duration PROBE_TIMEOUT = Duration.ofSeconds(5);

    /**
     * Classifies the connection at the moment the snapshot was taken. {@link #DISCONNECTED} is the
     * starting state and is also re-entered when {@link #stop()} runs; {@link #RESPONSIVE} is the
     * happy path; {@link #UNRESPONSIVE} carries a duration so the renderer can produce
     * {@code "unresponsive_for_47s"}.
     */
    public enum Status {
        /** Last observed traffic is within {@link #SILENCE_THRESHOLD}, no probe needed. */
        RESPONSIVE,
        /** Silence exceeded {@link #SILENCE_THRESHOLD}; the most recent probe outcome is in the snapshot. */
        UNRESPONSIVE,
        /** No live VM is registered with the monitor (pre-connect or after stop). */
        DISCONNECTED
    }

    /**
     * Immutable snapshot of monitor state for diagnostic rendering. All fields are nullable to
     * keep the type honest about the pre-first-event state: until the listener drains its first
     * event we genuinely don't know when the last packet flew.
     *
     * @param status            current health classification (see {@link Status})
     * @param lastTrafficAt     wall-clock time of the most recent observed packet, or {@code null}
     *                          before any traffic has been observed
     * @param silentFor         elapsed time since {@link #lastTrafficAt}, or {@code null} when
     *                          {@code lastTrafficAt} is null
     * @param lastProbeAt       wall-clock time of the most recent active probe, or {@code null}
     *                          when no probe has been needed since attach
     * @param lastProbeOutcome  human-readable probe outcome (e.g. {@code "OK"}, {@code "timed out after 5s"})
     *                          or {@code null} when no probe has been run
     */
    public record JdiHealthSnapshot(
        Status status,
        @Nullable Instant lastTrafficAt,
        @Nullable Duration silentFor,
        @Nullable Instant lastProbeAt,
        @Nullable String lastProbeOutcome
    ) {
        /** Convenience: snapshot used before any traffic is observed and no VM is registered. */
        public static JdiHealthSnapshot disconnected() {
            return new JdiHealthSnapshot(Status.DISCONNECTED, null, null, null, null);
        }
    }

    /**
     * Live VM the watchdog probes. Null when no session is active; reset to null on {@link #stop()}.
     * Published atomically so a concurrent tick reads either the new VM or null — never a torn
     * value.
     */
    private final AtomicReference<@Nullable VirtualMachine> vmRef = new AtomicReference<>();

    /**
     * Most recent observed traffic timestamp. Set on every {@link #notifyTraffic()} and read by
     * the probe thread; volatile-via-{@link AtomicReference} so the publish/observe pair is safe
     * across threads without taking a lock.
     */
    private final AtomicReference<@Nullable Instant> lastTrafficAt = new AtomicReference<>();

    /**
     * Most recent {@link JdiHealthSnapshot} the watchdog (or {@link #notifyTraffic}) computed.
     * {@link #snapshot()} returns the current value; updates are atomic so an agent's diagnose
     * call can never observe a snapshot in mid-write.
     */
    private final AtomicReference<JdiHealthSnapshot> snapshotRef =
        new AtomicReference<>(JdiHealthSnapshot.disconnected());

    /**
     * Daemon executor that owns the recurring probe tick. Single-threaded so a slow probe cannot
     * stack up against itself; the tick is rescheduled by Spring's lifecycle, not by the previous
     * tick.
     */
    @Nullable
    private volatile ScheduledExecutorService probeExecutor;

    /**
     * Handle to the scheduled probe task so {@link #stop()} can cancel it without shutting the
     * executor down — Spring's bean lifecycle keeps the bean instance alive across connect /
     * disconnect cycles, so we re-use the executor.
     */
    @Nullable
    private volatile ScheduledFuture<?> probeTask;

    /**
     * Registers a fresh VM with the monitor and arms the probe tick. Idempotent for the
     * same-VM case; a second call with a different VM tears down the previous task before
     * arming the new one. Called from {@link JDIConnectionService#connect} after a successful
     * attach.
     */
    public synchronized void start(VirtualMachine vm) {
        vmRef.set(vm);
        // Seed the last-traffic timestamp so the watchdog does not immediately flag a brand-new
        // connection as silent — the attach packet itself counts as traffic.
        notifyTraffic();
        if (probeExecutor == null) {
            probeExecutor = Executors.newSingleThreadScheduledExecutor(r -> {
                final Thread t = new Thread(r, "jdi-health-monitor");
                t.setDaemon(true);
                return t;
            });
        }
        cancelProbeTask();
        probeTask = probeExecutor.scheduleWithFixedDelay(
            this::probeTick,
            PROBE_TICK.toMillis(),
            PROBE_TICK.toMillis(),
            TimeUnit.MILLISECONDS
        );
        log.info("[Health] Watchdog armed (silence threshold {}s, probe timeout {}s)",
            SILENCE_THRESHOLD.toSeconds(), PROBE_TIMEOUT.toSeconds());
    }

    /**
     * Detaches the monitor from the current VM and cancels the probe task. The bean instance
     * survives — Spring keeps it across connect / disconnect cycles — but the snapshot reverts
     * to {@link Status#DISCONNECTED} so a post-disconnect diagnose call renders honestly.
     */
    public synchronized void stop() {
        cancelProbeTask();
        vmRef.set(null);
        lastTrafficAt.set(null);
        snapshotRef.set(JdiHealthSnapshot.disconnected());
    }

    /**
     * Records a fresh "we saw a JDWP packet just now" timestamp. Called by:
     * <ul>
     *   <li>{@link JdiEventListener#listen} every time the event queue yields an event (inbound side);</li>
     *   <li>{@link JDIConnectionService#connect} on successful attach (outbound side — the connect
     *       handshake counts as traffic);</li>
     *   <li>{@link #probeTick} on a successful active probe (closes the loop so the next tick does
     *       not re-probe in a tight window).</li>
     * </ul>
     * Synchronized on {@code this} to pair with {@link #stop()}: a concurrent {@code stop()} cannot
     * interleave between the {@code vmRef} check and the snapshot publish, so a late
     * {@code notifyTraffic} from an event drained right before disconnect cannot revive the monitor
     * by overwriting the DISCONNECTED snapshot. Bails early when {@code vmRef} is null because no
     * session is active.
     */
    public synchronized void notifyTraffic() {
        if (vmRef.get() == null) {
            return;
        }
        final Instant now = Instant.now();
        lastTrafficAt.set(now);
        // Optimistically classify as responsive — the watchdog will downgrade us if the next
        // silence interval lapses. Preserve the prior probe outcome so the diagnose renderer
        // can still surface "last probe at …" after the connection self-heals.
        final JdiHealthSnapshot prev = snapshotRef.get();
        snapshotRef.set(new JdiHealthSnapshot(
            Status.RESPONSIVE, now, Duration.ZERO,
            prev.lastProbeAt(), prev.lastProbeOutcome()
        ));
    }

    /**
     * Returns the most recent {@link JdiHealthSnapshot}. The returned record is immutable so the
     * caller can stash it without worrying about a subsequent tick overwriting fields underneath.
     */
    public JdiHealthSnapshot snapshot() {
        return snapshotRef.get();
    }

    /**
     * One iteration of the watchdog loop. Quick exits dominate the common path: no VM, no
     * traffic recorded yet, or silence within threshold all return early without issuing any
     * JDWP traffic. Only when silence exceeds the threshold do we escalate to an active probe.
     */
    private void probeTick() {
        try {
            final VirtualMachine vm = vmRef.get();
            if (vm == null) {
                return;
            }
            final Instant lastSeen = lastTrafficAt.get();
            if (lastSeen == null) {
                return;
            }
            final Duration silentFor = Duration.between(lastSeen, Instant.now());
            if (silentFor.compareTo(SILENCE_THRESHOLD) < 0) {
                // Still inside the responsive window — refresh the snapshot's silentFor so the
                // diagnose renderer shows a current value, but don't probe. Synchronized + re-check
                // pairs with stop() so a stopped monitor cannot be revived by a late tick.
                synchronized (this) {
                    if (vmRef.get() == vm) {
                        final JdiHealthSnapshot prev = snapshotRef.get();
                        snapshotRef.set(new JdiHealthSnapshot(
                            Status.RESPONSIVE, lastSeen, silentFor,
                            prev.lastProbeAt(), prev.lastProbeOutcome()
                        ));
                    }
                }
                return;
            }
            runActiveProbe(vm, lastSeen, silentFor);
        } catch (Throwable t) {
            // Probe failures must never kill the watchdog thread — log and move on.
            log.warn("[Health] Probe tick threw {}: {}", t.getClass().getSimpleName(), t.getMessage());
        }
    }

    /**
     * Issues a {@code vm.allThreads()} probe in a future and blocks on its completion for at most
     * {@link #PROBE_TIMEOUT}. The probe is intentionally a cheap, side-effect-free JDI command
     * that multiplexes alongside any in-progress {@code invokeMethod} — if it returns promptly,
     * the JDI socket is responsive; if it times out, the socket is conclusively wedged.
     */
    private void runActiveProbe(VirtualMachine vm, Instant lastSeen, Duration silentFor) {
        final Instant probeStartedAt = Instant.now();
        // Single-shot executor; cleaned up via the finally block so a slow JDI thread cannot
        // leak workers on the probe path.
        final java.util.concurrent.ExecutorService probeWorker = Executors.newSingleThreadExecutor(r -> {
            final Thread t = new Thread(r, "jdi-health-probe");
            t.setDaemon(true);
            return t;
        });
        try {
            final Future<Integer> probe = probeWorker.submit(() -> vm.allThreads().size());
            try {
                final int count = probe.get(PROBE_TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
                // Successful probe — the connection is alive after all. The check+publish is
                // wrapped in synchronized(this) so a concurrent stop() cannot interleave between
                // the vmRef re-check and either write below (notifyTraffic or snapshotRef.set).
                // Without the lock, the check-then-act would race stop() and revive a stopped
                // monitor by overwriting its DISCONNECTED snapshot.
                synchronized (this) {
                    if (vmRef.get() == vm) {
                        notifyTraffic();
                        snapshotRef.set(new JdiHealthSnapshot(
                            Status.RESPONSIVE,
                            lastTrafficAt.get(),
                            Duration.ZERO,
                            probeStartedAt,
                            String.format("OK (%d threads)", count)
                        ));
                    }
                }
                log.debug("[Health] Active probe succeeded after {}s silence (threads={})",
                    silentFor.toSeconds(), count);
            } catch (TimeoutException timeout) {
                probe.cancel(true);
                synchronized (this) {
                    if (vmRef.get() == vm) {
                        snapshotRef.set(new JdiHealthSnapshot(
                            Status.UNRESPONSIVE,
                            lastSeen,
                            silentFor,
                            probeStartedAt,
                            String.format("vm.allThreads() timed out after %ds", PROBE_TIMEOUT.toSeconds())
                        ));
                    }
                }
                log.warn("[Health] JDI probe timed out — connection appears wedged (silent for {}s)",
                    silentFor.toSeconds());
            } catch (ExecutionException ee) {
                final Throwable cause = ee.getCause() != null ? ee.getCause() : ee;
                synchronized (this) {
                    if (vmRef.get() == vm) {
                        snapshotRef.set(new JdiHealthSnapshot(
                            Status.UNRESPONSIVE,
                            lastSeen,
                            silentFor,
                            probeStartedAt,
                            String.format("vm.allThreads() failed: %s", cause.getClass().getSimpleName())
                        ));
                    }
                }
                log.warn("[Health] JDI probe failed with {}: {}",
                    cause.getClass().getSimpleName(), cause.getMessage());
            } catch (InterruptedException ie) {
                Thread.currentThread().interrupt();
            }
        } finally {
            probeWorker.shutdownNow();
        }
    }

    /**
     * Cancels and discards the current scheduled probe task; caller must hold the monitor's
     * intrinsic lock so a concurrent {@link #start} cannot race the cancel. {@code mayInterruptIfRunning}
     * is set to {@code true} so an in-flight probe parked in {@code Future.get(timeout)} wakes up
     * promptly — the snapshot write inside {@code runActiveProbe} is still {@code vmRef}-guarded
     * for the case where the interrupt arrives too late to short-circuit the publish.
     */
    private void cancelProbeTask() {
        final ScheduledFuture<?> task = probeTask;
        if (task != null) {
            task.cancel(true);
            probeTask = null;
        }
    }
}
