import * as plugins from './smartexit.plugins.js';
import { ProcessLifecycle } from './smartexit.classes.lifecycle.js';
import {
  probeProcessGroupByPid,
  signalProcessTreeByPid,
} from './smartexit.functions.processtree.js';
import type { TProcessSignal } from './smartexit.types.js';

export type { TProcessSignal } from './smartexit.types.js';

const defaultProcessTerminationGraceMs = 1_000;
const maximumTimerDelayMs = 2_147_483_647;

export interface ISmartExitOptions {
  /** Completely disable logging for this instance. */
  silent?: boolean;
  /** Time between POSIX SIGTERM and SIGKILL escalation. Defaults to 1,000ms. */
  processTerminationGraceMs?: number;
}

/** A stable object identity representing one process root owned by SmartExit. */
export interface IProcessHandle {
  readonly pid?: number;
}

interface IProcessHandleTracking {
  readonly pid: number;
  readonly generation: symbol;
}

export interface IKillAllResult {
  processesKilled: number;
  cleanupFunctionsRan: number;
}

export interface IKillAllOptions {
  /** Abort the graceful interval and proceed directly to force-kill escalation. */
  terminationSignal?: AbortSignal;
}

class TrackedPidSet extends Set<number> {
  constructor(
    private readonly onAdd: (pidArg: number) => void,
    private readonly onDelete: (pidArg: number) => void,
  ) {
    super();
  }

  public override add(pidArg: number): this {
    const alreadyTracked = this.has(pidArg);
    super.add(pidArg);
    if (!alreadyTracked) {
      this.onAdd(pidArg);
    }
    return this;
  }

  public override delete(pidArg: number): boolean {
    const deleted = super.delete(pidArg);
    if (deleted) {
      this.onDelete(pidArg);
    }
    return deleted;
  }

  public override clear(): void {
    for (const pid of Array.from(this)) {
      this.delete(pid);
    }
  }
}

/**
 * SmartExit — process and cleanup tracker.
 *
 * Lightweight: the constructor does NOT register signal handlers.
 * Each instance auto-registers with ProcessLifecycle so that global
 * shutdown (triggered by ProcessLifecycle.install()) kills all tracked processes.
 *
 * Libraries should create instances freely. Only the application entry point
 * should call `ProcessLifecycle.install()`.
 */
export class SmartExit {
  private static readonly cleanupCallContext =
    new plugins.asyncHooks.AsyncLocalStorage<SmartExit>();

  private static releaseTrackedPidFromAllInstances(pidArg: number): void {
    for (const instance of ProcessLifecycle.getInstances()) {
      instance.releaseTrackedPid(pidArg);
    }
  }

  /**
   * Kill an owned process tree. On POSIX, the PID must identify a detached
   * process-group leader; Windows treats it as the taskkill root PID.
   */
  public static async killTreeByPid(
    pidArg: number,
    signalArg: TProcessSignal = 'SIGKILL',
  ): Promise<void> {
    const result = signalProcessTreeByPid(pidArg, signalArg);
    const forcedTerminationSucceeded = result.status === 'signalled'
      && (process.platform === 'win32' || signalArg === 'SIGKILL');
    if (result.status === 'notFound' || forcedTerminationSucceeded) {
      SmartExit.releaseTrackedPidFromAllInstances(pidArg);
    }
    // Preserve the established public contract: missing trees are harmless,
    // Windows taskkill failures are swallowed, and other POSIX errors reject.
    if (process.platform !== 'win32' && result.status === 'failed') {
      throw result.error;
    }
  }

  // Instance state
  /**
   * Node ChildProcess instances retained for backward compatibility. New
   * structural process handles are kept in the internal lifecycle registry.
   */
  public processesToEnd = new plugins.lik.ObjectMap<plugins.childProcess.ChildProcess>();
  public cleanupFunctions = new plugins.lik.ObjectMap<() => Promise<unknown>>();
  private structuralProcessHandlesToEnd = new plugins.lik.ObjectMap<IProcessHandle>();
  private trackedPidGenerations = new Map<number, symbol>();
  private processHandleTracking = new WeakMap<IProcessHandle, IProcessHandleTracking>();
  private terminatingPidGenerations = new Map<number, symbol>();
  private suppressTrackedPidSetNotifications = false;
  /**
   * PIDs tracked independently for tree-killing during shutdown. Removing a
   * PID from this public ownership set also withdraws it from later escalation.
   */
  public trackedPids: Set<number> = new TrackedPidSet(
    (pidArg) => {
      this.ensureTrackedPidGeneration(pidArg);
    },
    (pidArg) => {
      if (!this.suppressTrackedPidSetNotifications) {
        this.releaseTrackedPid(pidArg);
      }
    },
  );
  private options: Required<ISmartExitOptions>;
  private killAllPromise: Promise<IKillAllResult> | undefined;
  private killAllTerminationController: AbortController | undefined;
  private killAllSignalDisposers = new Set<() => void>();

