import { spawn, ChildProcess } from 'child_process';
import psTree from 'ps-tree';
import pidusage from 'pidusage';

interface IMonitorConfig {
  name?: string;                 // Optional name to identify the instance
  projectDir: string;            // Directory where the command will run
  command: string;               // Full command to run (e.g., "npm run xyz")
  args?: string[];               // Optional: arguments for the command
  memoryLimitBytes: number;      // Maximum allowed memory (in bytes) for the process group
  monitorIntervalMs?: number;    // Interval (in ms) at which memory is checked (default: 5000)
}

class ProcessMonitor {
  private child: ChildProcess | null = null;
  private config: IMonitorConfig;
  private intervalId: NodeJS.Timeout | null = null;
  private stopped: boolean = true; // Initially stopped until start() is called

  constructor(config: IMonitorConfig) {
    this.config = config;
  }

  public start(): void {
    // Reset the stopped flag so that new processes can spawn.
    this.stopped = false;
    this.log(`Starting process monitor.`);
    this.spawnChild();

    // Set the monitoring interval.
    const interval = this.config.monitorIntervalMs || 5000;
    this.intervalId = setInterval(() => {
      if (this.child && this.child.pid) {
        this.monitorProcessGroup(this.child.pid, this.config.memoryLimitBytes);
      }
    }, interval);
  }

  private spawnChild(): void {
    // Don't spawn if the monitor has been stopped.
    if (this.stopped) return;

    if (this.config.args && this.config.args.length > 0) {
      this.log(
        `Spawning command "${this.config.command}" with args [${this.config.args.join(
          ', '
        )}] in directory: ${this.config.projectDir}`
      );
      this.child = spawn(this.config.command, this.config.args, {
        cwd: this.config.projectDir,
        detached: true,
        stdio: 'inherit',
      });
    } else {
      this.log(
        `Spawning command "${this.config.command}" in directory: ${this.config.projectDir}`
      );
      // Use shell mode to allow a full command string.
      this.child = spawn(this.config.command, {
        cwd: this.config.projectDir,
        detached: true,
        stdio: 'inherit',
        shell: true,
      });
    }

    this.log(`Spawned process with PID ${this.child.pid}`);

    // When the child process exits, restart it if the monitor isn't stopped.
    this.child.on('exit', (code, signal) => {
      this.log(`Child process exited with code ${code}, signal ${signal}.`);
      if (!this.stopped) {
        this.log('Restarting process...');
        this.spawnChild();
      }
    });
  }

  /**
   * Monitor the process group’s memory usage. If the total memory exceeds the limit,
   * kill the process group so that the 'exit' handler can restart it.
   */
  private async monitorProcessGroup(pid: number, memoryLimit: number): Promise<void> {
    try {
      const memoryUsage = await this.getProcessGroupMemory(pid);
      this.log(
        `Current memory usage for process group (PID ${pid}): ${this.humanReadableBytes(
          memoryUsage
        )} (${memoryUsage} bytes)`
      );
      if (memoryUsage > memoryLimit) {
        this.log(
          `Memory usage ${this.humanReadableBytes(
            memoryUsage
          )} exceeds limit of ${this.humanReadableBytes(memoryLimit)}. Restarting process.`
        );
        // Kill the entire process group by sending a signal to -PID.
        process.kill(-pid, 'SIGKILL');
      }
    } catch (error) {
      this.log('Error monitoring process group: ' + error);
    }
  }

  /**
   * Get the total memory usage (in bytes) for the process group (the main process and its children).
   */
  private getProcessGroupMemory(pid: number): Promise<number> {
    return new Promise((resolve, reject) => {
      psTree(pid, (err, children) => {
        if (err) return reject(err);
        // Include the main process and its children.
        const pids: number[] = [pid, ...children.map(child => Number(child.PID))];
        pidusage(pids, (err, stats) => {
          if (err) return reject(err);
          let totalMemory = 0;
          for (const key in stats) {
            totalMemory += stats[key].memory;
          }
          resolve(totalMemory);
        });
      });
    });
  }

  /**
   * Convert a number of bytes into a human-readable string (e.g. "1.23 MB").
   */
  private humanReadableBytes(bytes: number, decimals: number = 2): string {
    if (bytes === 0) return '0 Bytes';
    const k = 1024;
    const dm = decimals < 0 ? 0 : decimals;
    const sizes = ['Bytes', 'KB', 'MB', 'GB', 'TB', 'PB', 'EB', 'ZB', 'YB'];
    const i = Math.floor(Math.log(bytes) / Math.log(k));
    return parseFloat((bytes / Math.pow(k, i)).toFixed(dm)) + ' ' + sizes[i];
  }

  /**
   * Stop the monitor and prevent any further respawns.
   */
  public stop(): void {
    this.log('Stopping process monitor.');
    this.stopped = true;
    if (this.intervalId) {
      clearInterval(this.intervalId);
    }
    if (this.child && this.child.pid) {
      process.kill(-this.child.pid, 'SIGKILL');
    }
  }

  /**
   * Helper method for logging messages with the instance name.
   */
  private log(message: string): void {
    const prefix = this.config.name ? `[${this.config.name}] ` : '';
    console.log(prefix + message);
  }
}

// Example usage:
const config: IMonitorConfig = {
  name: 'Project XYZ Monitor',        // Identifier for the instance
  projectDir: '/path/to/your/project',  // Set the project directory here
  command: 'npm run xyz',               // Full command string (no need for args)
  memoryLimitBytes: 500 * 1024 * 1024,    // 500 MB memory limit
  monitorIntervalMs: 5000,              // Check memory usage every 5 seconds
};

const monitor = new ProcessMonitor(config);
monitor.start();

// Ensure that on process exit (e.g. Ctrl+C) we clean up the child process and prevent respawns.
process.on('SIGINT', () => {
  monitor.log('Received SIGINT, stopping monitor...');
  monitor.stop();
  process.exit();
});