import 'dart:io';
import 'dart:typed_data';
import 'package:flutter/services.dart' show rootBundle;
import 'package:path/path.dart' as p;
import 'package:path_provider/path_provider.dart';
import 'package:sherpa_onnx/sherpa_onnx.dart' as sherpa;

/// Wraps sherpa-onnx's GTCRN model for speech enhancement (noise reduction).
///
/// Uses the GTCRN model (ICASSP 2024, ~48K params, 523KB) which provides
/// lightweight ML-based denoising without over-processing — ideal for
/// preprocessing audio before sending to backend STT (faster-whisper).
///
/// Usage:
///   final enhancer = SpeechEnhancer();
///   await enhancer.init();
///   final denoised = enhancer.enhance(pcm16Bytes);
///   enhancer.dispose();
class SpeechEnhancer {
  static const String _modelAsset = 'assets/gtcrn_simple.onnx';

  sherpa.OfflineSpeechDenoiser? _denoiser;
  bool _initialized = false;
  String? _modelPath;

  /// Initialize the GTCRN denoiser. Must be called before [enhance].
  /// Copies the model from Flutter assets to the filesystem on first run.
  Future<void> init() async {
    if (_initialized) return;

    _modelPath = await _copyAssetToFile(_modelAsset);

    final config = sherpa.OfflineSpeechDenoiserConfig(
      model: sherpa.OfflineSpeechDenoiserModelConfig(
        gtcrn: sherpa.OfflineSpeechDenoiserGtcrnModelConfig(
          model: _modelPath!,
        ),
        numThreads: 1,
        debug: false,
        provider: 'cpu',
      ),
    );

    _denoiser = sherpa.OfflineSpeechDenoiser(config);
    _initialized = true;
  }

  /// Denoise PCM 16-bit signed little-endian mono audio at 16 kHz.
  ///
  /// Returns denoised PCM bytes in the same format.
  /// If the enhancer is not initialized, returns the input unchanged.
  Uint8List enhance(Uint8List pcm16Bytes) {
    if (_denoiser == null || pcm16Bytes.length < 640) return pcm16Bytes;

    // Convert Int16 PCM → Float32 normalized [-1.0, 1.0]
    final sampleCount = pcm16Bytes.length ~/ 2;
    final float32 = Float32List(sampleCount);
    final byteData = ByteData.sublistView(pcm16Bytes);
    for (int i = 0; i < sampleCount; i++) {
      float32[i] = byteData.getInt16(i * 2, Endian.little) / 32768.0;
    }

    // Run GTCRN denoiser
    final result = _denoiser!.run(samples: float32, sampleRate: 16000);

    // Convert Float32 → Int16 PCM bytes
    final output = Uint8List(result.samples.length * 2);
    final outView = ByteData.sublistView(output);
    for (int i = 0; i < result.samples.length; i++) {
      final val = (result.samples[i] * 32768.0).round().clamp(-32768, 32767);
      outView.setInt16(i * 2, val, Endian.little);
    }
    return output;
  }

  /// Release native resources. Must be called when done.
  void dispose() {
    _denoiser?.free();
    _denoiser = null;
    _initialized = false;
  }

  /// Copy a Flutter asset to the app support directory (cached).
  static Future<String> _copyAssetToFile(String assetPath) async {
    final dir = await getApplicationSupportDirectory();
    final target = p.join(dir.path, p.basename(assetPath));
    final file = File(target);

    if (!await file.exists()) {
      final data = await rootBundle.load(assetPath);
      await file.writeAsBytes(
        data.buffer.asUint8List(data.offsetInBytes, data.lengthInBytes),
      );
    }
    return target;
  }
}
