import 'dart:typed_data';

import 'package:tflite_flutter/tflite_flutter.dart';

enum YoloTensorType { float32, uint8, int8, float16, unsupported }

class YoloRuntime {
  final Interpreter _interpreter;
  final List<int> inputShape;
  final List<int> outputShape;
  final YoloTensorType inputType;
  final YoloTensorType outputType;
  final int outputByteLength;

  YoloRuntime._({
    required Interpreter interpreter,
    required this.inputShape,
    required this.outputShape,
    required this.inputType,
    required this.outputType,
    required this.outputByteLength,
  }) : _interpreter = interpreter;

  static Future<YoloRuntime> load(String assetPath) async {
    final options = InterpreterOptions()..threads = 2;
    final interpreter =
        await Interpreter.fromAsset(assetPath, options: options);
    final input = interpreter.getInputTensor(0);
    final output = interpreter.getOutputTensor(0);
    return YoloRuntime._(
      interpreter: interpreter,
      inputShape: input.shape,
      outputShape: output.shape,
      inputType: _mapTensorType(input.type),
      outputType: _mapTensorType(output.type),
      outputByteLength: output.numBytes(),
    );
  }

  void run(Uint8List input, Uint8List output) {
    _interpreter.run(input, output);
  }

  void close() {
    _interpreter.close();
  }
}

YoloTensorType _mapTensorType(TensorType type) {
  return switch (type) {
    TensorType.float32 => YoloTensorType.float32,
    TensorType.uint8 => YoloTensorType.uint8,
    TensorType.int8 => YoloTensorType.int8,
    TensorType.float16 => YoloTensorType.float16,
    _ => YoloTensorType.unsupported,
  };
}
