import unittest
import os
import shutil
import asyncio
from importlib.resources import files
from engines.f5_tts import F5TTSEngine

class TestF5TTSEngine(unittest.TestCase):
    async def asyncSetUp(self):
        self.engine = F5TTSEngine()
        self.voices_dir = "engines/f5-tts-voices"
        self.test_speaker_name = "test_speaker"
        self.test_speaker_wav = os.path.join(self.voices_dir, f"{self.test_speaker_name}.wav")
        self.test_speaker_txt = os.path.join(self.voices_dir, f"{self.test_speaker_name}.txt")

        # Create a dummy speaker for testing
        if not await asyncio.to_thread(os.path.exists, self.test_speaker_wav):
            default_wav_path = str(files("f5_tts").joinpath("infer/examples/basic/basic_ref_en.wav"))
            await asyncio.to_thread(shutil.copy, default_wav_path, self.test_speaker_wav)
        if not await asyncio.to_thread(os.path.exists, self.test_speaker_txt):
            await asyncio.to_thread(lambda: open(self.test_speaker_txt, "w").write("Some call me nature, others call me mother nature."))

        # Reload speakers to include the new test speaker
        self.engine._load_speakers()

    async def asyncTearDown(self):
        # Clean up the dummy speaker files
        if await asyncio.to_thread(os.path.exists, self.test_speaker_wav):
            await asyncio.to_thread(os.remove, self.test_speaker_wav)
        if await asyncio.to_thread(os.path.exists, self.test_speaker_txt):
            await asyncio.to_thread(os.remove, self.test_speaker_txt)

    async def test_synthesize_default_wav(self):
        if not self.engine.model:
            self.skipTest("F5-TTS model not initialized.")
            
        text = "Hello, this is a test with the default voice."
        audio_file = await self.engine.synthesize(text, fmt="wav")
        
        self.assertTrue(await asyncio.to_thread(os.path.exists, audio_file))
        self.assertTrue(await asyncio.to_thread(os.path.getsize, audio_file) > 0)
        
        await asyncio.to_thread(os.remove, audio_file)

    async def test_synthesize_custom_speaker_wav(self):
        if not self.engine.model:
            self.skipTest("F5-TTS model not initialized.")
            
        text = "Hello, this is a test with a custom voice."
        audio_file = await self.engine.synthesize(text, speaker=self.test_speaker_name, fmt="wav")
        
        self.assertTrue(await asyncio.to_thread(os.path.exists, audio_file))
        self.assertTrue(await asyncio.to_thread(os.path.getsize, audio_file) > 0)
        
        await asyncio.to_thread(os.remove, audio_file)
        
    async def test_synthesize_ogg(self):
        if not self.engine.model:
            self.skipTest("F5-TTS model not initialized.")
            
        text = "Hello, this is a test in ogg format."
        audio_file = await self.engine.synthesize(text, fmt="ogg")
        
        self.assertTrue(await asyncio.to_thread(os.path.exists, audio_file))
        self.assertTrue(await asyncio.to_thread(os.path.getsize, audio_file) > 0)
        
        await asyncio.to_thread(os.remove, audio_file)

if __name__ == '__main__':
    unittest.main()