// Reference SMaLL-100 translator for Flutter (onnxruntime Dart package).
// This is the actual implementation shipped in a production S2S translator
// app (Android + iOS) — fully runnable and verified, not pseudo-code.
//
// pubspec.yaml:
//   dependencies:
//     onnxruntime: ^1.4.1
//
// Ship onnx/encoder_model.onnx, onnx/decoder_model_merged.onnx, tokenizer.json,
// lang_tokens.json as assets (or download them into app storage — sherpa/
// onnxruntime need real filesystem paths, not asset-bundle URIs). Implements
// the algorithm in ../USAGE.md: encode + prepend target-lang token -> encoder
// -> greedy merged-decoder loop with KV cache -> decode. Needs bpe_tokenizer.dart
// (same directory) to read tokenizer.json — Dart has no off-the-shelf
// HuggingFace-tokenizers binding, so this repo ships a small pure-Dart one.
import 'dart:convert';
import 'dart:io';
import 'dart:typed_data';

import 'package:onnxruntime/onnxruntime.dart';

import 'bpe_tokenizer.dart';

class Small100Translator {
  static const int _nLayers = 3;
  static const int _nHeads = 16;
  static const int _headDim = 64;
  static const int _decoderStart = 2;
  static const int _eos = 2;
  static const int _pad = 1;
  static const int _maxLen = 64;

  final BpeTokenizer _tok;
  final Map<String, int> _langToId;
  final OrtSession _encoder;
  final OrtSession _decoder;

  Small100Translator._(this._tok, this._langToId, this._encoder, this._decoder);

  static Future<Small100Translator> create({
    required String tokenizerJsonPath,
    required String langTokensPath,
    required String encoderPath,
    required String decoderPath,
  }) async {
    OrtEnv.instance.init();
    final tok = await BpeTokenizer.fromFile(tokenizerJsonPath);
    final lt = jsonDecode(await File(langTokensPath).readAsString()) as Map<String, dynamic>;
    final langToId = (lt['lang_to_id'] as Map).map((k, v) => MapEntry(k as String, (v as num).toInt()));
    final opts = OrtSessionOptions()..setIntraOpNumThreads(2);
    // On-device acceleration: unsupported ops silently fall back to CPU.
    // Some emulators (e.g. those that run on qemu but report a real device's
    // Build.FINGERPRINT) crash inside libonnxruntime.so with NNAPI enabled —
    // real devices are unaffected. If you hit that, detect the emulator and
    // skip appendNnapiProvider rather than disabling acceleration everywhere.
    if (Platform.isAndroid) {
      opts.appendNnapiProvider(NnapiFlags.useNone);
    } else if (Platform.isIOS) {
      opts.appendCoreMLProvider(CoreMLFlags.useNone);
    }
    final encoder = OrtSession.fromFile(File(encoderPath), opts);
    final decoder = OrtSession.fromFile(File(decoderPath), opts);
    return Small100Translator._(tok, langToId, encoder, decoder);
  }

  bool supportsLang(String code) => _langToId.containsKey(code);

