import { loggableError } from "@mcpx/toolkit-core/logging";
import { Logger } from "winston";
import { Server } from "@modelcontextprotocol/sdk/server/index.js";
import { Clock } from "@mcpx/toolkit-core/time";
import {
  CloseSessionReason,
  McpxSession,
  SessionLivenessConfig,
  SessionLivenessStore,
  TouchSource,
} from "../model/sessions.js";

type PingResult =
  | { type: "timeout" }
  | { type: "failure"; error: Error }
  | { type: "success" };

const PING_TIMEOUT_FACTOR = 0.8;

export class SessionLivenessManager {
  private _store: SessionLivenessStore;
  private _config: SessionLivenessConfig;
  private _logger: Logger;
  private _gcStopper?: () => void;
  private _clock: Clock;

  constructor(
    store: SessionLivenessStore,
    config: SessionLivenessConfig,
    logger: Logger,
    clock: Clock,
  ) {
    this._store = store;
    this._config = config;
    this._logger = logger.child({ component: "SessionLiveness" });
    this._clock = clock;
  }

  initialize(): void {
    this.startSessionGc({
      ttlMin: this._config.sessionTtlMin,
      sweepIntervalMin: this._config.sessionSweepIntervalMin,
    });
  }

  shutdown(): void {
    this.stopGc();
    this.stopAllSessionLiveness();
  }

  touchSession(sessionId: string, source?: TouchSource): void {
    const session = this._store.getSession(sessionId);
    if (!session?.liveness) {
      return;
    }
    session.liveness.lastSeenAt = this._clock.now().getTime();
    this._logger.silly("Session liveness touch", {
      sessionId,
      metadata: session.metadata,
      source,
    });
  }

  onSessionAdded(sessionId: string): void {
    const session = this._store.getSession(sessionId);
    if (!session) {
      return;
    }

    if (session.metadata.isProbe) {
      this.scheduleProbeTransportTermination(sessionId, {
        probeClientsGraceLivenessPeriodMs:
          this._config.probeClientsGraceLivenessPeriodMs,
      });
    }

    this.startLivenessMonitoring(sessionId, {
      pingIntervalMs: this._config.pingIntervalMs,
    });
  }

  onSessionRemoved(sessionId: string): void {
    this.stopSessionLiveness(sessionId);
  }

  private startSessionGc(options: {
    ttlMin: number;
    sweepIntervalMin?: number;
  }): void {
    if (this._gcStopper) {
      return;
    }
    const { ttlMin, sweepIntervalMin } = options;
    if (ttlMin <= 0) {
      this._logger.debug("Session GC disabled", { ttlMin, sweepIntervalMin });
      return;
    }
    const ttlMs = ttlMin * 60_000;
    const sweepIntervalMs = (sweepIntervalMin ?? ttlMin) * 60_000;
    if (sweepIntervalMs <= 0) {
      this._logger.debug("Session GC disabled by sweep interval", {
        ttlMin,
        sweepIntervalMin,
      });
      return;
    }
    const interval = setInterval(async (): Promise<void> => {
      const now = this._clock.now().getTime();
      for (const [sessionId, session] of this._store.listSessions()) {
        const liveness = session?.liveness;
        if (!liveness) {
          continue;
        }
        const idleMs = now - liveness.lastSeenAt;
        if (idleMs <= ttlMs) {
          continue;
        }
        this._logger.debug("Session idle TTL exceeded, closing session", {
          sessionId,
          metadata: session.metadata,
          idleMs,
          ttlMs,
        });

        await this._store
          .closeSession(sessionId, CloseSessionReason.IdleTtlExceeded)
          .catch((error) => {
            this._logger.warn("Failed to close expired session", {
              sessionId,
              metadata: session.metadata,
              error: loggableError(error),
            });
          });
      }
    }, sweepIntervalMs);

    this._gcStopper = (): void => {
      clearInterval(interval);
      this._gcStopper = undefined;
    };
  }

  private stopGc(): void {
    this._gcStopper?.();
  }

  private stopSessionLiveness(sessionId: string): void {
    const session = this._store.getSession(sessionId);
    if (!session?.liveness) {
      return;
    }
    session.liveness.stopPing();
    delete session.liveness;
  }

