import json, re
import gradio as gr
import numpy as np
import chardet

from taiko import (
    preprocess,
    generate_taiko_wav,
    COURSE
)

def txt_loads(data_str:str):
    data = []
    current_course = None
    chart = []

    for line in data_str.split('\n'):
        if line.startswith("course:"):
            if current_course is not None:
                data.append({"course": current_course, "chart": chart})
                chart = []
            current_course = int(re.search(r'\d+', line).group())
        elif line:
            line = line.strip("[]")
            nums = line.split(",")
            chart.append([int(nums[0]), float(nums[1])])

    if current_course is not None:
        data.append({"course": current_course, "chart": chart})

    output_data = {"data": data}
    return output_data

def handle(chart_path:str, music_path:str):
    with open(chart_path, "rb") as file:
        raw_data = file.read(4096)
        result = chardet.detect(raw_data)
        encoding = result["encoding"]

    if chart_path.endswith(".json"):
        data = json.loads(open(chart_path, "r", encoding=encoding).read())
    elif chart_path.endswith(".txt"):
        data = txt_loads(open(chart_path, "r", encoding=encoding).read())
    
    if len(data["data"]) < 0 and len(data["data"]) > 5:
        raise("Issue occur: chart data")

    for d in data["data"]:
        chart = preprocess(d["chart"])
        audio = generate_taiko_wav(chart, music_path)
        COURSE[d["course"]]["audio"] = audio
    
    c = 2
    if music_path is None: c = 1

    return \
        (COURSE[0]["audio"].frame_rate * c, np.array(COURSE[0]["audio"].get_array_of_samples())), \
        (COURSE[1]["audio"].frame_rate * c, np.array(COURSE[1]["audio"].get_array_of_samples())), \
        (COURSE[2]["audio"].frame_rate * c, np.array(COURSE[2]["audio"].get_array_of_samples())), \
        (COURSE[3]["audio"].frame_rate * c, np.array(COURSE[3]["audio"].get_array_of_samples())), \
        (COURSE[4]["audio"].frame_rate * c, np.array(COURSE[4]["audio"].get_array_of_samples()))

if __name__ == "__main__":
    inputs = [
        gr.File(label="太鼓達人譜面/Taiko Chart Data (json or txt)", file_types=[".json", ".txt"]),
        gr.File(label="譜面音樂/Chart Music (ogg) (Optional)", file_types=[".ogg"])
    ]
    outputs = [gr.Audio(label=course["label"], format="mp3") for course in COURSE]

    demo = gr.Interface(fn=handle, inputs=inputs, outputs=outputs, title="程設二作業HW0105 / Taiko Music Generator")
    demo.launch()