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

# --- 1. CẤU HÌNH FILE ---
# Tên file model và config phải khớp chính xác với file bạn đã upload
MODEL_PATH = "model.onnx"
CONFIG_PATH = "config.json" 

# --- 2. LOAD MODEL & CONFIG ---
try:
    with open(CONFIG_PATH, "r", encoding="utf-8") as f:
        config = json.load(f)
    symbols = config.get("symbols", [])
    symbol_to_id = {s: i for i, s in enumerate(symbols)}
    sample_rate = config.get("audio", {}).get("sample_rate", 22050)
    print("✅ Đã tải Config và Symbols thành công.")

    sess = ort.InferenceSession(MODEL_PATH, providers=['CPUExecutionProvider'])
    print("✅ Đã tải Model ONNX thành công.")
except Exception as e:
    print(f"❌ Lỗi khởi tạo: {e}")

# --- 3. HÀM XỬ LÝ VĂN BẢN ---
def clean_text(text):
    # Loại bỏ thẻ HTML trong SRT (<i>, <b>...)
    text = re.sub(r'<[^>]*>', '', text)
    # Loại bỏ xuống dòng và khoảng trắng thừa
    text = text.replace('\n', ' ').strip()
    # Chuẩn hóa tiếng Việt (số, ngày tháng...)
    return TTSnorm(text)

def text_to_ids(text):
    text = clean_text(text).lower()
    ids = []
    for char in text:
        if char in symbol_to_id:
            ids.append(symbol_to_id[char])
        elif f" {char} " in symbol_to_id:
            ids.append(symbol_to_id[f" {char} "])
    return ids if ids else [1]

# --- 4. HÀM CHUYỂN SRT SANG AUDIO ---
def process_srt(srt_file, speed, noise, noise_w):
    try:
        if srt_file is None:
            return "Vui lòng chọn file SRT!", None
        
        # Đọc file phụ đề
        subs = pysrt.open(srt_file.name, encoding='utf-8')
        combined_audio = []
        
        for sub in subs:
            ids = text_to_ids(sub.text)
            input_ids = np.array([ids], dtype=np.int64)
            input_lens = np.array([len(ids)], dtype=np.int64)
            
            # Tốc độ mặc định 1.30 truyền vào scales
            scales = np.array([noise, noise_w, 1.0/speed], dtype=np.float32)
            
            # Chạy model ONNX
            outputs = sess.run(None, {
                "input": input_ids, 
                "input_lengths": input_lens, 
                "scales": scales
            })
            
            # Làm phẳng mảng 4D thành 1D (Sửa lỗi im lặng)
            audio_segment = outputs[0].flatten().astype(np.float32)
            combined_audio.append(audio_segment)
            
            # Thêm 0.2s im lặng giữa các câu cho tự nhiên
            silence = np.zeros(int(sample_rate * 0.2), dtype=np.float32)
            combined_audio.append(silence)

        # Ghép tất cả các đoạn audio
        final_audio = np.concatenate(combined_audio)
        
        # Xuất file wav
        output_path = "output_srt_voice.wav"
        wavfile.write(output_path, sample_rate, final_audio)
        
        return "Hoàn thành! Bạn có thể nghe hoặc tải audio bên dưới.", output_path
    except Exception as e:
        return f"Lỗi xử lý: {str(e)}", None

# --- 5. GIAO DIỆN GRADIO ---
with gr.Blocks(title="SRT to Audio VN") as demo:
    gr.Markdown("# 📂 Chuyển đổi SRT sang Audio Tiếng Việt")
    gr.Markdown("Hỗ trợ làm sạch văn bản và đặt tốc độ nói tùy chỉnh.")
    
    with gr.Row():
        with gr.Column():
            # Thành phần cho phép chọn file từ SD Card trên Android
            file_input = gr.File(label="Chọn file .srt từ thiết bị", file_types=[".srt"])
            
            with gr.Row():
                speed_sld = gr.Slider(0.5, 2.0, value=1.30, label="Tốc độ nói (Speed)")
                noise_sld = gr.Slider(0.1, 1.0, value=0.667, label="Noise Scale")
            
            noise_w_sld = gr.Slider(0.1, 1.0, value=0.8, label="Noise W")
            btn = gr.Button("🚀 Bắt đầu chuyển đổi", variant="primary")
            
        with gr.Column():
            status = gr.Textbox(label="Trạng thái hệ thống")
            audio_out = gr.Audio(label="Audio hoàn chỉnh", type="filepath")

    btn.click(
        fn=process_srt, 
        inputs=[file_input, speed_sld, noise_sld, noise_w_sld], 
        outputs=[status, audio_out]
    )

if __name__ == "__main__":
    # server_name="0.0.0.0" là bắt buộc để chạy trên Hugging Face
    demo.launch(server_name="0.0.0.0")