amical/apps/desktop/src/modules/models/model-manager.ts

import { EventEmitter } from "events";
import * as fs from "fs";
import * as path from "path";
import * as crypto from "crypto";
import { app } from "electron";
import {
  Model,
  DownloadProgress,
  ModelManagerState,
  AVAILABLE_MODELS,
} from "../../constants/models";
import { DownloadedModel } from "../../db/schema";
import {
  getDownloadedModelsRecord,
  createDownloadedModel,
  deleteDownloadedModel,
  validateDownloadedModels,
  validateModelFile,
  getValidDownloadedModels,
} from "../../db/downloaded-models";
import { logger } from "../../main/logger";

interface ModelManagerEvents {
  "download-progress": (modelId: string, progress: DownloadProgress) => void;
  "download-complete": (
    modelId: string,
    downloadedModel: DownloadedModel,
  ) => void;
  "download-error": (modelId: string, error: Error) => void;
  "download-cancelled": (modelId: string) => void;
  "model-deleted": (modelId: string) => void;
}

declare interface ModelManagerService {
  on<U extends keyof ModelManagerEvents>(
    event: U,
    listener: ModelManagerEvents[U],
  ): this;
  emit<U extends keyof ModelManagerEvents>(
    event: U,
    ...args: Parameters<ModelManagerEvents[U]>
  ): boolean;
}

class ModelManagerService extends EventEmitter {
  private state: ModelManagerState;
  private modelsDirectory: string;

  constructor() {
    super();
    this.state = {
      activeDownloads: new Map(),
    };

    // Create models directory in app data
    this.modelsDirectory = path.join(app.getPath("userData"), "models");
    this.ensureModelsDirectory();
  }

  // Initialize and validate models on startup
  async initialize(): Promise<void> {
    try {
      const validation = await validateDownloadedModels();

      if (validation.cleaned > 0) {
        logger.main.info("Cleaned up missing model records", {
          cleaned: validation.cleaned,
          valid: validation.valid.length,
          missing: validation.missing.map((m) => ({
            id: m.id,
            path: m.localPath,
          })),
        });
      }

      logger.main.info("Model manager initialized", {
        validModels: validation.valid.length,
        cleanedRecords: validation.cleaned,
      });
    } catch (error) {
      logger.main.error("Error initializing model manager", { error });
    }
  }

  private ensureModelsDirectory(): void {
    if (!fs.existsSync(this.modelsDirectory)) {
      fs.mkdirSync(this.modelsDirectory, { recursive: true });
      logger.main.info("Created models directory", {
        path: this.modelsDirectory,
      });
    }
  }

  // Get all available models from manifest
  getAvailableModels(): Model[] {
    return AVAILABLE_MODELS;
  }

  // Get downloaded models from database
  async getDownloadedModels(): Promise<Record<string, DownloadedModel>> {
    return await getDownloadedModelsRecord();
  }

  // Get only valid downloaded models (files that exist on disk)
  async getValidDownloadedModels(): Promise<Record<string, DownloadedModel>> {
    const validModels = await getValidDownloadedModels();
    const record: Record<string, DownloadedModel> = {};

    for (const model of validModels) {
      record[model.id] = model;
    }

    return record;
  }

  // Check if a model is downloaded and file exists
  async isModelDownloaded(modelId: string): Promise<boolean> {
    return await validateModelFile(modelId);
  }

  // Get download progress for a model
  getDownloadProgress(modelId: string): DownloadProgress | null {
    return this.state.activeDownloads.get(modelId) || null;
  }

  // Get all active downloads
  getActiveDownloads(): DownloadProgress[] {
    return Array.from(this.state.activeDownloads.values());
  }

