import os

from typing import Union

from fastapi import FastAPI
from fastapi.responses import FileResponse
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from datetime import datetime

model_id = "damo/speech_sambert-hifigan_tts_zhitian_emo_zh-cn_16k"
sambert_hifigan_tts = pipeline(task=Tasks.text_to_speech, model=model_id)

app = FastAPI()


def generate_audio(text: str):
    output = sambert_hifigan_tts(input=text)
    wav = output[OutputKeys.OUTPUT_WAV]
    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")  # Generate timestamp
    filename = f"media/audio/{timestamp}.wav"  # Use timestamp as filename
    os.makedirs(os.path.dirname(filename), exist_ok=True)  # Create directories if they don't exist
    with open(filename, "wb") as f:
        f.write(wav)
    return filename


@app.get("/tts", responses={200: {"content": {"audio/wav": {}}}})
def get_tts(text: Union[str, None] = None):
    filename = generate_audio(text)
    return FileResponse(filename, filename=f"{os.path.basename(filename)}")