#include "transcription.hpp"
#include "tokenizer.hpp"

#include <algorithm>
#include <cctype>
#include <set>
#include <string>
#include <vector>

namespace pk {

namespace {

// U+2581 LOWER ONE EIGHTH BLOCK — SentencePiece meta-space marker (3 bytes).
const char    META_SPACE[]  = "\xe2\x96\x81";
const size_t  META_SPACE_LEN = 3;

bool starts_with_meta(const std::string& piece) {
    return piece.size() >= META_SPACE_LEN &&
           (unsigned char)piece[0] == 0xE2 &&
           (unsigned char)piece[1] == 0x96 &&
           (unsigned char)piece[2] == 0x81;
}

// decode_ids_to_str([id]) for a single piece: replace `▁`->space, strip a single
// leading space (mirrors detokenize() on a 1-element id list).
std::string piece_to_text(const std::string& piece) {
    std::string out;
    out.reserve(piece.size());
    for (size_t i = 0; i < piece.size();) {
        if (i + META_SPACE_LEN <= piece.size() &&
            (unsigned char)piece[i]   == 0xE2 &&
            (unsigned char)piece[i+1] == 0x96 &&
            (unsigned char)piece[i+2] == 0x81) {
            out += ' ';
            i += META_SPACE_LEN;
        } else {
            out += piece[i++];
        }
    }
    if (!out.empty() && out[0] == ' ') out.erase(0, 1);
    return out;
}

// Mirror of NeMo extract_punctuation_from_vocab: a single-char string is a
// supported punctuation mark if it is Unicode category 'P*' AND its containing
// piece is NOT a special token ([...] / <...> / ## / ▁-prefixed / whitespace).
// For these subword vocabularies the punctuation marks are ASCII, so we classify
// the ASCII punctuation code points (Unicode category P) directly. The result is
// the set of single-character punctuation *strings* (NeMo iterates over each
// char of each non-special piece).
bool is_ascii_punct(unsigned char c) {
    // Unicode category starts with 'P' for these ASCII code points:
    //   ! " # % & ' ( ) * , - . / : ; ? @ [ \ ] _ { }
    switch (c) {
        case '!': case '"': case '#': case '%': case '&': case '\'':
        case '(': case ')': case '*': case ',': case '-': case '.':
        case '/': case ':': case ';': case '?': case '@': case '[':
        case '\\': case ']': case '_': case '{': case '}':
            return true;
        default:
            return false;
    }
}

bool is_special_token(const std::string& tok) {
    if (tok.empty()) return true;                      // whitespace-only -> special
    if (tok.front() == '[' && tok.back() == ']') return true;
    if (tok.front() == '<' && tok.back() == '>') return true;
    if (tok.size() >= 2 && tok[0] == '#' && tok[1] == '#') return true;
    if (starts_with_meta(tok)) return true;            // ▁-prefixed
    // whitespace-only
    if (std::all_of(tok.begin(), tok.end(),
                    [](char ch){ return std::isspace((unsigned char)ch); }))
        return true;
    return false;
}

std::set<std::string> extract_punctuation(const std::vector<std::string>& pieces) {
    std::set<std::string> punct;
    for (const std::string& tok : pieces) {
        if (is_special_token(tok)) continue;
        for (unsigned char c : tok) {
            if (is_ascii_punct(c)) punct.insert(std::string(1, (char)c));
        }
    }
    return punct;
}

} // namespace

std::vector<Word> group_words(const std::vector<TokenInfo>& tokens,
                              const std::vector<std::string>& pieces,
                              float frame_sec) {
    std::vector<Word> words;
    const int n = (int)tokens.size();
    if (n == 0) return words;

    const std::set<std::string> punct = extract_punctuation(pieces);
    const std::string DELIM = " "; // word delimiter (NeMo word_seperator default)

    auto piece_of = [&](int i) -> const std::string& {
        static const std::string empty;
        int id = tokens[i].id;
        if (id >= 0 && (size_t)id < pieces.size()) return pieces[(size_t)id];
        return empty;
    };
    auto is_punct = [&](const std::string& s) {
        return !s.empty() && s != DELIM && punct.count(s) > 0;
    };

    // Per-token char offsets (NeMo char_offsets): text = decode_ids_to_str([id]),
    // token = raw piece, start/end in encoder frames. end_offset = frame + span
    // (TDT duration; CTC run-length supplied by the caller via TokenInfo.span).
    std::vector<std::string> text(n), tok(n);
    std::vector<int32_t>      start(n), end(n);
    std::vector<float>        conf(n);
    for (int i = 0; i < n; ++i) {
        tok[i]   = piece_of(i);
        text[i]  = piece_to_text(tok[i]);
        start[i] = tokens[i].frame;
        end[i]   = tokens[i].frame + tokens[i].span;
        conf[i]  = tokens[i].conf;
    }

    // NeMo _refine_timestamps / _refine_timestamps_tdt: a punctuation token
    // (text[0] in supported_punctuation, i>0) is pinned to the previous token's
    // end (start = prev end; end = start). This is the TDT refinement; NeMo's CTC
    // refine only sets end = start, but here start is never used for a punct token
    // (it merges into the previous word and only its end matters), so the unified
    // rule reproduces both heads' dumped baseline word offsets exactly.
    for (int i = 0; i < n; ++i) {
        if (!text[i].empty() && i > 0) {
            std::string first(1, text[i][0]);
            if (is_punct(first)) {
                start[i] = end[i - 1];
                end[i]   = start[i];
            }
        }
    }

    // NeMo get_words_offsets (tokenizer_type='bpe', word_delimiter=' '):
    //   word-start condition: token != token_text  (i.e. the piece differs from
    //   its decoded text -> a `▁`-prefixed sub-word starts a new word), OR the
    //   token IS the delimiter and the next non-delimiter token is not punct.
    // built holds the indices of the current word's tokens; previous_token_index
    // marks the word's first token.
    std::vector<int> built;
    int prev = 0;

    auto detok_built = [&](const std::vector<int>& idxs) {
        std::vector<int32_t> ids;
        ids.reserve(idxs.size());
        for (int k : idxs) ids.push_back(tokens[k].id);
        return detokenize(pieces, ids);
    };
    auto min_conf = [&](const std::vector<int>& idxs) {
        float m = 1.0f;
        for (int k : idxs) m = std::min(m, conf[k]);
        return m;
    };

    for (int i = 0; i < n; ++i) {
        const std::string& ct = text[i];
        const std::string& tk = tok[i];
        const bool curr_punct = is_punct(ct);

        // next non-delimiter token text (NeMo lookahead).
        std::string next_non_delim;
        int j = i;
        while (next_non_delim.empty() && j < n - 1) {
            ++j;
            if (text[j] != DELIM) next_non_delim = text[j];
        }
        const bool next_is_punct = !next_non_delim.empty() && is_punct(next_non_delim);

        const bool word_start_cond =
            (tk != ct) || (ct == DELIM && !next_is_punct);

        if (word_start_cond && !curr_punct) {
            if (!built.empty()) {
                Word w;
                w.text  = detok_built(built);
                w.start = (float)start[prev]      * frame_sec;
                w.end   = (float)end[built.back()] * frame_sec;
                w.conf  = min_conf(built);
                words.push_back(std::move(w));
            }
            built.clear();
            if (ct != DELIM) {
                built.push_back(i);
                prev = i;
            }
        } else if (curr_punct && built.empty() && !words.empty()) {
            // Punctuation with no open word: attach to the previous word, extend
            // its end, drop a trailing space, append the punctuation char.
            Word& lw = words.back();
            lw.end = (float)end[i] * frame_sec;
            if (!lw.text.empty() && lw.text.back() == ' ') lw.text.pop_back();
            lw.text += ct;
            lw.conf = std::min(lw.conf, conf[i]);
        } else if (curr_punct && !built.empty()) {
            // Punctuation closing an open word: drop a trailing delimiter token,
            // then append this token.
            if (!built.empty()) {
                const std::string& last = tok[built.back()];
                if (last == " " || last == "_" || last == META_SPACE) built.pop_back();
            }
            built.push_back(i);
        } else {
            // Continuation sub-word: extend the current word.
            if (built.empty()) prev = i;
            built.push_back(i);
        }
    }

    // NeMo tail handling: force the first word's start to the first token's start,
    // and flush any remaining built tokens as the final word.
    if (!words.empty()) {
        words[0].start = (float)start[0] * frame_sec;
        if (!built.empty()) {
            Word w;
            w.text  = detok_built(built);
            w.start = (float)start[prev]      * frame_sec;
            w.end   = (float)end[built.back()] * frame_sec;
            w.conf  = min_conf(built);
            words.push_back(std::move(w));
        }
    } else if (!built.empty()) {
        Word w;
        w.text  = detok_built(built);
        w.start = (float)start[0]          * frame_sec;
        w.end   = (float)end[built.back()] * frame_sec;
        w.conf  = min_conf(built);
        words.push_back(std::move(w));
    }

    return words;
}

} // namespace pk