import os
import sys
import subprocess
import tempfile
import soundfile as sf
import librosa
from fastapi import FastAPI, File, Form, UploadFile
from fastapi.responses import FileResponse
from fastapi.middleware.cors import CORSMiddleware

# 1. АВТОМАТИЧЕСКОЕ СКАЧИВАНИЕ ОРИГИНАЛЬНОГО РЕПОЗИТОРИЯ
REPO_DIR = "MOSS-TTS-Nano"
if not os.path.exists(REPO_DIR):
    print("Клонируем оригинальный репозиторий MOSS-TTS-Nano...")
    subprocess.run(["git", "clone", "https://github.com/OpenMOSS/MOSS-TTS-Nano.git"])

# Добавляем скачанную папку в пути Python, чтобы импорты работали
sys.path.append(REPO_DIR)

# 2. ИМПОРТИРУЕМ ОРИГИНАЛЬНЫЙ ДВИЖОК
from moss_tts_nano_runtime import NanoTTSService

app = FastAPI()

# Разрешаем CORS для твоего React-интерфейса
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 3. ИНИЦИАЛИЗАЦИЯ МОДЕЛИ (скачает веса при первом запуске)
print("Загрузка весов MOSS-TTS-Nano (это может занять пару минут)...")
tts_service = NanoTTSService(
    checkpoint_path="OpenMOSS-Team/MOSS-TTS-Nano",
    audio_tokenizer_path="OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano",
    device="auto" # Автоматически включит GPU, если он есть
)
print("Модель успешно загружена и готова к работе!")

@app.get("/")
async def root():
    return {"message": "Бэкенд MOSS-TTS-Nano работает! Жду запросов от React UI на /api/predict"}

@app.post("/api/predict")
async def predict(
    text: str = Form(...),
    speed: float = Form(1.0),
    mode: str = Form("standard"),
    speaker: str = Form("default"),
    reference_audio: UploadFile = File(None)
):
    prompt_audio_path = None
    
    try:
        # Если загружен файл для клонирования голоса
        if mode == "clone" and reference_audio is not None:
            temp_audio = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
            content = await reference_audio.read()
            temp_audio.write(content)
            temp_audio.close()
            prompt_audio_path = temp_audio.name
        
        # Выбор стандартного голоса
        voice_preset = "Junhao" # Дефолтный голос в MOSS
        
        # 4. ЗАПУСК ОРИГИНАЛЬНОГО СИНТЕЗА
        result = tts_service.synthesize(
            text=text,
            voice=voice_preset if not prompt_audio_path else None,
            prompt_audio_path=prompt_audio_path,
            mode="voice_clone", # В MOSS это универсальный режим
        )
        
        audio_data = result['audio']
        sample_rate = result['sample_rate']
        
        # 5. ОБРАБОТКА СКОРОСТИ (если ползунок скорости изменен)
        if speed != 1.0:
            audio_data = librosa.effects.time_stretch(audio_data, rate=speed)
        
        # 6. СОХРАНЕНИЕ И ОТПРАВКА РЕЗУЛЬТАТА
        output_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
        sf.write(output_file.name, audio_data, sample_rate)
        
        return FileResponse(
            output_file.name, 
            media_type="audio/wav", 
            filename="output.wav"
        )
        
    except Exception as e:
        import traceback
        traceback.print_exc()
        return {"error": str(e)}
    finally:
        # Убираем за собой временные файлы
        if prompt_audio_path and os.path.exists(prompt_audio_path):
            os.remove(prompt_audio_path)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)