  private stopAllSessionLiveness(): void {
    for (const [sessionId] of this._store.listSessions()) {
      this.stopSessionLiveness(sessionId);
    }
  }

  private startLivenessMonitoring(
    sessionId: string,
    options: {
      pingIntervalMs: number;
    },
  ): void {
    const session = this._store.getSession(sessionId);
    if (!session) {
      return;
    }

    if (this._config.sessionTtlMin <= 0) {
      return;
    }

    const entry: McpxSession["liveness"] = {
      lastSeenAt: this._clock.now().getTime(),
      stopPing: this.setupPingMonitoring(
        session.server,
        sessionId,
        session.metadata,
        options,
      ),
    };
    session.liveness = entry;
    return;
  }

  private setupPingMonitoring(
    server: Server,
    sessionId: string,
    metadata: McpxSession["metadata"],
    options: {
      pingIntervalMs: number;
    },
  ): () => void {
    const { pingIntervalMs } = options;
    if (this.shouldSkipPingMonitoring(metadata, sessionId, pingIntervalMs)) {
      return () => {};
    }

    const pingTimeoutMs = Math.floor(pingIntervalMs * PING_TIMEOUT_FACTOR);
    let timeoutId: ReturnType<typeof setTimeout> | undefined;
    let stopped = false;

    const schedule = (): void => {
      timeoutId = setTimeout(() => runPing(), pingIntervalMs);
    };

    const runPing = async (): Promise<void> => {
      const result = await this.executePingWithTimeout(
        server,
        pingTimeoutMs,
        sessionId,
        metadata,
      );
      switch (result.type) {
        case "timeout":
          // ignore (keep for case we want to add a flow for ping timeout)
          break;
        case "failure":
          this._logger.debug("Ping failed", {
            sessionId,
            metadata,
            error: loggableError(result.error),
          });
          break;
        case "success":
          this.touchSession(sessionId, TouchSource.Ping);
          break;
      }
      if (!stopped) schedule();
    };

    schedule();

    return () => {
      if (stopped) {
        return;
      }
      this._logger.debug("Stopping ping monitoring", { sessionId, metadata });
      stopped = true;
      if (timeoutId !== undefined) clearTimeout(timeoutId);
    };
  }

  private shouldSkipPingMonitoring(
    metadata: McpxSession["metadata"],
    sessionId: string,
    pingIntervalMs: number,
  ): boolean {
    if (metadata.clientInfo.adapter?.support?.ping === false) {
      this._logger.info(
        "Client adapter does not support ping, skipping ping monitoring",
        { sessionId, metadata },
      );
      return true;
    }
    if (pingIntervalMs <= 0) {
      this._logger.info("Ping monitoring disabled by interval", {
        sessionId,
        metadata,
        pingIntervalMs,
      });
      return true;
    }
    return false;
  }

  private async executePingWithTimeout(
    server: Server,
    pingTimeoutMs: number,
    sessionId: string,
    metadata: McpxSession["metadata"],
  ): Promise<PingResult> {
    let timeoutId: ReturnType<typeof setTimeout> | undefined;
    const timeoutPromise = new Promise<PingResult>((resolve) => {
      timeoutId = setTimeout(
        () => resolve({ type: "timeout" as const }),
        pingTimeoutMs,
      );
    });

    const pingPromise = server
      .ping()
      .then(() => ({ type: "success" as const }))
      .catch((error) => ({ type: "failure" as const, error }))
      .finally(() => {
        if (timeoutId !== undefined) {
          clearTimeout(timeoutId);
        }
      });

    return Promise.race([pingPromise, timeoutPromise]).catch((error) => {
      this._logger.error("Unexpected error during ping", {
        error: loggableError(error),
        sessionId,
        metadata,
      });
      return { type: "failure" as const, error };
    });
  }

  private scheduleProbeTransportTermination(
    sessionId: string,
    options: {
      probeClientsGraceLivenessPeriodMs: number;
    },
  ): void {
    const session = this._store.getSession(sessionId);
    if (!session) {
      return;
    }
    setTimeout(async () => {
      await this._store
        .closeSession(sessionId, CloseSessionReason.ProbeTermination)
        .catch(() => {
          // ignore
        });
    }, options.probeClientsGraceLivenessPeriodMs);
  }
}
