hf-dataset.js

const { HfApi } = require('@huggingface/hub');

class HFDatasetManager {
  constructor() {
    this.hfToken = process.env.HF_TOKEN;
    this.datasetId = process.env.HF_DATASET_ID || "Detomo/houzou-devices";
    this.fileName = "devices.json";
    this.isEnabled = !!this.hfToken;
    
    if (this.isEnabled) {
      this.hfApi = new HfApi({ accessToken: this.hfToken });
    }
    
    console.log(`🤗 HF Dataset Manager initialized`);
    console.log(`   Dataset: ${this.datasetId}`);
    console.log(`   Enabled: ${this.isEnabled}`);
  }

  async loadDevices() {
    if (!this.isEnabled) {
      console.log('⚠️ HF Dataset disabled - no token provided');
      return new Map();
    }

    try {
      console.log('📥 Loading devices from HF dataset...');
      
      const fileUrl = `https://huggingface.co/datasets/${this.datasetId}/resolve/main/${this.fileName}`;
      const response = await fetch(fileUrl, {
        headers: {
          'Authorization': `Bearer ${this.hfToken}`
        }
      });

      if (response.status === 404) {
        console.log('📁 No devices file found in dataset, creating new one');
        await this.saveDevices(new Map());
        return new Map();
      }

      if (!response.ok) {
        throw new Error(`HTTP ${response.status}: ${response.statusText}`);
      }

      const data = await response.json();
      const deviceMap = new Map(data);
      
      console.log(`📥 Loaded ${deviceMap.size} devices from HF dataset`);
      return deviceMap;
      
    } catch (error) {
      console.error('❌ Error loading devices from HF dataset:', error);
      console.log('⚠️ Falling back to empty device list');
      return new Map();
    }
  }

  async saveDevices(deviceMap) {
    if (!this.isEnabled) {
      console.log('⚠️ HF Dataset disabled - cannot save devices');
      return false;
    }

    try {
      console.log(`💾 Saving ${deviceMap.size} devices to HF dataset...`);
      
      const devicesArray = Array.from(deviceMap.entries());
      const jsonData = JSON.stringify(devicesArray, null, 2);
      
      // Create a buffer from the JSON data
      const fileContent = Buffer.from(jsonData, 'utf8');
      
      // Upload using HfApi.uploadFile
      await this.hfApi.uploadFile({
        repo: this.datasetId,
        file: {
          path: this.fileName,
          content: fileContent
        },
        repoType: 'dataset',
        commitMessage: `Update devices data (${deviceMap.size} devices)`
      });
      
      console.log(`✅ Successfully saved ${deviceMap.size} devices to HF dataset`);
      return true;
      
    } catch (error) {
      console.error('❌ Error saving devices to HF dataset:', error);
      return false;
    }
  }

  async createDatasetIfNotExists() {
    if (!this.isEnabled) {
      console.log('⚠️ HF Dataset disabled - cannot create dataset');
      return false;
    }

    try {
      console.log('🔍 Checking if dataset exists...');
      
      // Try to check if dataset exists
      try {
        await this.hfApi.repoInfo({ repo: this.datasetId, repoType: 'dataset' });
        console.log('✅ Dataset already exists');
        return true;
      } catch (error) {
        if (error.message.includes('404')) {
          console.log('📁 Dataset not found, creating new one...');
          
          // Create new dataset
          await this.hfApi.createRepo({
            repo: this.datasetId,
            repoType: 'dataset',
            description: 'Device tokens for Houzou Medical App notifications'
          });
          
          console.log('✅ Dataset created successfully');
          
          // Create initial empty devices file
          await this.saveDevices(new Map());
          return true;
        }
        throw error;
      }
      
    } catch (error) {
      console.error('❌ Error initializing dataset:', error);
      console.log('⚠️ Dataset initialization failed, will use fallback');
      return false;
    }
  }

  isReady() {
    return this.isEnabled;
  }

  getStatus() {
    return {
      enabled: this.isEnabled,
      datasetId: this.datasetId,
      hasToken: !!this.hfToken
    };
  }
}

module.exports = HFDatasetManager;