import gradio as gr
import onnxruntime as ort
import numpy as np
import json
import pysrt
import re
from vinorm import TTSnorm
from phonemizer import phonemize
from scipy.io import wavfile

# --- 1. TẢI MODEL & CONFIG ---
try:
    with open("config.json", "r", encoding="utf-8") as f:
        config = json.load(f)
    # Lấy bản đồ ID âm vị từ config của bạn
    phoneme_id_map = config["phoneme_id_map"]
    sample_rate = config.get("audio", {}).get("sample_rate", 22050)
    
    sess = ort.InferenceSession("model.onnx", providers=['CPUExecutionProvider'])
    print("✅ Hệ thống đã sẵn sàng")
except Exception as e:
    print(f"❌ Lỗi khởi động: {e}")

# --- 2. LÀM SẠCH VĂN BẢN & CHUYỂN ÂM VỊ ---
def process_text(text):
    # Loại bỏ thẻ HTML <i>, <b> thường có trong SRT
    text = re.sub(r'<[^>]*>', '', text)
    # Chuẩn hóa tiếng Việt (100k -> một trăm nghìn)
    text = TTSnorm(text.replace('\n', ' ').strip())
    # Chuyển sang âm vị IPA (Piper yêu cầu)
    phonemes = phonemize(text, language='vi', backend='espeak', strip=True)
    
    ids = [phoneme_id_map.get("^", [1])[0]] # Bắt đầu
    for p in phonemes:
        if p in phoneme_id_map:
            ids.extend(phoneme_id_map[p])
            ids.append(phoneme_id_map.get("_", [0])[0]) # Ký tự trống
    ids.append(phoneme_id_map.get("$", [2])[0]) # Kết thúc
    return ids

# --- 3. XỬ LÝ SRT TO AUDIO ---
def srt_to_audio(srt_file, speed=1.30, noise=0.667, noise_w=0.8):
    if srt_file is None: return "Chưa có file!", None
    try:
        subs = pysrt.open(srt_file.name, encoding='utf-8')
        combined_audio = []
        
        for sub in subs:
            ids = process_text(sub.text)
            input_ids = np.array([ids], dtype=np.int64)
            input_lens = np.array([len(ids)], dtype=np.int64)
            # Tốc độ 1.30 được đưa vào tham số scale cuối cùng
            scales = np.array([noise, noise_w, 1.0/speed], dtype=np.float32)
            
            outputs = sess.run(None, {
                "input": input_ids, "input_lengths": input_lens, "scales": scales
            })
            
            # Làm phẳng mảng 4D thành 1D âm thanh
            combined_audio.append(outputs[0].flatten())
            # Nghỉ 0.2s giữa các câu
            combined_audio.append(np.zeros(int(sample_rate * 0.2), dtype=np.float32))

        final_wave = np.concatenate(combined_audio).astype(np.float32)
        wavfile.write("output.wav", sample_rate, final_wave)
        return "Thành công!", "output.wav"
    except Exception as e:
        return f"Lỗi: {str(e)}", None

# --- 4. GIAO DIỆN ---
demo = gr.Interface(
    fn=srt_to_audio,
    inputs=[
        gr.File(label="Tải file SRT từ điện thoại/SDCard"),
        gr.Slider(0.5, 2.0, value=1.30, label="Tốc độ nói"),
        gr.Slider(0.1, 1.0, value=0.667, label="Độ nhiễu (Noise)"),
        gr.Slider(0.1, 1.0, value=0.8, label="Độ nhiễu W")
    ],
    outputs=[gr.Textbox(label="Trạng thái"), gr.Audio(label="Kết quả lồng tiếng")],
    title="SRT to Audio Tiếng Việt - Tốc độ 1.30"
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0")