"""
NovaAi – TTS-Engine-Hub
f5_tts.py
Version: v0.0.2

Description:
  F5-TTS engine module.
  Implements the TTSEngineBase for F5-TTS text-to-speech synthesis.
  Now with robust speaker handling.

Author: Your Name (or leave as generated)
Date: 2025-12-03
"""

import os
import tempfile
import torch
import torchaudio
import numpy as np
import soundfile as sf
import asyncio
from .engine_base import TTSEngineBase
from importlib.resources import files

try:
    from f5_tts.api import F5TTS
except ImportError:
    print("Warning: F5TTS could not be imported. F5-TTS engine will not be available.")
    F5TTS = None

class F5TTSEngine(TTSEngineBase):
    def __init__(self):
        # Initialize F5-TTS specific resources, models, etc.
        print("F5-TTS Engine Initializing...")
        self.speakers = {}
        self.model = None

        if F5TTS:
            try:
                self.model = F5TTS(model="F5TTS_v1_Base")
                print("F5-TTS Engine Initialized.")
                self._load_speakers()
            except Exception as e:
                print(f"Error initializing F5-TTS Engine: {e}")
                self.model = None
        else:
            print("F5-TTS Engine not initialized because F5TTS is not available.")

    def _load_speakers(self):
        # Add the default speaker
        default_wav = str(files("f5_tts").joinpath("infer/examples/basic/basic_ref_en.wav"))
        default_txt = "engines/f5-tts-voices/default.txt"
        if os.path.exists(default_txt):
            self.speakers["default"] = {"wav": default_wav, "txt": default_txt}
        
        # Scan for custom speakers
        voices_dir = "engines/f5-tts-voices"
        if not os.path.isdir(voices_dir):
            return
        for file in os.listdir(voices_dir):
            if file.endswith(".wav"):
                speaker_name = file.rsplit('.', 1)[0]
                wav_path = os.path.join(voices_dir, file)
                txt_path = os.path.join(voices_dir, f"{speaker_name}.txt")
                if os.path.exists(txt_path):
                    self.speakers[speaker_name] = {"wav": wav_path, "txt": txt_path}
                    print(f"Found custom speaker: {speaker_name}")

    def _blocking_synthesize(self, text: str, speaker: str, fmt: str):
        """The actual blocking synthesis logic."""
        speaker_data = self.speakers[speaker]
        ref_file = speaker_data["wav"]
        
        with open(speaker_data["txt"], 'r') as f:
            ref_text = f.read()

        print(f"F5-TTS: Synthesizing '{text}' with reference voice from '{ref_file}'.")
        
        wav, sr, spec = self.model.infer(
            ref_file=ref_file,
            ref_text=ref_text,
            gen_text=text,
        )
        
        with tempfile.NamedTemporaryFile(delete=False, suffix=f".{fmt}") as temp_file:
            if fmt == "wav":
                torchaudio.save(temp_file.name, torch.from_numpy(wav).unsqueeze(0), sr, format="wav")
            else:
                # Convert to float32 for soundfile
                wav_float = wav.astype(np.float32) / np.iinfo(wav.dtype).max
                sf.write(temp_file.name, wav_float, sr)
            return temp_file.name

    async def synthesize(self, text: str, speaker: str = None, model: str = None, fmt: str = "ogg"):
        """
        Asynchronously generate speech audio from text input using F5-TTS.
        """
        if not self.model:
            raise RuntimeError("F5-TTS Engine not initialized.")
        
        speaker_to_use = speaker if speaker in self.speakers else "default"
        if speaker and speaker not in self.speakers:
            print(f"Warning: Speaker '{speaker}' not found. Falling back to default speaker.")

        if speaker_to_use not in self.speakers:
             raise RuntimeError("No default speaker found for F5-TTS. Please add a 'default.wav' and 'default.txt' to the 'engines/f5-tts-voices' directory.")

        try:
            # Run the blocking synthesis in a separate thread
            return await asyncio.to_thread(self._blocking_synthesize, text, speaker_to_use, fmt)
        except Exception as e:
            raise RuntimeError(f"F5-TTS synthesis failed: {e}")

    def list_models(self):
        """Return a list of available F5-TTS models."""
        if not self.model:
            return []
        return ["F5TTS_v1_Base"]

    def list_voices(self, model: str = None):
        """Return a list of available F5-TTS voices for a model."""
        return list(self.speakers.keys())

    def healthcheck(self):
        """Return health/status info for F5-TTS engine."""
        if self.model:
            return {"status": "ok", "message": "F5-TTS engine is ready"}
        else:
            return {"status": "error", "message": "F5-TTS engine failed to initialize"}

    async def selftest(self):
        """Run internal self-test for F5-TTS."""
        if not self.model:
            return {"status": "failed", "message": "F5-TTS Engine not initialized."}
        try:
            # Await the async synthesize method
            audio_file = await self.synthesize("this is a test.")
            selftest_passed = os.path.exists(audio_file) and os.path.getsize(audio_file) > 0
            if selftest_passed:
                os.remove(audio_file)
            return {"status": "passed" if selftest_passed else "failed", "message": "F5-TTS self-test successful"}
        except Exception as e:
            return {"status": "failed", "message": f"F5-TTS self-test failed: {e}"}