from flask import Flask, jsonify, request
import logging
import os
import numpy as np
import re
import subprocess
import tempfile
import textgrids
import base64
import scipy.io as scipy_io

from . import config

app = Flask(__name__)


MFA_PYTHON_CMD_PTN = (
    'bash -c "source {conda_sh_filepath}; conda activate {env_name}; '
    "mfa align --clean "
    "-t {temp_temp_dir} -j {num_cores} {pairs_dir} {dictionary_name} "
    '{acoustic_model_name} {out_dir} --retry_beam 4000"'
)


logger = logging.getLogger(__name__)


def _load_textgrid_file(filepath):
    grid = textgrids.TextGrid(filepath)
    word_aligment = []
    for t in grid["words"]:
        if len(t.text) > 0:
            word_aligment.append(((t.xmin, t.xmax), t.text))
    return word_aligment


def _fixup_mfa_output(word_aligment):
    # attach 's to preceeding word
    fixed_word_alignment = []
    found_merge = None
    for (start_s, end_s), text in word_aligment[::-1]:
        if found_merge is not None:
            fixed_word_alignment.append(
                ((start_s, found_merge[0][1]), text + found_merge[1])
            )
            found_merge = None
        elif text == "'s":
            found_merge = (start_s, end_s), text
        else:
            fixed_word_alignment.append(((start_s, end_s), text))
    return fixed_word_alignment[::-1]


def _mfa_align_text(
    audio_bytes_list,
    transcripts,
    num_cores=config.MAX_CPU,
    dictionary_name=config.MFA_DICTIONARY,
    acoustic_model_name=config.MFA_ACOUSTIC_MODEL,
):
    if len(audio_bytes_list) != len(transcripts):
        raise ValueError("different number of audio files and transcripts provided")
    audio_bytes_list = [base64.b64decode(ab) for ab in audio_bytes_list]
    audio_bytes_list = [np.frombuffer(ab, dtype=np.int16) for ab in audio_bytes_list]
    with tempfile.TemporaryDirectory() as temp_dir:
        pairs_dir = os.path.join(temp_dir, "pairs")
        out_dir = os.path.join(temp_dir, "out")
        os.mkdir(pairs_dir)
        temp_temp_dir = os.path.join(temp_dir, "tmp")
        os.mkdir(temp_temp_dir)

        # make one folder for each core
        for n in range(num_cores):
            sub_pairs_dir = os.path.join(pairs_dir, "{}".format(n))
            os.mkdir(sub_pairs_dir)
        for n, (audio_bytes, transcript) in enumerate(
            zip(audio_bytes_list, transcripts)
        ):
            assigned_core = n % num_cores
            sub_pairs_dir = os.path.join(pairs_dir, "{}".format(assigned_core))
            with open(os.path.join(sub_pairs_dir, "{}.txt".format(n)), "w") as f:
                f.write(transcript)
            audio_fname = os.path.join(sub_pairs_dir, "{}.wav".format(n))
            scipy_io.wavfile.write(audio_fname, 16_000, audio_bytes)

        cmd = MFA_PYTHON_CMD_PTN.format(
            conda_sh_filepath=config.CONDA_SH_FILEPATH,
            env_name=config.MFA_ENV_NAME,
            num_cores=num_cores,
            pairs_dir=pairs_dir,
            temp_temp_dir=temp_temp_dir,
            out_dir=out_dir,
            dictionary_name=dictionary_name,
            acoustic_model_name=acoustic_model_name,
        )
        out_code = subprocess.call(
            cmd, shell=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL
        )
        if out_code != 0:
            raise ValueError("mfa error")
        word_alignments = []
        for n in range(len(audio_bytes_list)):
            assigned_core = n % num_cores
            sub_out_dir = os.path.join(out_dir, "{}".format(assigned_core))
            textgrid_filepath = os.path.join(sub_out_dir, "{}.TextGrid".format(n))
            if not os.path.exists(textgrid_filepath):
                word_alignments.append(None)
                continue
            word_alignment = _load_textgrid_file(textgrid_filepath)
            word_alignments.append(word_alignment)
    # check alignment and fix if possible
    n_failed = 0
    for n in range(len(word_alignments)):
        word_alignment = word_alignments[n]
        if word_alignment is None:
            n_failed += 1
            continue
        transcript = transcripts[n]
        if transcript != " ".join([w for _, w in word_alignment]):
            b_align_failed = True
            # let's see if we can fix up alignment errors
            if "'s" in [w for _, w in word_alignment]:
                fixed_word_alignment = _fixup_mfa_output(word_alignment)
                if transcript == " ".join([w for _, w in fixed_word_alignment]):
                    word_alignments[n] = fixed_word_alignment
                    b_align_failed = False
            if b_align_failed:
                word_alignments[n] = None
                n_failed += 1
    if n_failed > 0:
        logger.warning(f"{n_failed}/{len(word_alignments)} failed.")
    return word_alignments


def _mfa_normalize(token):
    if re.search(r"\s", token):
        raise NotImplementedError("tokens must be whitespace-free")
    return re.sub(r"[^a-z0-9\']", "", token.lower()).strip("'")


@app.route("/align_tokens", methods=["POST"])
def align_tokens():
    content = request.json
    audio_bytes = content["audio_bytes"]
    transcript_tokens = content["transcript_tokens"]
    if len(audio_bytes) != len(transcript_tokens):
        raise ValueError("different number of audio files and transcripts provided")
    tokens_idx = []
    transcripts = []
    for tokens in transcript_tokens:
        tmp_idx = []
        tmp_tokens = []
        for n, token in enumerate(tokens):
            if token["type"] != "text":
                continue
            norm_token = _mfa_normalize(token["value"])
            if len(norm_token) == 0:
                continue
            tmp_idx.append(n)
            tmp_tokens.append(norm_token)
        tokens_idx.append(tmp_idx)
        transcripts.append(" ".join(tmp_tokens))
    word_alignments = _mfa_align_text(audio_bytes, transcripts)

    if not (len(tokens_idx) == len(word_alignments) == len(transcript_tokens)):
        raise ValueError("something went wrong")
    expanded_word_timestamps = []
    for token_idx, word_alignment, tokens in zip(
        tokens_idx, word_alignments, transcript_tokens
    ):
        word_timestamps = [None] * len(tokens)
        if word_alignment is not None:
            for a, (b, _) in zip(token_idx, word_alignment):
                word_timestamps[a] = b
        expanded_word_timestamps.append(word_timestamps)
    return jsonify(expanded_word_timestamps)


@app.route("/")
def hello_world():
    return "<h1>hello world</h1>"


if __name__ == "__main__":
    app.run(host="0.0.0.0", debug=True, port=7866)