import pytest
from test_data.texts import read_translation
from lib.utils import save_csv
from environment import REPORTS_DIR


@pytest.fixture(scope="module")
def tts(get_platform):
    if get_platform == "apple":
        from lib.models.mac.kokoro import KokoroTTS
        return KokoroTTS
    elif get_platform == "intel":
        from lib.models.intel.kokoro import KokoroTTS
        return KokoroTTS


def test_tts_zh(tts):
    tts = tts.from_language("zh")
    report = []
    for text in read_translation("zh", count_limit=20):
        samples, sr, time_cost = tts.generate(text)
        report.append([text, time_cost])
    save_csv(REPORTS_DIR / "test_tts_zh.csv", ["text", "time"], report)
    speed = [i[1] for i in report]
    print("tts zh average speed: ", sum(speed) / len(speed))


def test_tts_en(tts):
    tts = tts.from_language("en")
    report = []
    for text in read_translation("en",count_limit=20):
        samples, sr, time_cost = tts.generate(text, speed=1.4)
        report.append([text, time_cost])
    save_csv(REPORTS_DIR / "test_tts_en.csv", ["text", "time"], report)
    speed = [i[1] for i in report]
    print("tts en average speed: ", sum(speed) / len(speed))