  // Download a model
  async downloadModel(modelId: string): Promise<void> {
    const model = AVAILABLE_MODELS.find((m) => m.id === modelId);
    if (!model) {
      throw new Error(`Model not found: ${modelId}`);
    }

    if (await this.isModelDownloaded(modelId)) {
      throw new Error(`Model already downloaded: ${modelId}`);
    }

    if (this.state.activeDownloads.has(modelId)) {
      throw new Error(`Download already in progress: ${modelId}`);
    }

    const abortController = new AbortController();
    const downloadPath = path.join(this.modelsDirectory, model.filename);

    const progress: DownloadProgress = {
      modelId,
      progress: 0,
      status: "downloading",
      bytesDownloaded: 0,
      totalBytes: model.size,
      abortController,
    };

    this.state.activeDownloads.set(modelId, progress);
    this.emit("download-progress", modelId, progress);

    try {
      logger.main.info("Starting model download", {
        modelId,
        size: model.sizeFormatted,
        url: model.downloadUrl,
      });

      const response = await fetch(model.downloadUrl, {
        signal: abortController.signal,
      });

      if (!response.ok) {
        throw new Error(
          `Failed to download: ${response.status} ${response.statusText}`,
        );
      }

      const totalBytes =
        parseInt(response.headers.get("content-length") || "0") || model.size;
      progress.totalBytes = totalBytes;

      const fileStream = fs.createWriteStream(downloadPath);
      let bytesDownloaded = 0;
      let lastProgressEmit = 0;

      const reader = response.body?.getReader();
      if (!reader) {
        throw new Error("Failed to get response reader");
      }

      while (true) {
        const { done, value } = await reader.read();

        if (done) break;

        if (abortController.signal.aborted) {
          fileStream.close();
          fs.unlinkSync(downloadPath);
          throw new Error("Download cancelled");
        }

        fileStream.write(value);
        bytesDownloaded += value.length;

        progress.bytesDownloaded = bytesDownloaded;
        progress.progress = Math.round((bytesDownloaded / totalBytes) * 100);

        // Emit progress every 1% or 1MB to avoid too many events
        const progressPercent = progress.progress;
        if (
          progressPercent - lastProgressEmit >= 1 ||
          bytesDownloaded - (lastProgressEmit * totalBytes) / 100 >= 1024 * 1024
        ) {
          this.emit("download-progress", modelId, { ...progress });
          lastProgressEmit = progressPercent;
        }
      }

      fileStream.end();

      // Get actual file size (no validation against expected size)
      const stats = fs.statSync(downloadPath);
      logger.main.info("Download completed", {
        modelId,
        expectedSize: totalBytes,
        actualSize: stats.size,
        sizeDifference: Math.abs(stats.size - totalBytes),
      });

      // Verify checksum if provided
      if (model.checksum) {
        const fileChecksum = await this.calculateFileChecksum(downloadPath);
        if (fileChecksum !== model.checksum) {
          fs.unlinkSync(downloadPath);
          throw new Error(
            `Checksum mismatch. Expected: ${model.checksum}, Got: ${fileChecksum}`,
          );
        }
      }

      // Create downloaded model record in database
      const downloadedModel = await createDownloadedModel({
        id: model.id,
        name: model.name,
        type: model.type,
        localPath: downloadPath,
        downloadedAt: new Date(),
        size: stats.size,
        checksum: model.checksum,
      });

      // Clean up active download
      this.state.activeDownloads.delete(modelId);

      logger.main.info("Model download completed", {
        modelId,
        path: downloadPath,
        size: stats.size,
      });

      this.emit("download-complete", modelId, downloadedModel);
    } catch (error) {
      // Clean up on error
      this.state.activeDownloads.delete(modelId);

      if (fs.existsSync(downloadPath)) {
        fs.unlinkSync(downloadPath);
      }

      const err = error instanceof Error ? error : new Error(String(error));

      if (abortController.signal.aborted) {
        logger.main.info("Model download cancelled", { modelId });
        this.emit("download-cancelled", modelId);
      } else {
        logger.main.error("Model download failed", {
          modelId,
          error: err.message,
        });
        this.emit("download-error", modelId, err);
      }

      throw err;
    }
  }

  // Cancel a model download
  cancelDownload(modelId: string): void {
    const download = this.state.activeDownloads.get(modelId);
    if (!download) {
      throw new Error(`No active download found for model: ${modelId}`);
    }

    download.status = "cancelling";
    download.abortController?.abort();

    // Immediately remove from active downloads to prevent restart issues
    this.state.activeDownloads.delete(modelId);

    logger.main.info("Cancelled model download", { modelId });
    this.emit("download-cancelled", modelId);
  }

  // Delete a downloaded model
  async deleteModel(modelId: string): Promise<void> {
    const downloadedModels = await this.getDownloadedModels();
    const downloadedModel = downloadedModels[modelId];

    if (!downloadedModel) {
      throw new Error(`Model not found: ${modelId}`);
    }

    // Delete file
    if (fs.existsSync(downloadedModel.localPath)) {
      fs.unlinkSync(downloadedModel.localPath);
      logger.main.info("Deleted model file", {
        modelId,
        path: downloadedModel.localPath,
      });
    }

    // Remove from database
    await deleteDownloadedModel(modelId);

    this.emit("model-deleted", modelId);
  }

  // Calculate file checksum (SHA-1)
  private async calculateFileChecksum(filePath: string): Promise<string> {
    return new Promise((resolve, reject) => {
      const hash = crypto.createHash("sha1");
      const stream = fs.createReadStream(filePath);

      stream.on("data", (data) => hash.update(data));
      stream.on("end", () => resolve(hash.digest("hex")));
      stream.on("error", reject);
    });
  }

  // Get models directory path
  getModelsDirectory(): string {
    return this.modelsDirectory;
  }

  // Validate and clean up stale model records (can be called periodically)
  async validateAndCleanup(): Promise<{ cleaned: number; valid: number }> {
    try {
      const validation = await validateDownloadedModels();

      if (validation.cleaned > 0) {
        logger.main.info("Periodic cleanup completed", {
          cleaned: validation.cleaned,
          valid: validation.valid.length,
        });
      }

      return {
        cleaned: validation.cleaned,
        valid: validation.valid.length,
      };
    } catch (error) {
      logger.main.error("Error during model validation cleanup", { error });
      return { cleaned: 0, valid: 0 };
    }
  }

  // Cleanup - cancel all active downloads
  cleanup(): void {
    logger.main.info("Cleaning up model downloads", {
      activeDownloads: this.state.activeDownloads.size,
    });

    for (const [modelId] of this.state.activeDownloads) {
      try {
        this.cancelDownload(modelId);
      } catch (error) {
        logger.main.warn("Error cancelling download during cleanup", {
          modelId,
          error: error instanceof Error ? error.message : String(error),
        });
      }
    }
  }
}

export { ModelManagerService };