gerbil/src/main/managers/SillyTavernManager.ts

import { spawn } from 'child_process';
import { createServer, request, type Server } from 'http';
import { homedir } from 'os';
import { join } from 'path';
import { readFileSync, writeFileSync, existsSync } from 'fs';
import type { ChildProcess } from 'child_process';

import { LogManager } from './LogManager';
import { WindowManager } from './WindowManager';
import { SILLYTAVERN } from '@/constants';
import { terminateProcess } from '@/utils/process';

export interface SillyTavernConfig {
  name: string;
  port: number;
  proxyPort?: number;
}

export class SillyTavernManager {
  private sillyTavernProcess: ChildProcess | null = null;
  private proxyServer: Server | null = null;
  private logManager: LogManager;
  private windowManager: WindowManager;

  constructor(logManager: LogManager, windowManager: WindowManager) {
    this.logManager = logManager;
    this.windowManager = windowManager;

    process.on('SIGINT', () => {
      this.cleanup().catch(() => {
        void 0;
      });
    });

    process.on('SIGTERM', () => {
      this.cleanup().catch(() => {
        void 0;
      });
    });
  }

  private detectedDataRoot: string | null = null;

  private getFallbackDataRoot(): string {
    const platform = process.platform;
    const home = homedir();

    switch (platform) {
      case 'win32':
        return join(home, 'AppData', 'Local', 'SillyTavern', 'Data', 'data');
      case 'darwin':
        return join(
          home,
          'Library',
          'Application Support',
          'SillyTavern',
          'data'
        );
      case 'linux':
      default:
        return join(home, '.local', 'share', 'SillyTavern', 'data');
    }
  }

  private getSillyTavernDataRoot(): string {
    if (this.detectedDataRoot) {
      return this.detectedDataRoot;
    }

    const fallback = this.getFallbackDataRoot();

    if (existsSync(fallback)) {
      this.detectedDataRoot = fallback;
      return fallback;
    }

    const platform = process.platform;
    const home = homedir();
    const alternatePaths = [];

    switch (platform) {
      case 'win32':
        alternatePaths.push(
          join(home, 'AppData', 'Local', 'SillyTavern', 'data'),
          join(home, '.sillytavern', 'data')
        );
        break;
      case 'darwin':
        alternatePaths.push(join(home, '.sillytavern', 'data'));
        break;
      case 'linux':
      default:
        alternatePaths.push(join(home, '.sillytavern', 'data'));
        break;
    }

    for (const path of alternatePaths) {
      if (existsSync(path)) {
        this.detectedDataRoot = path;
        return path;
      }
    }

    this.detectedDataRoot = fallback;
    return fallback;
  }

  private getSillyTavernSettingsPath(): string {
    return join(this.getSillyTavernDataRoot(), 'default-user', 'settings.json');
  }

  private static readonly SILLYTAVERN_BASE_ARGS = [
    'sillytavern',
    '--listen',
    '--browserLaunchEnabled',
    'false',
    '--disableCsrf',
  ];

  async isNpxAvailable(): Promise<boolean> {
    try {
      const testProcess = spawn('npx', ['--version'], { stdio: 'pipe' });

      return new Promise<boolean>((resolve) => {
        const timeout = setTimeout(() => {
          testProcess.kill();
          resolve(false);
        }, 5000);

        testProcess.on('exit', (code) => {
          clearTimeout(timeout);
          resolve(code === 0);
        });

        testProcess.on('error', () => {
          clearTimeout(timeout);
          resolve(false);
        });
      });
    } catch {
      return false;
    }
  }

  private createNpxProcess(args: string[]): ChildProcess {
    return spawn('npx', args, {
      stdio: ['pipe', 'pipe', 'pipe'],
      detached: false,
      env: {
        ...process.env,
        DATA_ROOT: this.getSillyTavernDataRoot(),
      },
    });
  }

