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 CẤU HÌNH ---
try:
    with open("config.json", "r", encoding="utf-8") as f:
        config = json.load(f)
    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 đồng bộ thời gian đã sẵn sàng")
except Exception as e:
    print(f"❌ Lỗi khởi động: {e}")

# --- 2. XỬ LÝ VĂN BẢN ---
def process_text(text):
    # Làm sạch HTML và chuẩn hóa tiếng Việt
    text = re.sub(r'<[^>]*>', '', text)
    text = TTSnorm(text.replace('\n', ' ').strip())
    # Chuyển sang âm vị (IPA) để model Piper có thể đọc
    phonemes = phonemize(text, language='vi', backend='espeak', strip=True)
    
    ids = [phoneme_id_map.get("^", [1])[0]]
    for p in phonemes:
        if p in phoneme_id_map:
            ids.extend(phoneme_id_map[p])
            ids.append(phoneme_id_map.get("_", [0])[0])
    ids.append(phoneme_id_map.get("$", [2])[0])
    return ids

# --- 3. LOGIC ĐỒNG BỘ THỜI GIAN (SYNC) ---
def srt_to_audio_sync(srt_file, speed=1.30, noise=0.667, noise_w=0.8):
    if srt_file is None: return "Chưa chọn file SRT", None
    try:
        subs = pysrt.open(srt_file.name, encoding='utf-8')
        final_audio_stream = []
        current_sample_index = 0

        for sub in subs:
            # 1. Tính toán thời điểm bắt đầu tính bằng Sample
            start_time_seconds = sub.start.ordinal / 1000.0
            start_sample_target = int(start_time_seconds * sample_rate)

            # 2. Nếu thời điểm hiện tại chưa tới thời điểm bắt đầu, chèn khoảng lặng
            if start_sample_target > current_sample_index:
                silence_len = start_sample_target - current_sample_index
                final_audio_stream.append(np.zeros(silence_len, dtype=np.float32))
                current_sample_index = start_sample_target

            # 3. Tạo âm thanh cho câu phụ đề
            ids = process_text(sub.text)
            input_ids = np.array([ids], dtype=np.int64)
            input_lens = np.array([len(ids)], dtype=np.int64)
            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
            audio_segment = outputs[0].flatten().astype(np.float32)
            
            # 4. Thêm audio vào luồng và cập nhật vị trí hiện tại
            final_audio_stream.append(audio_segment)
            current_sample_index += len(audio_segment)

        # Kết hợp tất cả thành một file duy nhất
        full_audio = np.concatenate(final_audio_stream)
        wavfile.write("synced_dubbing.wav", sample_rate, full_audio)
        
        return "Đã tạo file dubbing khớp thời gian thành công!", "synced_dubbing.wav"
    except Exception as e:
        return f"Lỗi xử lý: {str(e)}", None

# --- 4. GIAO DIỆN ---
demo = gr.Interface(
    fn=srt_to_audio_sync,
    inputs=[
        gr.File(label="Chọn file SRT từ SD Card", file_types=[".srt"]),
        gr.Slider(0.5, 2.0, value=1.30, label="Tốc độ nói"),
        gr.Slider(0.1, 1.0, value=0.667, label="Noise Scale"),
        gr.Slider(0.1, 1.0, value=0.8, label="Noise W")
    ],
    outputs=[gr.Textbox(label="Trạng thái"), gr.Audio(label="Audio khớp chuẩn thời gian")],
    title="SRT Dubbing - Đồng bộ thời gian chuẩn",
    description="Tự động thêm khoảng lặng để âm thanh khớp hoàn toàn với mốc thời gian trong file SRT."
)

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