modelgrid/ts/cli/model-handler.ts

/**
 * Model Handler
 *
 * CLI commands for model management.
 */

import { logger } from '../logger.ts';
import { theme } from '../colors.ts';
import { ContainerManager } from '../containers/container-manager.ts';
import { ModelRegistry } from '../models/registry.ts';
import { ModelLoader } from '../models/loader.ts';
import type { ITableColumn } from '../logger.ts';

/**
 * Handler for model-related CLI commands
 */
export class ModelHandler {
  private containerManager: ContainerManager;
  private modelRegistry: ModelRegistry;
  private modelLoader: ModelLoader;

  constructor(
    containerManager: ContainerManager,
    modelRegistry: ModelRegistry,
  ) {
    this.containerManager = containerManager;
    this.modelRegistry = modelRegistry;
    this.modelLoader = new ModelLoader(modelRegistry, containerManager);
  }

  /**
   * List all available models
   */
  public async list(): Promise<void> {
    logger.log('');
    logger.info('Models');
    logger.log('');

    // Get loaded models from containers
    const loadedModels = await this.containerManager.getAllAvailableModels();

    // Get greenlit models
    const greenlitModels = await this.modelRegistry.getAllGreenlitModels();

    if (loadedModels.size === 0 && greenlitModels.length === 0) {
      logger.logBox(
        'No Models',
        [
          'No models are loaded or greenlit.',
          '',
          theme.dim('Pull a model with:'),
          `  ${theme.command('modelgrid model pull <name>')}`,
        ],
        60,
        'warning',
      );
      return;
    }

    // Show loaded models
    if (loadedModels.size > 0) {
      logger.info(`Loaded Models (${loadedModels.size}):`);
      logger.log('');

      const rows = [];
      for (const [name, info] of loadedModels) {
        rows.push({
          name,
          container: info.container,
          size: info.size ? this.formatSize(info.size) : theme.dim('N/A'),
          format: info.format || theme.dim('N/A'),
          modified: info.modifiedAt
            ? new Date(info.modifiedAt).toLocaleDateString()
            : theme.dim('N/A'),
        });
      }

      const columns: ITableColumn[] = [
        { header: 'Name', key: 'name', align: 'left', color: theme.highlight },
        { header: 'Container', key: 'container', align: 'left' },
        { header: 'Size', key: 'size', align: 'right', color: theme.info },
        { header: 'Format', key: 'format', align: 'left' },
        { header: 'Modified', key: 'modified', align: 'left', color: theme.dim },
      ];

      logger.logTable(columns, rows);
      logger.log('');
    }

    // Show greenlit models (not yet loaded)
    const loadedNames = new Set(loadedModels.keys());
    const unloadedGreenlit = greenlitModels.filter((m) => !loadedNames.has(m.name));

    if (unloadedGreenlit.length > 0) {
      logger.info(`Available to Pull (${unloadedGreenlit.length}):`);
      logger.log('');

      const rows = unloadedGreenlit.map((m) => ({
        name: m.name,
        container: m.container,
        vram: `${m.minVram} GB`,
        tags: m.tags?.join(', ') || theme.dim('None'),
      }));

      const columns: ITableColumn[] = [
        { header: 'Name', key: 'name', align: 'left' },
        { header: 'Container', key: 'container', align: 'left' },
        { header: 'Min VRAM', key: 'vram', align: 'right', color: theme.info },
        { header: 'Tags', key: 'tags', align: 'left', color: theme.dim },
      ];

      logger.logTable(columns, rows);
      logger.log('');
    }
  }

  /**
   * Pull a model
   */
  public async pull(modelName: string): Promise<void> {
    if (!modelName) {
      logger.error('Model name is required');
      return;
    }

    logger.log('');
    logger.info(`Pulling model: ${modelName}`);
    logger.log('');

    const result = await this.modelLoader.loadModel(modelName);

    if (result.success) {
      if (result.alreadyLoaded) {
        logger.success(`Model "${modelName}" is already loaded`);
      } else {
        logger.success(`Model "${modelName}" pulled successfully`);
      }
      if (result.container) {
        logger.dim(`Container: ${result.container}`);
      }
    } else {
      logger.error(`Failed to pull model: ${result.error}`);
    }

    logger.log('');
  }

  /**
   * Remove a model
   */
  public async remove(modelName: string): Promise<void> {
    if (!modelName) {
      logger.error('Model name is required');
      return;
    }

    logger.info(`Removing model: ${modelName}`);

    const success = await this.modelLoader.unloadModel(modelName);

    if (success) {
      logger.success(`Model "${modelName}" removed`);
    } else {
      logger.error(`Failed to remove model "${modelName}"`);
    }
  }

  /**
   * Show model loading status and recommendations
   */
  public async status(): Promise<void> {
    logger.log('');
    await this.modelLoader.printStatus();
  }

  /**
   * Refresh greenlist cache
   */
  public async refresh(): Promise<void> {
    logger.info('Refreshing greenlist...');

    await this.modelRegistry.refreshGreenlist();

    logger.success('Greenlist refreshed');
  }

  /**
   * Format file size
   */
  private formatSize(bytes: number): string {
    const units = ['B', 'KB', 'MB', 'GB', 'TB'];
    let size = bytes;
    let unitIndex = 0;

    while (size >= 1024 && unitIndex < units.length - 1) {
      size /= 1024;
      unitIndex++;
    }

    return `${size.toFixed(1)} ${units[unitIndex]}`;
  }
}