  private async ensureSillyTavernSettings(): Promise<void> {
    const settingsPath = this.getSillyTavernSettingsPath();

    if (existsSync(settingsPath)) {
      this.windowManager.sendKoboldOutput('SillyTavern settings found');
      return;
    }

    this.windowManager.sendKoboldOutput(
      'SillyTavern settings not found, starting SillyTavern briefly to generate config...'
    );

    return new Promise((resolve, reject) => {
      const initProcess = this.createNpxProcess(
        SillyTavernManager.SILLYTAVERN_BASE_ARGS
      );

      let hasResolved = false;

      initProcess.on('exit', (code: number | null, signal: string | null) => {
        this.windowManager.sendKoboldOutput(
          signal
            ? `SillyTavern init process terminated with signal ${signal}`
            : `SillyTavern init process exited with code ${code}`
        );

        if (!hasResolved) {
          hasResolved = true;

          if (code !== 0) {
            const errorMsg =
              code === 4294963214
                ? 'SillyTavern failed to install due to EBUSY error (resource busy or locked). This is a critical error.'
                : `SillyTavern initialization failed with exit code ${code}`;

            this.logManager.logError(
              'SillyTavern initialization failed:',
              new Error(errorMsg)
            );
            this.windowManager.sendKoboldOutput(`CRITICAL ERROR: ${errorMsg}`);
            reject(new Error(errorMsg));
          } else {
            this.windowManager.sendKoboldOutput(
              'SillyTavern settings should now be generated'
            );
            resolve();
          }
        }
      });

      initProcess.on('error', (error) => {
        if (!hasResolved) {
          hasResolved = true;
          this.logManager.logError(
            'Failed to initialize SillyTavern settings:',
            error
          );
          this.windowManager.sendKoboldOutput(
            `SillyTavern initialization error: ${error.message}`
          );
          reject(error);
        }
      });

      if (initProcess.stdout) {
        initProcess.stdout.on('data', (data: Buffer) => {
          const output = data.toString();

          if (output.includes('SillyTavern is listening')) {
            setTimeout(async () => {
              if (!initProcess.killed && !hasResolved) {
                hasResolved = true;

                await terminateProcess(initProcess, {
                  logError: (message, error) =>
                    this.logManager.logError(message, error),
                });

                this.windowManager.sendKoboldOutput(
                  'SillyTavern settings should now be generated'
                );

                resolve();
              }
            }, 2000);
          }
        });
      }

      if (initProcess.stderr) {
        initProcess.stderr.on('data', (data: Buffer) => {
          this.windowManager.sendKoboldOutput(data.toString().trim());
        });
      }
    });
  }

  private parseKoboldConfig(args: string[]): { host: string; port: number } {
    let host = 'localhost';
    let port = 5001;

    for (let i = 0; i < args.length - 1; i++) {
      if (args[i] === '--hostname' || args[i] === '--host') {
        host = args[i + 1];
      } else if (args[i] === '--port') {
        const parsedPort = parseInt(args[i + 1], 10);
        if (!isNaN(parsedPort)) {
          port = parsedPort;
        }
      }
    }

    return { host, port };
  }

  private async setupSillyTavernConfig(
    koboldHost: string,
    koboldPort: number
  ): Promise<void> {
    try {
      const configPath = this.getSillyTavernSettingsPath();

      this.windowManager.sendKoboldOutput(
        `Configuring SillyTavern settings at: ${configPath}`
      );

      let settings: Record<string, unknown> = {};

      if (existsSync(configPath)) {
        try {
          const content = readFileSync(configPath, 'utf-8');
          settings = JSON.parse(content) as Record<string, unknown>;
          this.windowManager.sendKoboldOutput(
            `Loaded existing SillyTavern settings`
          );
        } catch {
          this.windowManager.sendKoboldOutput(
            `Could not read existing settings, creating new ones`
          );
        }
      }

      const koboldUrl = `http://${koboldHost}:${koboldPort}`;

      if (!settings.power_user) settings.power_user = {};
      const powerUser = settings.power_user as Record<string, unknown>;

      if (!settings.textgenerationwebui_settings)
        settings.textgenerationwebui_settings = {};
      const textgenSettings = settings.textgenerationwebui_settings as Record<
        string,
        unknown
      >;

      if (!textgenSettings.server_urls) textgenSettings.server_urls = {};
      const serverUrls = textgenSettings.server_urls as Record<string, unknown>;
      serverUrls.koboldcpp = koboldUrl;

      settings.main_api = 'textgenerationwebui';
      textgenSettings.type = 'koboldcpp';
      powerUser.auto_connect = true;

      writeFileSync(configPath, JSON.stringify(settings, null, 2), 'utf-8');

      this.windowManager.sendKoboldOutput(
        `SillyTavern configuration updated successfully!`
      );
    } catch (error) {
      this.logManager.logError(
        'Failed to setup SillyTavern config:',
        error as Error
      );
      this.windowManager.sendKoboldOutput(
        `Failed to configure SillyTavern: ${error instanceof Error ? error.message : String(error)}`
      );
    }
  }

  private async waitForSillyTavernToStart(_port: number): Promise<void> {
    this.windowManager.sendKoboldOutput('Waiting for SillyTavern to start...');

    return new Promise((resolve, reject) => {
      const timeout = setTimeout(() => {
        reject(new Error('SillyTavern failed to start within 30 seconds'));
      }, 30000);

      const checkForOutput = (data: Buffer) => {
        if (data.toString().includes('SillyTavern is listening')) {
          clearTimeout(timeout);
          this.windowManager.sendKoboldOutput('SillyTavern is now running!');
          resolve();

          if (this.sillyTavernProcess?.stdout) {
            this.sillyTavernProcess.stdout.removeListener(
              'data',
              checkForOutput
            );
          }
        }
      };

      if (this.sillyTavernProcess?.stdout) {
        this.sillyTavernProcess.stdout.on('data', checkForOutput);
      } else {
        clearTimeout(timeout);
        reject(new Error('SillyTavern process stdout not available'));
      }
    });
  }

