"""THIS FILE JUST DOES ALIGNMENT...FAST

Takes the saved words times and aligns them.
This is flexible as we can adjust the algiment's max duration.
"""

import os
import json
import argparse
import tqdm
from tqdm.contrib.concurrent import process_map
from multiprocessing import cpu_count
import random
import time
from suno_utils.tasks.hoot import get_aligned_lyrics
from suno_utils.utils.text import write_jsonl


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", type=str)
    args = parser.parse_args()
    return args


OUTPUT_BASE_DIR = "/home/christian/data/hoot/alignments"
MAX_DURATION = "2min"

if __name__ == "__main__":
    input_args = parse_args()
    if input_args.dataset == "genius":
        INPUT_PATH = "/home/tony/Data/Hoot/metas.json"
        DOWNLOAD_PATH = "/app/suno/data/hoot"
        OUTPUT_PATH = os.path.join(OUTPUT_BASE_DIR, "outputs_3")
        ALIGNED_OUTPUT_PATH = os.path.join(OUTPUT_BASE_DIR, f"genius_{MAX_DURATION}")
    elif input_args.dataset == "ytm":
        INPUT_PATH = "/home/tony/Data/Hoot/ytm_metas.json"
        DOWNLOAD_PATH = "/app/suno/data/hoot_ytm"
        OUTPUT_PATH = os.path.join(OUTPUT_BASE_DIR, "outputs_3_ytm")
        ALIGNED_OUTPUT_PATH = os.path.join(OUTPUT_BASE_DIR, f"ytm_{MAX_DURATION}")
    elif input_args.dataset == "deezer":
        INPUT_PATH = "/home/tony/Data/Hoot/deezer_metas.json"
        DOWNLOAD_PATH = "/app/suno/data/hoot_deezer"
        OUTPUT_PATH = os.path.join(OUTPUT_BASE_DIR, "outputs_3_deezer")
        ALIGNED_OUTPUT_PATH = os.path.join(OUTPUT_BASE_DIR, f"deezer_{MAX_DURATION}")
    else:
        raise ValueError("Unknown dataset")

    # load the infomration
    with open(INPUT_PATH, "r") as fp:
        metas = json.load(fp)
    meta_id_map = {meta["id"]: meta for meta in metas}
    # with open("/home/tony/Data/Hoot/alignments/all_end_time.json", "r") as fp:
    #     all_end_time = json.load(fp)
    # make the aligned output path
    os.makedirs(ALIGNED_OUTPUT_PATH, exist_ok=True)

    def generatge_alinged_lyrics_map(index):
        random.seed(index)
        try:
            output_jsonl_path = f"{ALIGNED_OUTPUT_PATH}/batch_{index}.jsonl"
            if os.path.exists(output_jsonl_path):
                # let's avoid the race condition
                time.sleep(1)
                return
            # print("index")
            with open(output_jsonl_path, "w") as fp:
                fp.write("")
            with open(f"{OUTPUT_PATH}/batch_{index}.json", "r") as fp:
                hooted_outputs = json.load(fp)
            # # if you want to refine the endtime
            # refined_hoot_outputs = []
            for meta in hooted_outputs:
                meta_id = meta["id"]
                # up to v3 we still need this hack..as the original lyrics are cased
                meta["lyrics"] = (
                    meta_id_map[meta_id]["lyrics"]
                    if "lyrics" in meta_id_map[meta_id]
                    else meta_id_map[meta_id]["text"]
                )
            # filter some inputs
            aligned_lyrics_map = get_aligned_lyrics(
                hooted_outputs,
                silent=True,
                min_cer=0.8,
                min_p_align=0.0,
                aligned_text_frac=0.9,
                return_only_valid=True,
            )
            # print(len(hooted_outputs), len(aligned_lyrics_map))
            write_jsonl(
                [(k, v) for k, v in aligned_lyrics_map.items()],
                output_jsonl_path,
            )
        except Exception as e:
            print(f"WTF {index}, {e}")
        return

    print("Start aligning lyrics.")
    input_jobs = []
    for i in tqdm.tqdm(range(0, len(metas), 100)):
        # this is for testing only
        # for i in tqdm.tqdm(range(0, 100, 100)):
        input_jobs.append(i)
    print(f"Total jobs: {len(input_jobs)}")

    process_map(
        generatge_alinged_lyrics_map,
        input_jobs,
        max_workers=cpu_count() - 4,
        chunksize=1,
    )
    print("DONE!!!")