# Copyright (c)  2023  Xiaomi Corporation

import logging

try:
    import click
except ImportError:
    print("Please run")
    print("  pip install click")
    print("before you continue")
    raise

from pathlib import Path
from sherpa_onnx import text2token


@click.group()
def cli():
    """
    The shell entry point to sherpa-onnx.
    """
    logging.basicConfig(
        format="%(asctime)s %(levelname)s [%(filename)s:%(lineno)d] %(message)s",
        level=logging.INFO,
    )


@cli.command(name="text2token")
@click.argument("input", type=click.Path(exists=True, dir_okay=False))
@click.argument("output", type=click.Path())
@click.option(
    "--tokens",
    type=str,
    required=True,
    help="The path to tokens.txt.",
)
@click.option(
    "--tokens-type",
    type=click.Choice(
        [
            "cjkchar",
            "bpe",
            "cjkchar+bpe",
            "fpinyin",
            "ppinyin",
            "phone+ppinyin",
        ],
        case_sensitive=True,
    ),
    required=True,
    help="""The type of modeling units, should be cjkchar, bpe, cjkchar+bpe, fpinyin, ppinyin or phone+ppinyin.
    fpinyin means full pinyin, each cjkchar has a pinyin(with tone).
    ppinyin means partial pinyin, it splits pinyin into initial and final,
    phone means English phonemes in CMU dictionary format.
    """,
)
@click.option(
    "--bpe-model",
    type=str,
    help="The path to bpe.model. Only required when tokens-type is bpe or cjkchar+bpe.",
)
@click.option(
    "--lexicon",
    type=str,
    help="The path to lexicon.txt. Only required when tokens-type is phone+ppinyin.",
)
def encode_text(
    input: Path,
    output: Path,
    tokens: Path,
    tokens_type: str,
    bpe_model: Path,
    lexicon: Path,
):
    """
    Encode the texts given by the INPUT to tokens and write the results to the OUTPUT.
    Each line in the texts contains the original phrase, it might also contain some
    extra items, for example, the boosting score (starting with :), the triggering
    threshold (starting with #, only used in keyword spotting task) and the original
    phrase (starting with @). Note: the extra items will be kept same in the output.

    example input 1 (tokens_type = ppinyin):

    小爱同学 :2.0 #0.6 @小爱同学
    你好问问 :3.5 @你好问问
    小艺小艺 #0.6 @小艺小艺

    example output 1:

    x iǎo ài t óng x ué :2.0 #0.6 @小爱同学
    n ǐ h ǎo w èn w èn :3.5 @你好问问
    x iǎo y ì x iǎo y ì #0.6 @小艺小艺

    example input 2 (tokens_type = bpe):

    HELLO WORLD :1.5 #0.4
    HI GOOGLE :2.0 #0.8
    HEY SIRI #0.35

    example output 2:

    ▁HE LL O ▁WORLD :1.5 #0.4
    ▁HI ▁GO O G LE :2.0 #0.8
    ▁HE Y ▁S I RI #0.35
    """
    texts = []
    # extra information like boosting score (start with :), triggering threshold (start with #)
    # original keyword (start with @)
    extra_info = []
    with open(input, "r", encoding="utf8") as f:
        for line in f:
            extra = []
            text = []
            toks = line.strip().split()
            for tok in toks:
                if tok[0] == ":" or tok[0] == "#" or tok[0] == "@":
                    extra.append(tok)
                else:
                    text.append(tok)
            texts.append(" ".join(text))
            extra_info.append(extra)

    encoded_texts = text2token(
        texts,
        tokens=tokens,
        tokens_type=tokens_type,
        bpe_model=bpe_model,
        lexicon=lexicon,
    )
    with open(output, "w", encoding="utf8") as f:
        for i, txt in enumerate(encoded_texts):
            txt += extra_info[i]
            f.write(" ".join(txt) + "\n")