  private ensureTrackedPidGeneration(pidArg: number): symbol {
    let generation = this.trackedPidGenerations.get(pidArg);
    if (!generation) {
      generation = Symbol(`tracked-pid-${pidArg}`);
      this.trackedPidGenerations.set(pidArg, generation);
    }
    return generation;
  }

  private getProcessHandlesToEnd(): IProcessHandle[] {
    return [
      ...this.processesToEnd.getArray(),
      ...this.structuralProcessHandlesToEnd.getArray(),
    ];
  }

  private addProcessHandleToRegistry(processHandleArg: IProcessHandle): void {
    if (processHandleArg instanceof plugins.childProcess.ChildProcess) {
      this.processesToEnd.add(processHandleArg);
    } else {
      this.structuralProcessHandlesToEnd.add(processHandleArg);
    }
  }

  private removeProcessHandleFromRegistry(processHandleArg: IProcessHandle): void {
    this.structuralProcessHandlesToEnd.remove(processHandleArg);
    if (processHandleArg instanceof plugins.childProcess.ChildProcess) {
      this.processesToEnd.remove(processHandleArg);
    }
  }

  private validateProcessHandlePid(processHandleArg: IProcessHandle): number | undefined {
    const pid = processHandleArg.pid;
    if (pid === undefined) {
      return undefined;
    }
    if (!Number.isSafeInteger(pid) || pid <= 0) {
      throw new TypeError('A process handle PID must be a positive safe integer when provided.');
    }
    return pid;
  }

  private releaseTrackedPid(pidArg: number, expectedGenerationArg?: symbol): void {
    const currentGeneration = this.trackedPidGenerations.get(pidArg);
    if (expectedGenerationArg && currentGeneration !== expectedGenerationArg) {
      return;
    }
    this.suppressTrackedPidSetNotifications = true;
    try {
      this.trackedPids.delete(pidArg);
    } finally {
      this.suppressTrackedPidSetNotifications = false;
    }
    this.trackedPidGenerations.delete(pidArg);
    if (this.terminatingPidGenerations.get(pidArg) === currentGeneration) {
      this.terminatingPidGenerations.delete(pidArg);
    }
    for (const processHandle of this.getProcessHandlesToEnd()) {
      const processHandleTracking = this.processHandleTracking.get(processHandle);
      if ((processHandleTracking?.pid ?? processHandle.pid) === pidArg) {
        this.removeProcessHandleFromRegistry(processHandle);
        this.processHandleTracking.delete(processHandle);
      }
    }
  }

  private log(message: string, isError = false): void {
    if (this.options.silent) {
      return;
    }
    const prefix = '[smartexit]';
    if (isError) {
      console.error(`${prefix} ${message}`);
    } else {
      console.log(`${prefix} ${message}`);
    }
  }

  constructor(optionsArg: ISmartExitOptions = {}) {
    const processTerminationGraceMs =
      optionsArg.processTerminationGraceMs ?? defaultProcessTerminationGraceMs;
    if (
      !Number.isSafeInteger(processTerminationGraceMs)
      || processTerminationGraceMs < 0
      || processTerminationGraceMs > maximumTimerDelayMs
    ) {
      throw new TypeError(
        `processTerminationGraceMs must be an integer between 0 and ${maximumTimerDelayMs}.`,
      );
    }
    this.options = {
      silent: optionsArg.silent ?? false,
      processTerminationGraceMs,
    };
    // Auto-register with the global ProcessLifecycle registry
    ProcessLifecycle.registerInstance(this);
  }

  /** Register a PID-bearing process handle for cleanup on shutdown. */
  public addProcess(processHandleArg: IProcessHandle): void {
    const pid = this.validateProcessHandlePid(processHandleArg);
    this.addProcessHandleToRegistry(processHandleArg);
    if (pid !== undefined) {
      const generation = Symbol(`tracked-pid-${pid}`);
      this.trackedPids.add(pid);
      this.trackedPidGenerations.set(pid, generation);
      this.processHandleTracking.set(processHandleArg, { pid, generation });
    }
  }

  /** Register an async cleanup function to run on shutdown. */
  public addCleanupFunction(cleanupFunctionArg: () => Promise<unknown>) {
    this.cleanupFunctions.add(cleanupFunctionArg);
  }

