src/services/hfStorage.ts

/**
 * src/services/hfStorage.ts
 *
 * HuggingFace storage using the official @huggingface/hub SDK.
 *
 * WHY @huggingface/hub instead of raw axios:
 *   The commit API (NDJSON) rejects binary files on repos with Xet storage enabled.
 *   The official SDK handles Xet automatically — same protocol the web UI uses.
 *   Works on both regular repos and Xet-enabled repos transparently.
 */

import { uploadFile, fileExists, createRepo, RepoDesignation } from '@huggingface/hub';
import axios from 'axios';
import { logger } from '../utils/logger';
import { ENV } from '../config/env';

// ─── Repo config ──────────────────────────────────────────────────────────────

// Using "model" type — avoids Xet enforcement that's mandatory on new datasets.
// The user creates the repo as a Model repo on HuggingFace.
const REPO: RepoDesignation = {
  type: 'model',
  name: ENV.HF_DATASET_REPO,   // e.g. "M-hv1/Titan-CDN"
};

const CREDENTIALS = { accessToken: ENV.HF_TOKEN };

const HF_BASE = 'https://huggingface.co';

// ─── Service ──────────────────────────────────────────────────────────────────

export class HFStorageService {

  // ── Startup check ─────────────────────────────────────────────────────────

  /**
   * Verifies the repo exists and the token has write access.
   * Creates the repo automatically if it doesn't exist.
   * Called once at bootstrap before any job is processed.
   */
  async ensureRepoExists(): Promise<void> {
    const repo = ENV.HF_DATASET_REPO;
    try {
      // fileExists() makes a lightweight HEAD request — cheap existence check
      await fileExists({ repo: REPO, credentials: CREDENTIALS, path: 'README.md' });
      logger.info(`HFStorage: model repo "${repo}" exists ✅`);
    } catch {
      // Repo might not exist yet — try to create it
      logger.info(`HFStorage: creating model repo "${repo}"...`);
      try {
        await createRepo({
          repo:        REPO,
          credentials: CREDENTIALS,
          private:     false,
        });
        logger.info(`HFStorage: model repo "${repo}" created ✅`);
      } catch (createErr) {
        // Repo might already exist but fileExists threw for another reason — log and continue
        logger.warn(`HFStorage: createRepo warning (repo may already exist): ${String(createErr)}`);
      }
    }
  }

  // ── Upload buffer ──────────────────────────────────────────────────────────

  /**
   * Uploads any Buffer (image, ZIP, binary) to the HF model repo.
   * Uses @huggingface/hub which handles Xet storage automatically.
   * Returns a permanent public URL.
   */
  async uploadBuffer(
    buffer: Buffer,
    filename: string,
    mimeType = 'application/octet-stream'
  ): Promise<string> {
    const repo = ENV.HF_DATASET_REPO;
    logger.info(
      `HFStorage: uploading "${filename}" ` +
      `(${(buffer.length / 1024).toFixed(1)} KB) → ${repo}`
    );

    // @huggingface/hub expects a Blob — Node 20 has Blob built-in
    const blob = new Blob([buffer], { type: mimeType });

    await uploadFile({
      repo:         REPO,
      credentials:  CREDENTIALS,
      file:         { path: filename, content: blob },
      commitTitle:  `Upload ${filename}`,
    });

    // Public download URL for model repos
    const url = `${HF_BASE}/${repo}/resolve/main/${filename}`;
    logger.info(`HFStorage: ✅ uploaded → ${url}`);
    return url;
  }

  // ── Upload icon from external URL ─────────────────────────────────────────

  /**
   * Fetches an icon from an external URL and uploads it to HF.
   * If the URL is already a HF URL — uses it directly (no re-upload).
   */
  async uploadIconFromUrl(iconUrl: string, trackingId: string): Promise<string> {
    // Already on HF? Use directly
    if (iconUrl.includes('huggingface.co')) {
      logger.info(`HFStorage: icon already on HF — using directly`);
      return iconUrl;
    }

    logger.info(`HFStorage: fetching icon from ${iconUrl}`);
    const res = await axios.get<ArrayBuffer>(iconUrl, {
      responseType: 'arraybuffer',
      timeout: 30_000,
    });

    const buf      = Buffer.from(res.data);
    const ext      = iconUrl.split('?')[0].split('.').pop()?.toLowerCase() || 'png';
    const filename = `icons/${trackingId}/icon.${ext}`;
    const mime     = String(res.headers['content-type'] || 'image/png');

    return this.uploadBuffer(buf, filename, mime);
  }

  // ── Upload icon from base64 ────────────────────────────────────────────────

  async uploadIconFromBase64(
    base64: string,
    mime: string,
    trackingId: string
  ): Promise<string> {
    const buf      = Buffer.from(base64, 'base64');
    const ext      = mime.split('/')[1]?.split(';')[0] || 'png';
    const filename = `icons/${trackingId}/icon.${ext}`;
    return this.uploadBuffer(buf, filename, mime);
  }
}

export const hfStorage = new HFStorageService();