  private createProxyServer(targetPort: number, proxyPort: number): void {
    this.proxyServer = createServer((req, res) => {
      const options = {
        hostname: 'localhost',
        port: targetPort,
        path: req.url,
        method: req.method,
        headers: req.headers,
      };

      const proxyReq = request(options, (proxyRes) => {
        const headers = { ...proxyRes.headers };
        delete headers['x-frame-options'];
        res.writeHead(proxyRes.statusCode || 200, headers);
        proxyRes.pipe(res);
      });

      proxyReq.on('error', (err) => {
        this.logManager.logError('Proxy request error:', err);
        res.writeHead(500);
        res.end('Proxy error');
      });

      req.pipe(proxyReq);
    });

    this.proxyServer.listen(proxyPort, () => {
      this.windowManager.sendKoboldOutput(
        `Proxy server started on port ${proxyPort}, forwarding to SillyTavern on port ${targetPort}`
      );
    });

    this.proxyServer.on('error', (err) => {
      this.logManager.logError('Proxy server error:', err);
    });
  }

  async startFrontend(args: string[]): Promise<void> {
    try {
      const config = {
        name: 'sillytavern',
        port: SILLYTAVERN.PORT,
        proxyPort: SILLYTAVERN.PROXY_PORT,
      };
      const { host: koboldHost, port: koboldPort } =
        this.parseKoboldConfig(args);

      await this.stopFrontend();

      this.windowManager.sendKoboldOutput(
        `Preparing SillyTavern to connect to KoboldCpp at ${koboldHost}:${koboldPort}...`
      );

      await this.ensureSillyTavernSettings();
      await this.setupSillyTavernConfig(koboldHost, koboldPort);

      this.windowManager.sendKoboldOutput(
        `Starting ${config.name} frontend on port ${config.port}...`
      );

      const sillyTavernArgs = [
        ...SillyTavernManager.SILLYTAVERN_BASE_ARGS,
        '--port',
        config.port.toString(),
      ];

      this.windowManager.sendKoboldOutput(
        'Final port check before starting SillyTavern...'
      );

      this.sillyTavernProcess = this.createNpxProcess(sillyTavernArgs);

      if (this.sillyTavernProcess.stdout) {
        this.sillyTavernProcess.stdout.on('data', (data: Buffer) => {
          this.windowManager.sendKoboldOutput(data.toString(), true);
        });
      }

      if (this.sillyTavernProcess.stderr) {
        this.sillyTavernProcess.stderr.on('data', (data: Buffer) => {
          this.windowManager.sendKoboldOutput(data.toString(), true);
        });
      }

      this.sillyTavernProcess.on(
        'exit',
        (code: number | null, signal: string | null) => {
          const message = signal
            ? `SillyTavern terminated with signal ${signal}`
            : `SillyTavern exited with code ${code}`;
          this.windowManager.sendKoboldOutput(message);
          this.sillyTavernProcess = null;
        }
      );

      this.sillyTavernProcess.on('error', (error) => {
        this.logManager.logError('SillyTavern process error:', error);
        this.windowManager.sendKoboldOutput(
          `SillyTavern error: ${error.message}`
        );

        this.sillyTavernProcess = null;
      });

      await this.waitForSillyTavernToStart(config.port);
      this.createProxyServer(config.port, config.proxyPort);
    } catch (error) {
      this.logManager.logError(
        `Failed to start SillyTavern: ${error instanceof Error ? error.message : String(error)}`,
        error as Error
      );
      throw error;
    }
  }

  async stopFrontend(): Promise<void> {
    this.windowManager.sendKoboldOutput('Stopping SillyTavern frontend...');

    try {
      this.windowManager.sendKoboldOutput(
        `Killing processes on port ${SILLYTAVERN.PORT}...`
      );
    } catch (error) {
      this.logManager.logError(
        'Error killing processes on port:',
        error as Error
      );
    }

    const promises: Promise<void>[] = [];

    if (this.sillyTavernProcess) {
      promises.push(
        terminateProcess(this.sillyTavernProcess, {
          logError: (message, error) =>
            this.logManager.logError(message, error),
        }).then(() => {
          this.sillyTavernProcess = null;
          this.windowManager.sendKoboldOutput('SillyTavern process terminated');
        })
      );
    }

    if (this.proxyServer) {
      promises.push(
        new Promise((resolve) => {
          this.proxyServer?.close(() => {
            this.proxyServer = null;
            this.windowManager.sendKoboldOutput('Proxy server closed');
            resolve();
          });
        })
      );
    }

    await Promise.all(promises);
    this.windowManager.sendKoboldOutput('SillyTavern frontend stopped');
  }

  async cleanup(): Promise<void> {
    if (this.sillyTavernProcess) {
      try {
        await this.stopFrontend();
      } catch (error) {
        this.logManager.logError(
          'Error during SillyTavernManager cleanup:',
          error as Error
        );
      }
    }
  }
}