  /**
   * Unregister the same process-handle object. On POSIX, group ownership remains when
   * descendants still exist under the tracked detached group leader PID.
   */
  public removeProcess(processHandleArg: IProcessHandle): void {
    const processHandleTracking = this.processHandleTracking.get(processHandleArg);
    this.removeProcessHandleFromRegistry(processHandleArg);
    this.processHandleTracking.delete(processHandleArg);
    if (!processHandleTracking) {
      return;
    }
    const { pid, generation: processGeneration } = processHandleTracking;
    const trackedGeneration = this.trackedPidGenerations.get(pid);
    if (
      trackedGeneration !== processGeneration
      || this.terminatingPidGenerations.get(pid) === processGeneration
    ) {
      return;
    }
    if (process.platform === 'win32') {
      this.releaseTrackedPid(pid, processGeneration);
    } else {
      const probeResult = probeProcessGroupByPid(pid);
      if (probeResult.status === 'notFound') {
        this.releaseTrackedPid(pid, processGeneration);
      } else if (probeResult.status === 'failed') {
        this.log(
          `Failed to confirm process group ${pid} ended while removing its process handle: ${probeResult.error}`,
          true,
        );
      }
    }
  }

  /**
   * Run cleanup functions, then kill all tracked process trees.
   * Called by ProcessLifecycle during global shutdown.
   * Can also be called manually.
   */
  public killAll(optionsArg: IKillAllOptions = {}): Promise<IKillAllResult> {
    if (SmartExit.cleanupCallContext.getStore() === this) {
      throw new Error(
        'killAll() cannot be called from a cleanup function registered on the same SmartExit instance.',
      );
    }

    if (this.killAllPromise) {
      this.linkTerminationSignal(optionsArg.terminationSignal);
      return this.killAllPromise;
    }

    const terminationController = new AbortController();
    this.killAllTerminationController = terminationController;
    this.linkTerminationSignal(optionsArg.terminationSignal);
    // Defer the actual run by one microtask so the shared promise is installed
    // before a cleanup function can synchronously call killAll again.
    const currentKillAllPromise = Promise.resolve().then(() =>
      this.runKillAll(terminationController.signal)
    );
    this.killAllPromise = currentKillAllPromise;
    const finishKillAll = (): void => {
      if (this.killAllPromise !== currentKillAllPromise) {
        return;
      }
      this.killAllPromise = undefined;
      this.killAllTerminationController = undefined;
      for (const disposeSignal of this.killAllSignalDisposers) {
        disposeSignal();
      }
      this.killAllSignalDisposers.clear();
    };
    currentKillAllPromise.then(finishKillAll, finishKillAll);
    return currentKillAllPromise;
  }

  private linkTerminationSignal(signalArg: AbortSignal | undefined): void {
    const terminationController = this.killAllTerminationController;
    if (!signalArg || !terminationController) {
      return;
    }
    if (signalArg.aborted) {
      terminationController.abort();
      return;
    }
    const linkedSignal = signalArg;
    const linkedController = terminationController;
    const signalDisposers = this.killAllSignalDisposers;
    function disposeSignal(): void {
      linkedSignal.removeEventListener('abort', abortHandler);
      signalDisposers.delete(disposeSignal);
    }
    function abortHandler(): void {
      linkedController.abort();
      disposeSignal();
    }
    linkedSignal.addEventListener('abort', abortHandler, { once: true });
    this.killAllSignalDisposers.add(disposeSignal);
  }

  private async waitForTerminationGrace(terminationSignalArg: AbortSignal): Promise<void> {
    if (terminationSignalArg.aborted) {
      return;
    }
    await new Promise<void>((resolve) => {
      let timeoutHandle: NodeJS.Timeout | undefined;
      const finishWait = (): void => {
        if (timeoutHandle) {
          clearTimeout(timeoutHandle);
          timeoutHandle = undefined;
        }
        terminationSignalArg.removeEventListener('abort', finishWait);
        resolve();
      };
      terminationSignalArg.addEventListener('abort', finishWait, { once: true });
      if (terminationSignalArg.aborted) {
        finishWait();
        return;
      }
      timeoutHandle = setTimeout(finishWait, this.options.processTerminationGraceMs);
    });
  }

