/*!
 *  Copyright (c) 2024 by Contributors
 * \file xgrammar/testing.cc
 */
#include "testing.h"

#include <xgrammar/xgrammar.h>

#include <cstdint>
#include <sstream>
#include <string>
#include <vector>

#include "grammar_impl.h"
#include "grammar_parser.h"
#include "support/encoding.h"

namespace xgrammar {

std::string PrintTokenByIds(
    const std::vector<int32_t>& token_ids, const TokenizerInfo& tokenizer_info, int max_print_num
) {
  std::stringstream ss;
  const auto& sorted_decoded_vocab = tokenizer_info.GetDecodedVocab();
  ss << "[";
  int print_num = std::min(static_cast<int>(token_ids.size()), max_print_num);
  for (int i = 0; i < print_num; ++i) {
    ss << "#" << token_ids[i] << " <" << EscapeString(sorted_decoded_vocab[token_ids[i]]) << ">";
    if (i < print_num - 1) {
      ss << ", ";
    }
  }
  if (static_cast<int>(token_ids.size()) > max_print_num) {
    ss << ", ...";
  }
  ss << "]";
  return ss.str();
}

Grammar _EBNFToGrammarNoNormalization(
    const std::string& ebnf_string, const std::string& root_rule_name
) {
  return ParseEBNF(ebnf_string, root_rule_name);
}

std::string _PrintGrammarFSMs(const Grammar& grammar) {
  std::string result;
  for (int i = 0; i < grammar->NumRules(); i++) {
    result += "Rule " + std::to_string(i) + ": " + grammar->GetRule(i).name + ", FSM: ";
    if (grammar->per_rule_fsms[i].has_value()) {
      result += grammar->per_rule_fsms[i]->ToString();
    } else {
      result += "None";
    }
    result += "\n";
  }
  return result;
}

namespace details {

void DFS(
    int32_t curr,
    int32_t parent_pos,
    const int64_t* retrieve_next_token,
    const int64_t* retrieve_next_sibling,
    const int64_t* draft_tokens,
    GrammarMatcher& matcher,
    DLTensor* bitmask
) {
  int32_t* bitmask_data = reinterpret_cast<int32_t*>(bitmask->data);
  int32_t bitmask_size = static_cast<int32_t>(bitmask->shape[1]);

  bool accepted;
  if (curr == 0) {
    // The first token generated by the target model, always accepted
    accepted = true;
  } else {
    int32_t curr_token_id = draft_tokens[curr];
    int32_t* parent_bitmask = bitmask_data + parent_pos * bitmask_size;
    // 32 boolean bitmask values are packed into 32-bit integers
    accepted = (parent_bitmask[curr_token_id / 32] & (1 << (curr_token_id % 32))) != 0;
  }

  if (accepted) {
    if (curr != 0) {
      matcher.AcceptToken(draft_tokens[curr]);
    }

    if (!matcher.IsTerminated()) {
      matcher.FillNextTokenBitmask(bitmask, curr);

      if (retrieve_next_token[curr] != -1) {
        DFS(retrieve_next_token[curr],
            curr,
            retrieve_next_token,
            retrieve_next_sibling,
            draft_tokens,
            matcher,
            bitmask);
      }
    }

    if (curr != 0) {
      matcher.Rollback(1);
    }
  }

  if (retrieve_next_sibling[curr] != -1) {
    DFS(retrieve_next_sibling[curr],
        parent_pos,
        retrieve_next_token,
        retrieve_next_sibling,
        draft_tokens,
        matcher,
        bitmask);
  }
}

}  // namespace details

void TraverseDraftTree(
    const DLTensor* retrieve_next_token,
    const DLTensor* retrieve_next_sibling,
    const DLTensor* draft_tokens,
    GrammarMatcher& matcher,
    DLTensor* bitmask
) {
  // Check dtype
  XGRAMMAR_CHECK(retrieve_next_token->dtype.code == kDLInt && retrieve_next_token->dtype.bits == 64)
      << "The retrieve_next_token tensor must be int64";
  XGRAMMAR_CHECK(
      retrieve_next_sibling->dtype.code == kDLInt && retrieve_next_sibling->dtype.bits == 64
  ) << "The retrieve_next_sibling tensor must be int64";
  XGRAMMAR_CHECK(draft_tokens->dtype.code == kDLInt && draft_tokens->dtype.bits == 64)
      << "The draft_tokens tensor must be int64";
  XGRAMMAR_CHECK(bitmask->dtype.code == kDLInt && bitmask->dtype.bits == 32)
      << "The bitmask tensor must be int32";

  XGRAMMAR_CHECK(retrieve_next_token->shape[0] == retrieve_next_sibling->shape[0])
      << "The retrieve_next_token and retrieve_next_sibling tensors must have the same length";
  XGRAMMAR_CHECK(retrieve_next_token->shape[0] == draft_tokens->shape[0])
      << "The retrieve_next_token and draft_tokens tensors must have the same length";

  details::DFS(
      0,
      -1,
      reinterpret_cast<const int64_t*>(retrieve_next_token->data),
      reinterpret_cast<const int64_t*>(retrieve_next_sibling->data),
      reinterpret_cast<const int64_t*>(draft_tokens->data),
      matcher,
      bitmask
  );
}

}  // namespace xgrammar
