import os
import yaml
import torch
import numpy as np
import soundfile as sf
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from typing import Optional
import io

# --- Tận dụng mã nguồn gốc VieNeu-TTS ---
# Đảm bảo các thư mục vieneu_tts và utils nằm ở thư mục gốc để import thành công
from vieneu_tts import VieNeuTTS
from utils.core_utils import split_text_into_chunks

app = FastAPI(title="VieNeu-TTS API Server")

# Kích hoạt CORS để index.html trên máy bạn có thể gọi API này
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# --- Khởi tạo Model & Cấu hình ---
CONFIG_PATH = "config.yaml"
if not os.path.exists(CONFIG_PATH):
    raise FileNotFoundError("❌ Không tìm thấy file config.yaml ở thư mục gốc!")

with open(CONFIG_PATH, "r", encoding="utf-8") as f:
    config = yaml.safe_load(f)

# Lấy đường dẫn model GGUF từ config
# Lưu ý: config.yaml của bạn phải có: backbone_configs -> VieNeu-TTS-GGUF -> repo
try:
    model_path = config['backbone_configs']['VieNeu-TTS-GGUF']['repo']
except KeyError:
    raise KeyError("❌ config.yaml thiếu cấu hình 'backbone_configs/VieNeu-TTS-GGUF/repo'")

print(f"🔄 Đang nạp model từ: {model_path} (Chế độ CPU)...")

# Khởi tạo động cơ VieNeu-TTS
# Ép device về "cpu" vì Hugging Face Space miễn phí không có GPU
tts = VieNeuTTS(backbone_repo=model_path, backbone_device="cpu")
print("✅ VieNeu-TTS đã sẵn sàng phục vụ!")

# --- Định nghĩa Schema cho API ---
class TTSRequest(BaseModel):
    input: str
    voice: Optional[str] = "Vĩnh (nam miền Nam)"
    speed: Optional[float] = 1.0

# --- Các cổng kết nối (Endpoints) ---

@app.get("/")
def health_check():
    """Kiểm tra xem API có đang sống không"""
    return {
        "status": "online",
        "model": model_path,
        "device": "cpu",
        "message": "VieNeu-TTS API is running on Hugging Face Docker"
    }

@app.get("/v1/voices")
def get_voices():
    """Trả về danh sách giọng đọc từ config.yaml cho index.html chọn"""
    voices = list(config.get('voice_samples', {}).keys())
    return {"data": [{"id": v} for v in voices]}

@app.post("/v1/audio/speech")
async def text_to_speech(request: TTSRequest):
    """Cổng chính: Nhận text -> Trả về file âm thanh .wav"""
    try:
        # 1. Lấy thông tin giọng đọc mẫu
        voice_info = config['voice_samples'].get(request.voice)
        if not voice_info:
            raise HTTPException(status_code=400, detail=f"Giọng '{request.voice}' không tồn tại trong config")

        ref_audio_path = voice_info['audio']
        ref_text_path = voice_info['text']

        if not os.path.exists(ref_audio_path):
            raise HTTPException(status_code=404, detail=f"Không tìm thấy file audio mẫu: {ref_audio_path}")

        # Đọc nội dung text của giọng mẫu
        with open(ref_text_path, "r", encoding="utf-8") as f:
            ref_text_content = f.read().strip()

        # 2. Xử lý văn bản đầu vào (Chia nhỏ đoạn văn dài bằng hàm của VieNeu-TTS)
        chunks = split_text_into_chunks(request.input, max_chars=256)
        
        # 3. Mã hóa giọng đọc mẫu (Encoding reference)
        ref_codes = tts.encode_reference(ref_audio_path)
        
        # 4. Thực hiện chuyển đổi (Inference) cho từng đoạn và nối lại
        final_audio_segments = []
        for chunk in chunks:
            if not chunk.strip(): continue
            # tts.infer trả về numpy array
            wav = tts.infer(chunk, ref_codes, ref_text_content)
            if wav is not None:
                final_audio_segments.append(wav)
        
        if not final_audio_segments:
            raise HTTPException(status_code=500, detail="Không thể tạo ra âm thanh từ văn bản này")

        # Nối tất cả các đoạn audio lại thành một
        combined_wav = np.concatenate(final_audio_segments)

        # 5. Đóng gói vào file WAV và gửi về trình duyệt
        buffer = io.BytesIO()
        sf.write(buffer, combined_wav, 24000, format='WAV')
        buffer.seek(0)

        return StreamingResponse(buffer, media_type="audio/wav")

    except Exception as e:
        print(f"❌ Lỗi xử lý TTS: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Internal Server Error: {str(e)}")

# --- Chạy Server ---
if __name__ == "__main__":
    import uvicorn
    # Hugging Face Docker BẮT BUỘC dùng port 7860
    uvicorn.run(app, host="0.0.0.0", port=7860)