import { Injectable } from '@angular/core';
import * as ort from 'onnxruntime-web';

// Use 'any' to bypass TypeScript checking for global libraries
declare var tflite: any;
declare var tf: any;

@Injectable({ providedIn: 'root' })
export class LocalInferenceService {
  private onnxSession: ort.InferenceSession | null = null;
  private tfliteModel: any | null = null;

  private readonly GRADE_COLORS: { [key: string]: string } = {
    'Empty_Bunch': '#6C757D', 'Underripe': '#F9A825', 'Abnormal': '#DC3545',
    'Ripe': '#00A651', 'Unripe': '#9E9D24', 'Overripe': '#5D4037'
  };

  async loadModel(modelPath: string) {
    if (modelPath.endsWith('.onnx')) {
      ort.env.wasm.wasmPaths = '/assets/wasm/';
      this.onnxSession = await ort.InferenceSession.create(modelPath, { executionProviders: ['wasm'] });
      this.tfliteModel = null;
      console.log('✅ ONNX Engine Ready');
    } else {
      // Accessing global variables from angular.json scripts
      if (typeof tflite !== 'undefined') {
        tflite.setWasmPath('/assets/tflite-wasm/');
        this.tfliteModel = await tflite.loadTFLiteModel(modelPath);
        this.onnxSession = null;
        console.log('✅ TFLite Engine Ready');
      } else {
        throw new Error('TFLite library not loaded. Check angular.json scripts.');
      }
    }
  }

  async runInference(input: Float32Array): Promise<Float32Array | null> {
    try {
      if (this.onnxSession) {
        const tensor = new ort.Tensor('float32', input, [1, 3, 640, 640]);
        const output = await this.onnxSession.run({ [this.onnxSession.inputNames[0]]: tensor });
        return output[this.onnxSession.outputNames[0]].data as Float32Array;
      } else if (this.tfliteModel) {
        // Use Global 'tf' for tensor conversion
        const inputTensor = tf.tensor(input, [1, 3, 640, 640]);
        // YOLOv8 TFLite often requires HWC shape [1, 640, 640, 3]
        const hwcTensor = tf.transpose(inputTensor, [0, 2, 3, 1]);
        
        const result = await this.tfliteModel.predict(hwcTensor);
        const data = await result.data();
        
        tf.dispose([inputTensor, hwcTensor, result]); 
        return data as Float32Array;
      }
    } catch (err) {
      console.error('Inference Error:', err);
    }
    return null;
  }

  parseDetections(rawData: any, threshold: number, imgWidth: number, imgHeight: number): any[] {
    if (!rawData) return [];
    const tensorData = rawData instanceof Float32Array ? rawData : new Float32Array(rawData);
    const detections = [];
    
    // Safety check for data length
    const numBoxes = Math.floor(tensorData.length / 6);

    for (let i = 0; i < numBoxes; i++) {
      const offset = i * 6;
      const confidence = tensorData[offset + 4];

      if (confidence >= threshold) {
        const classId = Math.round(tensorData[offset + 5]);
        const className = this.getClassName(classId);
        detections.push({
          box: [
            tensorData[offset + 0] * imgWidth,
            tensorData[offset + 1] * imgHeight,
            tensorData[offset + 2] * imgWidth,
            tensorData[offset + 3] * imgHeight
          ],
          confidence,
          class: className,
          color: this.GRADE_COLORS[className] || '#000000'
        });
      }
    }
    return detections;
  }

  private getClassName(id: number): string {
    const classes = ['Empty_Bunch', 'Underripe', 'Abnormal', 'Ripe', 'Unripe', 'Overripe'];
    return classes[id] || 'Unknown';
  }
}