  private async runKillAll(terminationSignalArg: AbortSignal): Promise<IKillAllResult> {
    const cleanupFuncs = this.cleanupFunctions.getArray();
    let cleanupFunctionsRan = 0;

    // Phase 1: Run cleanup functions (processes still alive)
    for (const cleanupFunction of cleanupFuncs) {
      try {
        await SmartExit.cleanupCallContext.run(this, cleanupFunction);
        cleanupFunctionsRan++;
      } catch (err) {
        this.log(`Cleanup function failed: ${err}`, true);
      }
    }

    // Cleanup is intentionally complete before the shutdown snapshot. Processes
    // added or removed by cleanup functions are therefore reflected in this run.
    const trackedProcessSnapshot = this.getProcessHandlesToEnd().map((processHandle) => ({
      processHandle,
      generation: this.processHandleTracking.get(processHandle)?.generation,
    }));
    const trackedPidSnapshot = Array.from(this.trackedPids, (pid) => ({
      pid,
      generation: this.ensureTrackedPidGeneration(pid),
    }));
    const signalledPids = new Set<number>();
    const releasedPidGenerations = new Map<number, symbol>();
    for (const { pid, generation } of trackedPidSnapshot) {
      this.terminatingPidGenerations.set(pid, generation);
    }

    try {
      if (process.platform === 'win32') {
        // taskkill /T /F is already a forced tree termination. A later PID-based
        // escalation could target a recycled PID, so Windows has one phase only.
        for (const { pid, generation } of trackedPidSnapshot) {
          const signalResult = signalProcessTreeByPid(pid, 'SIGTERM');
          if (signalResult.status === 'signalled') {
            signalledPids.add(pid);
            releasedPidGenerations.set(pid, generation);
          } else if (signalResult.status === 'notFound') {
            releasedPidGenerations.set(pid, generation);
          } else if (signalResult.status === 'failed') {
            this.log(`Failed to terminate process tree ${pid}: ${signalResult.error}`, true);
          }
        }
      } else {
        // Signal all owned groups first, then share one bounded grace interval.
        const pendingPosixPidGenerations = new Map<number, symbol>();
        for (const { pid, generation } of trackedPidSnapshot) {
          const signalResult = signalProcessTreeByPid(pid, 'SIGTERM');
          if (signalResult.status === 'signalled') {
            signalledPids.add(pid);
            pendingPosixPidGenerations.set(pid, generation);
          } else if (signalResult.status === 'notFound') {
            releasedPidGenerations.set(pid, generation);
          } else if (signalResult.status === 'failed') {
            pendingPosixPidGenerations.set(pid, generation);
            this.log(`Failed to signal process group ${pid} with SIGTERM: ${signalResult.error}`, true);
          }
        }

        if (pendingPosixPidGenerations.size > 0) {
          await this.waitForTerminationGrace(terminationSignalArg);
        }

        for (const [pid, generation] of pendingPosixPidGenerations) {
          // A direct forced kill may have released this ownership while the
          // shared graceful interval was in flight. Do not target a PID or
          // process-group ID that may since have been reused.
          if (
            !this.trackedPids.has(pid)
            || this.trackedPidGenerations.get(pid) !== generation
          ) {
            continue;
          }
          const probeResult = probeProcessGroupByPid(pid);
          if (probeResult.status === 'notFound') {
            releasedPidGenerations.set(pid, generation);
            continue;
          }
          if (probeResult.status === 'failed') {
            this.log(`Failed to probe process group ${pid}: ${probeResult.error}`, true);
          }

          const signalResult = signalProcessTreeByPid(pid, 'SIGKILL');
          if (signalResult.status === 'signalled') {
            signalledPids.add(pid);
            releasedPidGenerations.set(pid, generation);
          } else if (signalResult.status === 'notFound') {
            releasedPidGenerations.set(pid, generation);
          } else if (signalResult.status === 'failed') {
            this.log(`Failed to signal process group ${pid} with SIGKILL: ${signalResult.error}`, true);
          }
        }
      }
    } finally {
      // Release only snapshot PIDs that are absent or received a successful
      // terminating OS operation. Failed kills and distinct PIDs registered
      // during grace remain owned for a later killAll or exit pass.
      for (const [pid, generation] of releasedPidGenerations) {
        this.releaseTrackedPid(pid, generation);
      }
      for (const { pid, generation } of trackedPidSnapshot) {
        if (this.terminatingPidGenerations.get(pid) === generation) {
          this.terminatingPidGenerations.delete(pid);
        }
      }
      for (const { processHandle, generation } of trackedProcessSnapshot) {
        if (this.processHandleTracking.get(processHandle)?.generation === generation) {
          this.removeProcessHandleFromRegistry(processHandle);
        }
      }
    }

    return { processesKilled: signalledPids.size, cleanupFunctionsRan };
  }

  /** Remove this instance from the global ProcessLifecycle registry. */
  public deregister(): void {
    ProcessLifecycle.deregisterInstance(this);
  }
}