/**
 * Process Real Photos for VIMS
 * Creates training dataset from uploaded photos
 */

const fs = require('fs').promises;
const path = require('path');
const sharp = require('sharp');

class RealPhotoProcessor {
  constructor() {
    this.inputDir = './data/real-photos';
    this.outputDir = './data/training';
  }

  async processPhotoSet(photoSetId) {
    const photoDir = path.join(this.inputDir, photoSetId);
    const annotationFile = path.join(photoDir, 'annotations.json');
    
    console.log(`Processing ${photoSetId}...`);
    
    // Read annotations
    const annotations = JSON.parse(await fs.readFile(annotationFile, 'utf-8'));
    
    // Process each image
    for (const [filename, data] of Object.entries(annotations.annotations)) {
      const imagePath = path.join(photoDir, filename);
      
      try {
        await this.processImage(imagePath, data, photoSetId);
      } catch (error) {
        console.error(`Failed to process ${filename}:`, error.message);
      }
    }
    
    console.log(`✓ ${photoSetId} processed`);
  }

  async processImage(imagePath, annotation, photoSetId) {
    // Read image
    const image = sharp(imagePath);
    const metadata = await image.metadata();
    
    // Resize to training size
    const resized = await image
      .resize(640, 640, { fit: 'contain', background: { r: 114, g: 114, b: 114 } })
      .jpeg({ quality: 95 })
      .toBuffer();
    
    // Save to training directory
    const outputName = `${photoSetId}_${annotation.angle}.jpg`;
    const outputPath = path.join(this.outputDir, 'atm', 'images', 'train', outputName);
    
    await fs.mkdir(path.dirname(outputPath), { recursive: true });
    await fs.writeFile(outputPath, resized);
    
    // Create YOLO label file
    const labelPath = outputPath.replace('/images/', '/labels/').replace('.jpg', '.txt');
    const labels = this.convertToYOLO(annotation.components, metadata.width, metadata.height);
    
    await fs.mkdir(path.dirname(labelPath), { recursive: true });
    await fs.writeFile(labelPath, labels);
    
    console.log(`  ✓ ${outputName}`);
  }

  convertToYOLO(components, imgWidth, imgHeight) {
    const classMap = {
      'card_reader': 0,
      'pin_pad': 1,
      'display': 2,
      'cash_dispenser': 3,
      'nfc_reader': 4,
      'receipt_printer': 5,
      'camera': 6,
      'speaker': 7,
      'button': 8
    };
    
    return components.map(comp => {
      const classId = classMap[comp.type] || 0;
      const { x, y, w, h } = comp.bbox;
      
      // YOLO format: class x_center y_center width height (all normalized)
      return `${classId} ${x + w/2} ${y + h/2} ${w} ${h}`;
    }).join('\n');
  }

  async createDatasetYaml() {
    const yaml = {
      path: path.resolve(this.outputDir, 'atm'),
      train: 'images/train',
      val: 'images/val',
      test: 'images/test',
      nc: 9,
      names: [
        'card_reader',
        'pin_pad',
        'display',
        'cash_dispenser',
        'nfc_reader',
        'receipt_printer',
        'camera',
        'speaker',
        'button'
      ]
    };
    
    const yamlPath = path.join(this.outputDir, 'atm', 'dataset.yaml');
    await fs.writeFile(yamlPath, JSON.stringify(yaml, null, 2));
    
    console.log('✓ dataset.yaml created');
  }

  async run() {
    console.log('🚀 Processing real photos for VIMS training\n');
    
    // Find all photo sets
    const entries = await fs.readdir(this.inputDir, { withFileTypes: true });
    const photoSets = entries.filter(e => e.isDirectory()).map(e => e.name);
    
    console.log(`Found ${photoSets.length} photo set(s): ${photoSets.join(', ')}\n`);
    
    for (const photoSet of photoSets) {
      await this.processPhotoSet(photoSet);
    }
    
    await this.createDatasetYaml();
    
    console.log('\n✅ All photos processed!');
    console.log('Next step: Train model with:');
    console.log('  python src/training/train-yolo.py atm --epochs 50');
  }
}

// Run if called directly
if (require.main === module) {
  const processor = new RealPhotoProcessor();
  processor.run().catch(console.error);
}

module.exports = { RealPhotoProcessor };