service/service.ts

#!/usr/bin/env bun

import { execSync, spawn } from 'child_process';
import * as fs from 'fs';
import { dirname, basename, join, resolve } from 'path';
import * as zlib from 'zlib';
import * as readline from 'readline';
import type { Readable } from 'stream';

type RunnerFunction = (options: RunnerOptions) => (() => void) | void;

type Logger = {
  log: (...args: any[]) => void;
  trace: (...args: any[]) => void;
  debug: (...args: any[]) => void;
  info: (...args: any[]) => void;
  warn: (...args: any[]) => void;
  error: (...args: any[]) => void;
};

type RunnerOptions = {
  logger: Logger;
};

type ServiceOptions = {
  name: string;
  description: string;
  runner: RunnerFunction;
};

export default function service(options: ServiceOptions) {
  const serviceName = options.name;
  const serviceDescription = options.description;
  const runner = options.runner;

  // Paths
  const lockFile = `/tmp/${serviceName}.pid`;
  const logDir = `/tmp/`;
  const logFile = join(logDir, `${serviceName}.log`);
  const maxLogSize = 5 * 1024 * 1024; // 5 MB per log file
  const maxTotalLogSize = 20 * 1024 * 1024; // 20 MB total for all logs
  const serviceConfig = `/etc/systemd/system/${serviceName}.service`;

  // Get the absolute path to the current executable
  const selfFile = resolve(process.argv0);
  const selfDir = dirname(selfFile);
  const selfFilename = basename(selfFile);

  // Function to check if systemd is available
  function hasSystemd() {
    try {
      execSync('which systemctl', { stdio: 'ignore' });
      return fs.existsSync('/run/systemd/system');
    } catch {
      return false;
    }
  }

  // Function to check if the systemd service file exists
  function hasSystemdService() {
    return fs.existsSync(serviceConfig);
  }

  // Log stream variable
  let logStream: fs.WriteStream | null = null;

  // Function to create the logger
  function createLogger(): Logger {
    // Determine if running under systemd
    const useSystemdLogging = hasSystemdService();

    // If not using systemd, set up log file
    if (!useSystemdLogging) {
      // Open the log file for appending
      logStream = fs.createWriteStream(logFile, { flags: 'a' });
    }

    // Function to write logs and handle rotation
    function writeLog(message: string, severity: string = 'INFO') {
      const formattedMessage = useSystemdLogging
        // Systemd includes its own timestamp
        ? `[${severity}] ${message}\n`
        // Custom timestamp
        : `${new Date().toISOString()} [${severity}] ${message}\n`;

      if (useSystemdLogging) {
        // Write to stdout or stderr based on severity
        if (severity === 'ERROR' || severity === 'WARN') {
          process.stderr.write(formattedMessage);
        } else {
          process.stdout.write(formattedMessage);
        }
      } else {
        rotateLogs();
        if (logStream) {
          try {
            logStream.write(formattedMessage);
          } catch {
            // If writing fails, fallback to stdout/stderr
            process.stdout.write(formattedMessage);
          }
        } else {
          // If logStream is closed, write to stdout/stderr
          process.stdout.write(formattedMessage);
        }
      }
    }

    // Return the logger object
    const logger: Logger = {
      log: (...args: any[]) => {
        writeLog(args.join(' '), 'INFO');
      },
      trace: (...args: any[]) => {
        writeLog(args.join(' '), 'TRACE');
      },
      debug: (...args: any[]) => {
        writeLog(args.join(' '), 'DEBUG');
      },
      info: (...args: any[]) => {
        writeLog(args.join(' '), 'INFO');
      },
      warn: (...args: any[]) => {
        writeLog(args.join(' '), 'WARN');
      },
      error: (...args: any[]) => {
        writeLog(args.join(' '), 'ERROR');
      },
    };

    return logger;
  }

  // Function to close the logger
  function closeLogger() {
    if (logStream) {
      logStream.end();
      logStream = null;
    }
  }

  // Function to rotate logs
  function rotateLogs() {
    if (fs.existsSync(logFile)) {
      const stats = fs.statSync(logFile);
      if (stats.size >= maxLogSize) {
        // Close the existing log stream before renaming
        if (logStream) {
          logStream.end(); // Close the stream
          logStream = null;
        }

        // Determine the rotated log file name
        const timestamp = new Date().toISOString().replace(/[:.]/g, '-');
        const rotatedLogFile = `${logFile}.${timestamp}`;

        // Rename the current log file
        fs.renameSync(logFile, rotatedLogFile);

        // Compress the rotated log file
        try {
          const gzip = zlib.createGzip();
          const source = fs.createReadStream(rotatedLogFile);
          const destination = fs.createWriteStream(`${rotatedLogFile}.gz`);

          source.pipe(gzip).pipe(destination);

          source.on('end', () => {
            fs.unlinkSync(rotatedLogFile);
          });
        } catch (error) {
          console.error('Error compressing log file:', error);
        }

        // Manage total log size
        manageTotalLogSize();

        // Re-open the log stream
        logStream = fs.createWriteStream(logFile, { flags: 'a' });
      }
    }
  }

  // Function to manage total log size
  function manageTotalLogSize() {
    try {
      const files = fs.readdirSync(logDir);
      const logFiles = files
        .filter((file) => file.startsWith(`${serviceName}.log`))
        .map((file) => {
          const filePath = join(logDir, file);
          const stats = fs.lstatSync(filePath);
          return {
            name: file,
            path: filePath,
            time: stats.mtime.getTime(),
            size: stats.size,
          };
        });

      logFiles.sort((a, b) => b.time - a.time); // Sort by most recent

      // Calculate total size
      let totalSize = logFiles.reduce((acc, file) => acc + file.size, 0);

      // Delete oldest logs until total size is under maxTotalLogSize
      for (let i = logFiles.length - 1; i >= 0 && totalSize > maxTotalLogSize; i--) {
        try {
          fs.unlinkSync(logFiles[i].path);
          totalSize -= logFiles[i].size;
        } catch (error) {
          console.error('Error deleting log file:', error);
        }
      }
    } catch (err) {
      console.error('Error managing total log size:', err);
    }
  }

  // Daemon function
  function daemon() {
    // Create the logger
    const logger = createLogger();

    logger.info('Daemon started...');

    // Run the user-provided runner function
    let cleanupFunction: (() => void) | void;
    try {
      cleanupFunction = runner({ logger });
    } catch (err) {
      logger.error('Error in runner function:', err);
      process.exit(1);
    }

    // Signal handling
    function cleanUp() {
      logger.info('Daemon is stopping...');
      try {
        if (cleanupFunction) {
          cleanupFunction();
        }
      } catch (err) {
        logger.error('Error in cleanup function:', err);
      }

      closeLogger();

      try {
        fs.unlinkSync(lockFile);
      } catch {
        // Ignore errors
      }

      process.exit(0);
    }

    process.on('SIGTERM', cleanUp);
    process.on('SIGINT', cleanUp);
    process.on('SIGQUIT', cleanUp);
    process.on('uncaughtException', (err) => {
      logger.error('Uncaught exception:', err);
      cleanUp();
    });

    // Keep the process running
    process.stdin.resume();
  }

  // Install function
  function install() {
    console.log(`Installing ${serviceName}...`);
    if (!hasSystemd()) {
      console.error('Systemd is not available on this system.');
      process.exit(1);
    }

    const serviceConfigContent = `
[Unit]
Description=${serviceDescription}
After=network.target

[Service]
ExecStart=${selfFile} daemon
Environment="NODE_ENV=production" "_=${selfFile}"
WorkingDirectory=${selfDir}
User=root
Restart=always
KillSignal=SIGTERM
TimeoutStopSec=5
SyslogIdentifier=${serviceName}

[Install]
WantedBy=multi-user.target
`;

    try {
      fs.writeFileSync(`/tmp/${serviceName}.service`, serviceConfigContent);
      execSync(`sudo mv /tmp/${serviceName}.service ${serviceConfig}`, {
        stdio: 'inherit',
      });
      execSync('sudo systemctl daemon-reload', { stdio: 'inherit' });
      execSync(`sudo systemctl enable ${serviceName}`, { stdio: 'inherit' });
      start();
    } catch (error) {
      console.error('Error during installation:', error);
      process.exit(1);
    }
  }

  // Uninstall function
  function uninstall() {
    console.log(`Uninstalling ${serviceName}...`);
    if (!hasSystemdService()) {
      console.error('Service is not installed via systemd.');
      process.exit(1);
    }

    try {
      execSync(`sudo systemctl stop ${serviceName}`, { stdio: 'inherit' });
      execSync(`sudo systemctl disable ${serviceName}`, { stdio: 'inherit' });
      execSync(`sudo rm ${serviceConfig}`, { stdio: 'inherit' });
      execSync('sudo systemctl daemon-reload', { stdio: 'inherit' });
    } catch (error) {
      console.error('Error during uninstallation:', error);
      process.exit(1);
    }
  }

  // Start function
  function start() {
    if (hasSystemdService()) {
      console.log(`Starting ${serviceName} via systemd...`);
      try {
        execSync(`sudo systemctl start ${serviceName}`, { stdio: 'inherit' });
      } catch (error) {
        console.error('Error starting service via systemd:', error);
        process.exit(1);
      }
    } else {
      console.log(`Starting ${serviceName} manually...`);
      if (isRunning()) {
        console.log(`${serviceName} is already running.`);
        process.exit(0);
      }

      const child = spawn(selfFile, ['daemon'], {
        detached: true,
        cwd: selfDir,
        stdio: ['ignore', 'ignore', 'ignore'],
      });

      if (!child.pid) {
        console.error('Failed to start the service.');
        process.exit(1);
      }

      fs.writeFileSync(lockFile, child.pid.toString());
      console.log(`${serviceName} started with PID ${child.pid}`);
      child.unref();
    }
  }

  // Stop function
  function stop() {
    if (hasSystemdService()) {
      console.log(`Stopping ${serviceName} via systemd...`);
      try {
        execSync(`sudo systemctl stop ${serviceName}`, { stdio: 'inherit' });
      } catch (error) {
        console.error('Error stopping service via systemd:', error);
        process.exit(1);
      }
    } else {
      console.log(`Stopping ${serviceName} manually...`);
      if (!fs.existsSync(lockFile)) {
        console.log(`${serviceName} is not running.`);
        process.exit(0);
      }

      const pidStr = fs.readFileSync(lockFile, 'utf-8');
      const pid = parseInt(pidStr, 10);
      try {
        process.kill(pid, 'SIGTERM');
        console.log(`Stopped ${serviceName} with PID ${pid}.`);
      } catch (err) {
        console.error(`Failed to stop ${serviceName}:`, err);
      }

      try {
        fs.unlinkSync(lockFile);
      } catch {
        // Ignore errors
      }
    }
  }

  // Restart function
  function restart() {
    if (hasSystemdService()) {
      console.log(`Restarting ${serviceName} via systemd...`);
      try {
        execSync(`sudo systemctl restart ${serviceName}`, { stdio: 'inherit' });
      } catch (error) {
        console.error('Error restarting service via systemd:', error);
        process.exit(1);
      }
    } else {
      console.log(`Restarting ${serviceName} manually...`);
      stop();
      // Give some time for the process to stop
      setTimeout(() => {
        start();
      }, 1000);
    }
  }

  // Status function
  function status() {
    if (hasSystemdService()) {
      try {
        execSync(`systemctl status ${serviceName}`, { stdio: 'inherit' });
      } catch {
        console.log(`${serviceName} is not running via systemd.`);
      }
    } else {
      if (isRunning()) {
        const pid = fs.readFileSync(lockFile, 'utf-8');
        console.log(`${serviceName} is running with PID ${pid}.`);
      } else {
        console.log(`${serviceName} is not running.`);
      }
    }
  }

  // Check if the process is running
  function isRunning() {
    try {
      const pidStr = fs.readFileSync(lockFile, 'utf-8');
      const pid = parseInt(pidStr, 10);
      process.kill(pid, 0);
      return true;
    } catch {
      return false;
    }
  }

  // Function to parse command-line options
  function parseLogOptions(args?: string): { follow: boolean; lines: number } {
    const options = { follow: false, lines: 10 }; // Default values
    const argsArray = args ? args.split(' ') : [];

    argsArray.forEach((arg, index) => {
      if (arg === '-f' || arg === '--follow') {
        options.follow = true;
      } else if (arg === '-n' || arg === '--lines') {
        const value = argsArray[index + 1];
        if (value && !isNaN(parseInt(value, 10))) {
          options.lines = parseInt(value, 10);
        }
      }
    });

    return options;
  }

  // Custom log viewer function
  function logs(args?: string) {
    const options = parseLogOptions(args);

    if (hasSystemdService()) {
      try {
        execSync(`journalctl -u ${serviceName} ${args || ''}`, {
          stdio: 'inherit',
        });
      } catch (error) {
        console.error('Error fetching logs via systemd:', error);
        process.exit(1);
      }
    } else {
      console.log(`Fetching logs manually...`);

      // Get all log files (including compressed ones)
      try {
        const files = fs.readdirSync(logDir);
        const logFiles = files
          .filter((file) => file.startsWith(`${serviceName}.log`))
          .map((file) => {
            const filePath = join(logDir, file);
            const stats = fs.lstatSync(filePath);
            return {
              name: file,
              path: filePath,
              time: stats.mtime.getTime(),
            };
          });

        logFiles.sort((a, b) => a.time - b.time); // Sort by oldest first

        if (logFiles.length > 0) {
          if (options.follow) {
            // Real-time log following
            followLogs();
          } else {
            // Display historical logs
            displayLogs(logFiles, options.lines);
          }
        } else {
          console.log('No logs found.');
        }
      } catch (error) {
        console.error('Error fetching logs:', error);
      }
    }
  }

  // Function to display historical logs
  function displayLogs(logFiles: { path: string }[], lines: number) {
    const allLines: string[] = [];

    for (const file of logFiles) {
      let fileStream: Readable;
      if (file.path.endsWith('.gz')) {
        fileStream = fs.createReadStream(file.path).pipe(zlib.createGunzip());
      } else {
        fileStream = fs.createReadStream(file.path);
      }

      const rl = readline.createInterface({
        input: fileStream,
        crlfDelay: Infinity,
      });

      rl.on('line', (line) => {
        allLines.push(line);
      });

      rl.on('close', () => {
        // After reading all lines, display the last 'n' lines
        const start = Math.max(allLines.length - lines, 0);
        for (let i = start; i < allLines.length; i++) {
          console.log(allLines[i]);
        }
      });
    }
  }

  // Function to follow logs in real-time
  function followLogs() {
    let currentLogFile = logFile;
    let position = 0;

    const readNewData = () => {
      fs.stat(currentLogFile, (err, stats) => {
        if (err) {
          // File might not exist yet
          return;
        }

        if (stats.size < position) {
          // Log file was rotated
          position = 0;
        }

        const stream = fs.createReadStream(currentLogFile, {
          encoding: 'utf-8',
          flags: 'r',
          start: position,
        });

        stream.on('data', (chunk) => {
          process.stdout.write(chunk as string);
          position += Buffer.byteLength(chunk, 'utf-8');
        });

        stream.on('error', (err) => {
          console.error('Error reading log file:', err);
        });
      });
    };

    // Initial read to get any existing data
    readNewData();

    // Watch the current log file
    let watcher = fs.watch(currentLogFile, (eventType) => {
      if (eventType === 'change') {
        readNewData();
      } else if (eventType === 'rename') {
        // Log file was rotated
        watcher.close();
        position = 0;
        currentLogFile = logFile; // The new log file
        // Start watching the new log file
        watcher = fs.watch(currentLogFile, (eventType) => {
          if (eventType === 'change') {
            readNewData();
          }
        });
      }
    });

    // Keep the process running
    process.stdin.resume();
  }

  // Command-line interface
  const action = process.argv[2];
  const args = process.argv.slice(3).join(' ');

  switch (action) {
    case 'install':
      install();
      break;
    case 'uninstall':
      uninstall();
      break;
    case 'start':
      start();
      break;
    case 'stop':
      stop();
      break;
    case 'restart':
      restart();
      break;
    case 'status':
      status();
      break;
    case 'logs':
      logs(args);
      break;
    case 'daemon':
      daemon();
      break;
    default:
      console.log(
        `Usage: ${selfFilename} {install|uninstall|start|stop|restart|status|logs [args]}`
      );
      console.log(`Description: ${serviceDescription}`);
  }
}