  /// Translate [text] into [tgtLang] (e.g. 'en'). Returns '' for empty input or
  /// an unknown target language.
  Future<String> translate(String text, String tgtLang) async {
    if (text.trim().isEmpty) return '';
    final langId = _langToId[tgtLang];
    if (langId == null) return '';

    final runOpts = OrtRunOptions();
    // [tgt lang] + subwords + </s>  — SMaLL-100 puts the target-language token
    // on the *source* side, not via forced_bos_token_id.
    final ids = <int>[langId, ..._tok.encode(text), _eos];
    final seqLen = ids.length;

    final inputIds = OrtValueTensor.createTensorWithDataList(Int64List.fromList(ids), [1, seqLen]);
    final attnMask = OrtValueTensor.createTensorWithDataList(
        Int64List.fromList(List<int>.filled(seqLen, 1)), [1, seqLen]);

    final encOut = _encoder.run(runOpts, {'input_ids': inputIds, 'attention_mask': attnMask});
    final encoderHidden = encOut[0]!; // frozen across steps

    final past = <String, OrtValue>{};
    for (var i = 0; i < _nLayers; i++) {
      for (final kind in ['decoder', 'encoder']) {
        for (final kv in ['key', 'value']) {
          past['past_key_values.$i.$kind.$kv'] =
              OrtValueTensor.createTensorWithDataList(Float32List(0), [1, _nHeads, 0, _headDim]);
        }
      }
    }

    final outIds = <int>[];
    var cur = _decoderStart;
    var useCache = false;

    try {
      for (var step = 0; step < _maxLen; step++) {
        final stepInputIds = OrtValueTensor.createTensorWithDataList(Int64List.fromList([cur]), [1, 1]);
        final useCacheT = OrtValueTensor.createTensorWithDataList([useCache], [1]);

        final feeds = <String, OrtValue>{
          'encoder_attention_mask': attnMask,
          'input_ids': stepInputIds,
          'encoder_hidden_states': encoderHidden,
          'use_cache_branch': useCacheT,
          ...past,
        };

        final outputs = _decoder.run(runOpts, feeds);
        final outByName = <String, OrtValue?>{};
        for (var i = 0; i < _decoder.outputNames.length; i++) {
          outByName[_decoder.outputNames[i]] = outputs[i];
        }

        final next = _argmaxLastLogits(outByName['logits']!);
        stepInputIds.release();
        useCacheT.release();
        if (next == _eos) {
          _releaseOutputs(outputs, keep: {});
          break;
        }
        outIds.add(next);
        // Greedy decode can degenerate into a repeated token/short cycle on
        // out-of-distribution input; bail out instead of running to _maxLen.
        if (_degenerate(outIds)) {
          _releaseOutputs(outputs, keep: {});
          break;
        }

        final keep = <OrtValue>{};
        for (var i = 0; i < _nLayers; i++) {
          for (final kv in ['key', 'value']) {
            final decKey = 'past_key_values.$i.decoder.$kv';
            final newDec = outByName['present.$i.decoder.$kv']!;
            past[decKey]!.release();
            past[decKey] = newDec;
            keep.add(newDec);
            if (!useCache) {
              // Encoder KV is only produced on step 1 (use_cache_branch=false),
              // then frozen and reused for every later step.
              final encKey = 'past_key_values.$i.encoder.$kv';
              final newEnc = outByName['present.$i.encoder.$kv']!;
              past[encKey]!.release();
              past[encKey] = newEnc;
              keep.add(newEnc);
            }
          }
        }
        _releaseOutputs(outputs, keep: keep);

        cur = next;
        useCache = true;
      }
    } finally {
      for (final t in past.values) {
        t.release();
      }
      encoderHidden.release();
      inputIds.release();
      attnMask.release();
      runOpts.release();
    }

    return _tok.decode(outIds);
  }

  int _argmaxLastLogits(OrtValue logits) {
    final v = logits.value as List;
    final row = (v[0] as List).last as List; // [1, seq, vocab] -> last position
    var bestIdx = 0;
    var bestVal = double.negativeInfinity;
    for (var i = 0; i < row.length; i++) {
      if (i == _pad) continue; // never emit <pad>
      final x = (row[i] as num).toDouble();
      if (x > bestVal) {
        bestVal = x;
        bestIdx = i;
      }
    }
    return bestIdx;
  }

  // Bail out of greedy decoding if it degenerates into a repeated token or a
  // short repeating cycle (period 2-4) — cheap safety net beyond ../USAGE.md's
  // base algorithm, worth adopting on any platform.
  static bool _degenerate(List<int> ids) {
    final n = ids.length;
    if (n >= 6) {
      var same = true;
      for (var i = n - 6; i < n - 1; i++) {
        if (ids[i] != ids[i + 1]) {
          same = false;
          break;
        }
      }
      if (same) return true;
    }
    for (var p = 2; p <= 4; p++) {
      if (n >= p * 3) {
        var cyc = true;
        for (var i = 0; i < p * 2; i++) {
          if (ids[n - 1 - i] != ids[n - 1 - i - p]) {
            cyc = false;
            break;
          }
        }
        if (cyc) return true;
      }
    }
    return false;
  }

  void _releaseOutputs(List<OrtValue?> outputs, {required Set<OrtValue> keep}) {
    for (final o in outputs) {
      if (o != null && !keep.contains(o)) o.release();
    }
  }

  void dispose() {
    _encoder.release();
    _decoder.release();
  }
}

// Usage (e.g. from a Flutter widget's initState, after copying the asset
// files to real paths — see ModelManager-style staging in the app examples):
//
//   final t = await Small100Translator.create(
//     tokenizerJsonPath: '$modelsRoot/tokenizer.json',
//     langTokensPath: '$modelsRoot/lang_tokens.json',
//     encoderPath: '$modelsRoot/onnx/encoder_model.onnx',
//     decoderPath: '$modelsRoot/onnx/decoder_model_merged.onnx',
//   );
//   print(await t.translate('你好，请问最近的地铁站怎么走？', 'en'));
//   t.dispose();