// SentencePiece-BPE tokenizer for SMaLL-100 (M2M-100 vocab), reading this
// repo's tokenizer.json directly (vocab + merges) — no native/FFI tokenizer
// binding needed, pure Dart. Verified against the Python `tokenizers` output.
//
// pubspec.yaml: no extra dependency (dart:convert + dart:io only).
import 'dart:convert';
import 'dart:io';

/// [encode] returns the subword ids only — the caller prepends the target
/// language token and appends `</s>` (id 2, per lang_tokens.json). [decode]
/// joins content tokens back to text, skipping special tokens and turning the
/// metaspace marker `▁` into spaces.
class BpeTokenizer {
  static const String _metaspace = '▁';

  final Map<String, int> _vocab;
  final List<String> _idToToken;
  final Map<String, int> _mergeRank; // "a b" -> rank (lower = higher priority)
  final Set<int> _specialIds;
  final int _unkId;

  BpeTokenizer._(this._vocab, this._idToToken, this._mergeRank, this._specialIds, this._unkId);

  // Fullwidth (U+FFxx) forms the SentencePiece nmt_nfkc normalizer folds to
  // ASCII (NFKC). CJK punctuation like `。`(U+3002) / `、`(U+3001) is NOT
  // folded — leave it as-is.
  static const Map<String, String> _punct = {
    '，': ',', '？': '?', '！': '!', '：': ':', '；': ';',
    '（': '(', '）': ')', '％': '%', '　': ' ',
  };

  static Future<BpeTokenizer> fromFile(String tokenizerJsonPath) async {
    final json = jsonDecode(await File(tokenizerJsonPath).readAsString()) as Map<String, dynamic>;
    final model = json['model'] as Map<String, dynamic>;

    final rawVocab = (model['vocab'] as Map).cast<String, dynamic>();
    final vocab = <String, int>{};
    var maxId = 0;
    rawVocab.forEach((k, v) {
      final id = (v as num).toInt();
      vocab[k] = id;
      if (id > maxId) maxId = id;
    });
    final idToToken = List<String>.filled(maxId + 1, '');
    vocab.forEach((k, v) {
      if (v >= 0 && v <= maxId) idToToken[v] = k;
    });

    final merges = model['merges'] as List<dynamic>;
    final mergeRank = <String, int>{};
    for (var i = 0; i < merges.length; i++) {
      final m = merges[i];
      final String key;
      if (m is String) {
        key = m; // "a b"
      } else {
        final pair = (m as List).cast<String>();
        key = '${pair[0]} ${pair[1]}';
      }
      mergeRank[key] = i;
    }

    // Special ids (</s>, <pad>, <unk>, <s>, and the 100 __xx__ language tokens).
    final specialIds = <int>{};
    for (final t in (json['added_tokens'] as List<dynamic>? ?? const [])) {
      final m = t as Map<String, dynamic>;
      if (m['special'] == true) specialIds.add((m['id'] as num).toInt());
    }
    final unkId = vocab['<unk>'] ?? 3;

    return BpeTokenizer._(vocab, idToToken, mergeRank, specialIds, unkId);
  }

  String _normalize(String text) {
    final sb = StringBuffer();
    for (final r in text.runes) {
      final ch = String.fromCharCode(r);
      sb.write(_punct[ch] ?? ch);
    }
    return sb.toString();
  }

  // Greedy BPE: repeatedly merge the adjacent pair with the lowest merge rank.
  List<String> _bpe(String word) {
    final symbols = word.runes.map(String.fromCharCode).toList();
    if (symbols.length < 2) return symbols;
    while (true) {
      var bestRank = 1 << 30;
      var bestI = -1;
      for (var i = 0; i < symbols.length - 1; i++) {
        final r = _mergeRank['${symbols[i]} ${symbols[i + 1]}'];
        if (r != null && r < bestRank) {
          bestRank = r;
          bestI = i;
        }
      }
      if (bestI < 0) break;
      symbols[bestI] = symbols[bestI] + symbols[bestI + 1];
      symbols.removeAt(bestI + 1);
    }
    return symbols;
  }

  /// Text -> subword ids (no language token, no `</s>`).
  List<int> encode(String text) {
    final ids = <int>[];
    var prevUnk = false;
    for (final piece in _normalize(text).split(RegExp(r'\s+'))) {
      if (piece.isEmpty) continue;
      for (final sym in _bpe(_metaspace + piece)) {
        final id = _vocab[sym];
        if (id == null) {
          if (!prevUnk) ids.add(_unkId); // fuse_unk
          prevUnk = true;
        } else {
          ids.add(id);
          prevUnk = false;
        }
      }
    }
    return ids;
  }

  /// Content ids -> text (skips special tokens, `▁` -> space).
  String decode(List<int> ids) {
    final sb = StringBuffer();
    for (final id in ids) {
      if (_specialIds.contains(id)) continue;
      if (id >= 0 && id < _idToToken.length) sb.write(_idToToken[id]);
    }
    return sb.toString().replaceAll(_metaspace, ' ').trim();
  }
}