app.py

import os
import time
import random
import base64
import tempfile
import requests
import numpy as np
import soundfile as sf
import pysrt
import librosa
import pyrubberband as pyrb
import concurrent.futures
import gradio as gr
import edge_tts
import asyncio

# --- MAP ID GIỌNG ĐỌC TIKTOK ---
tiktok_voices = {
    "🇻🇳 Nữ (Việt Nam) - BV074": "BV074_streaming",
    "🇻🇳 Nam (Việt Nam) - BV075": "BV075_streaming",
    "🇺🇸 Nữ US (Jessie)": "en_us_002",
    "🇺🇸 Nam US (Joey)": "en_us_006",
    "🎭 Ghost Face (Kinh dị)": "en_us_ghostface"
}

# --- MAP ID GIỌNG ĐỌC MICROSOFT EDGE (MIỄN PHÍ) ---
edge_voices = {
    "🇻🇳 Nữ (Hoài My - Chuẩn, Nhẹ nhàng)": "vi-VN-HoaiMyNeural",
    "🇻🇳 Nam (Nam Minh - Chuẩn, Trầm ấm)": "vi-VN-NamMinhNeural",
    "🇺🇸 Nữ US (Aria - Sang trọng)": "en-US-AriaNeural",
    "🇺🇸 Nam US (Christopher - Lôi cuốn)": "en-US-ChristopherNeural",
    "🇬🇧 Nam Anh Quốc (Ryan)": "en-GB-RyanNeural"
}

# --- CÁC HÀM XỬ LÝ DSP TỐI ƯU HÓA ---
def safe_normalize(audio_array, target_peak=0.90):
    audio_array = np.array(audio_array, dtype=np.float32)
    if np.max(np.abs(audio_array)) > 1.0:
        audio_array = audio_array / 32768.0
    audio_array = audio_array - np.mean(audio_array)
    max_amp = np.max(np.abs(audio_array))
    if max_amp > 0:
         audio_array = (audio_array / max_amp) * target_peak
    return audio_array

def apply_fade(audio_array, sr, fade_duration=0.015):
    fade_samples = int(sr * fade_duration)
    if len(audio_array) < fade_samples * 2:
        return audio_array
    fade_in = np.linspace(0.0, 1.0, fade_samples, dtype=np.float32)
    fade_out = np.linspace(1.0, 0.0, fade_samples, dtype=np.float32)
    audio_array[:fade_samples] *= fade_in
    audio_array[-fade_samples:] *= fade_out
    return audio_array

# --- GỌI API TIKTOK ---
def get_tiktok_audio(text, v_code, target_sr=None, max_retries=5):
    url = "https://tiktok-tts.weilnet.workers.dev/api/generation"
    payload = {"text": text, "voice": v_code}

    for attempt in range(max_retries):
        time.sleep(random.uniform(0.1, 0.8))
        try:
            response = requests.post(url, json=payload, timeout=20)
            if response.status_code in [429, 500, 502, 503, 504]:
                time.sleep((2 ** attempt) + random.uniform(0.5, 1.5))
                continue

            response.raise_for_status()
            data = response.json()

            if "data" in data:
                audio_bytes = base64.b64decode(data["data"])
                with tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as f:
                    f.write(audio_bytes)
                    tmp_path = f.name
                audio_array, sr = librosa.load(tmp_path, sr=target_sr)
                os.remove(tmp_path)
                return audio_array, sr
        except requests.exceptions.RequestException:
            time.sleep((2 ** attempt) + random.uniform(0.5, 1.5))
    return None, None

# --- GỌI API MICROSOFT EDGE CÓ CHỈNH TỐC ĐỘ GỐC ---
def get_edge_audio(text, v_code, rate_str="+0%", target_sr=None):
    with tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as f:
        tmp_path = f.name
    
    async def _generate():
        communicate = edge_tts.Communicate(text, v_code, rate=rate_str)
        await communicate.save(tmp_path)
        
    try:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        loop.run_until_complete(_generate())
        loop.close()
        
        audio_array, sr = librosa.load(tmp_path, sr=target_sr)
        os.remove(tmp_path)
        return audio_array, sr
    except Exception:
        if os.path.exists(tmp_path):
            os.remove(tmp_path)
        return None, None

# --- HÀM XỬ LÝ LÕI ĐA LUỒNG ---
def process_single_line(index, text, start_sec, end_sec, provider, v_code, native_sr, speed_m, pitch_steps):
    srt_duration = end_sec - start_sec
    if srt_duration <= 0: srt_duration = 0.5

    # 1. PHÂN LUỒNG & GỌI API (Edge tự buff thêm tốc độ)
    if provider == "Microsoft Edge (Neural Cao Cấp)":
        base_boost = 15 
        user_boost = int((speed_m - 1.0) * 100)
        total_rate = base_boost + user_boost
        rate_str = f"+{total_rate}%" if total_rate >= 0 else f"{total_rate}%"
        audio_chunk, _ = get_edge_audio(text, v_code, rate_str=rate_str, target_sr=native_sr)
    else:
        audio_chunk, _ = get_tiktok_audio(text, v_code, target_sr=native_sr)

    if audio_chunk is None:
        return index, None, None, None

    # 2. XÉN KHOẢNG LẶNG VÀ BƠM "ĐỆM KHÍ" (BUFFER) 100ms
    audio_chunk_trimmed, _ = librosa.effects.trim(audio_chunk, top_db=35)
    
    # Chỉ xử lý nếu file đủ dài (> 0.2s)
    if len(audio_chunk_trimmed) > native_sr * 0.2:
        audio_chunk = audio_chunk_trimmed
        
        # Tạo 100 mili-giây (0.1 giây) âm thanh trống tuyệt đối
        silence_samples = int(0.1 * native_sr)
        
        # Chia đôi 100ms ra: 50ms gắn vào đầu, 50ms gắn vào cuối
        half_silence = np.zeros(silence_samples // 2, dtype=np.float32)
        
        # Ghép nối an toàn
        audio_chunk = np.concatenate((half_silence, audio_chunk, half_silence))

    audio_chunk = safe_normalize(audio_chunk, target_peak=0.90)
    
    # Tính lại tổng chiều dài sau khi đã bơm 100ms đệm khí
    orig_dur = len(audio_chunk) / native_sr

    # 3. TÍNH TOÁN BÙ TRỪ THỜI GIAN THEO SRT
    if provider == "Microsoft Edge (Neural Cao Cấp)":
        # Tốc độ đã được Microsoft ép trên server, ta chỉ ép thêm nếu nó vẫn dài hơn SRT
        stretch_factor = orig_dur / srt_duration
        stretch_factor = max(1.0, stretch_factor) 
    else:
        # TikTok không ép tốc độ trên server, ta phải ép tay
        stretch_factor = max(speed_m, orig_dur / srt_duration)

    # 4. ÉP THỜI GIAN/CAO ĐỘ AN TOÀN (CÓ VÒNG BẢO VỆ CHỐNG CRASH)
    if abs(stretch_factor - 1.0) > 0.01 or pitch_steps != 0:
        try:
            processed_chunk = pyrb.time_stretch(audio_chunk, native_sr, stretch_factor)
            if pitch_steps != 0:
                processed_chunk = pyrb.pitch_shift(processed_chunk, native_sr, pitch_steps)
        except Exception as e:
            # Nếu pyrubberband sập, hoàn tác dùng lại file gốc có đệm 100ms để không bị mất câu
            processed_chunk = audio_chunk
            print(f"Bỏ qua ép tốc độ dòng {index} do file quá ngắn hoặc lỗi pyrubberband: {e}")
    else:
        processed_chunk = audio_chunk

    processed_chunk = safe_normalize(processed_chunk, target_peak=0.90)
    processed_chunk = apply_fade(processed_chunk, native_sr, fade_duration=0.015)

    start_sample = int(start_sec * native_sr)
    return index, processed_chunk, start_sample, stretch_factor

# --- HÀM GIAO TIẾP VỚI GRADIO ---
def srt_to_speech_handler(srt_file, provider, chon_giong, toc_do, cao_do, so_luong_luong, progress=gr.Progress(track_tqdm=True)):
    if srt_file is None:
        return None, "❌ Vui lòng tải lên file phụ đề .srt trước!"

    log_messages = []
    def log(msg):
        log_messages.append(msg)
        return "\n".join(log_messages)

    if provider == "Microsoft Edge (Neural Cao Cấp)":
        voice_id = edge_voices.get(chon_giong)
    else:
        voice_id = tiktok_voices.get(chon_giong)
    
    yield None, log("⏳ Đang kiểm tra kết nối với máy chủ...")
    
    if provider == "Microsoft Edge (Neural Cao Cấp)":
        _, native_sample_rate = get_edge_audio("Test", voice_id)
    else:
        _, native_sample_rate = get_tiktok_audio("Test", voice_id)

    if native_sample_rate is None:
        yield None, log("❌ Không thể kết nối. Vui lòng thử lại sau.")
        return

    yield None, log(f"✅ Kết nối thành công! Tần số gốc: {native_sample_rate} Hz\n⏳ Đang đọc file phụ đề...")
    
    try:
        subs = pysrt.open(srt_file.name)
    except Exception as e:
        yield None, log(f"❌ Lỗi đọc file SRT: {str(e)}")
        return

    if len(subs) == 0:
        yield None, log("❌ File SRT trống, không tìm thấy câu phụ đề nào.")
        return

    max_end_sec = max([sub.end.ordinal / 1000.0 for sub in subs])
    total_samples = int(max_end_sec * native_sample_rate)
    final_audio = np.zeros(total_samples, dtype=np.float32)

    yield None, log(f"🚀 Bắt đầu tổng hợp đa luồng ({so_luong_luong} luồng) cho {len(subs)} câu...")

    tasks = []
    for i, sub in enumerate(subs):
        text = sub.text.replace('\n', ' ').strip()
        if not text: continue
        if len(text) > 290: text = text[:290]
        start_sec = sub.start.ordinal / 1000.0
        end_sec = sub.end.ordinal / 1000.0
        tasks.append((i, text, start_sec, end_sec))

    count = 0
    with concurrent.futures.ThreadPoolExecutor(max_workers=so_luong_luong) as executor:
        futures = {
            executor.submit(
                process_single_line,
                task[0], task[1], task[2], task[3],
                provider, voice_id, native_sample_rate, toc_do, cao_do
            ): task for task in tasks
        }

        for future in progress.tqdm(concurrent.futures.as_completed(futures), total=len(tasks), desc="Đang xử lý phụ đề"):
            idx, processed_chunk, start_sample, speed_applied = future.result()
            count += 1

            if processed_chunk is not None:
                end_sample = start_sample + len(processed_chunk)
                if end_sample > len(final_audio):
                    processed_chunk = processed_chunk[:(len(final_audio) - start_sample)]
                    end_sample = len(final_audio)

                final_audio[start_sample:end_sample] += processed_chunk
                yield None, log(f" [{count}/{len(subs)}] ✅ Dòng {idx+1}: Thành công | Biên độ giãn: {speed_applied:.2f}x")
            else:
                yield None, log(f" [{count}/{len(subs)}] ❌ Dòng {idx+1}: Thất bại hoàn toàn.")

    yield None, log("\n⏳ Đang áp dụng bộ lọc âm thanh Limiter lần cuối...")
    final_audio = np.tanh(final_audio)
    final_audio = safe_normalize(final_audio, target_peak=0.95)

    output_path = os.path.join(tempfile.gettempdir(), f"ketqua_TTS_Pro.wav")
    sf.write(output_path, final_audio, native_sample_rate)

    yield output_path, log("🎉 TỔNG HỢP HOÀN THÀNH! Bạn có thể tải file ở bên dưới.")

# --- UI UPDATE LOGIC ---
def update_provider(provider):
    if provider == "Microsoft Edge (Neural Cao Cấp)":
        return gr.update(choices=list(edge_voices.keys()), value=list(edge_voices.keys())[0])
    else:
        return gr.update(choices=list(tiktok_voices.keys()), value=list(tiktok_voices.keys())[0])

# --- KHỞI TẠO GIAO DIỆN GRADIO ---
with gr.Blocks(title="SRT to Speech (TikTok + Edge)") as demo:
    gr.Markdown("<h1 style='text-align: center; margin-bottom: 0;'>🎙️ Hệ Thống Chuyển Đổi Phụ Đề Sang Giọng Đọc</h1>")
    gr.Markdown("<p style='text-align: center; color: gray;'>Hỗ trợ TikTok TTS và Microsoft Edge TTS (Miễn phí, Không cần Key)</p>")
    
    with gr.Row():
        with gr.Column(scale=5):
            with gr.Group():
                gr.Markdown("### 📂 1. Tải Lên Phụ Đề")
                srt_input = gr.File(label="Tải lên duy nhất 1 file (.srt)", file_types=[".srt"])
            
            with gr.Group():
                gr.Markdown("### ⚙️ 2. Chọn Dịch Vụ")
                provider_input = gr.Radio(
                    choices=["TikTok (Giọng Phổ Biến)", "Microsoft Edge (Neural Cao Cấp)"],
                    value="Microsoft Edge (Neural Cao Cấp)", 
                    label="Nền tảng xử lý"
                )

            with gr.Group():
                gr.Markdown("### 🎛️ 3. Cấu Hình Giọng Đọc")
                voice_input = gr.Dropdown(
                    choices=list(edge_voices.keys()), 
                    value=list(edge_voices.keys())[0], 
                    label="Chọn Giọng",
                    interactive=True
                )
                with gr.Row():
                    speed_input = gr.Slider(minimum=1.0, maximum=1.5, value=1.1, step=0.1, label="Tốc độ mặc định")
                    pitch_input = gr.Slider(minimum=-12, maximum=12, value=0, step=1, label="Cao độ (Pitch)")
                
                threads_input = gr.Slider(minimum=1, maximum=20, value=10, step=1, label="Số lượng luồng tải song song")
            
            btn_run = gr.Button("🚀 BẮT ĐẦU CHUYỂN ĐỔI", variant="primary", size="lg")
            
        with gr.Column(scale=5):
            with gr.Group():
                gr.Markdown("### 🎧 4. Kết Quả Âm Thanh")
                audio_output = gr.Audio(label="File Phụ Đề Hoàn Thiện (.WAV)", type="filepath")
            
            with gr.Group():
                gr.Markdown("### 📊 Trạng Thái Xử Lý")
                log_output = gr.Textbox(label="Nhật ký (Logs)", lines=14, max_lines=14, interactive=False)

    provider_input.change(
        fn=update_provider,
        inputs=[provider_input],
        outputs=[voice_input]
    )

    btn_run.click(
        fn=srt_to_speech_handler,
        inputs=[srt_input, provider_input, voice_input, speed_input, pitch_input, threads_input],
        outputs=[audio_output, log_output]
    )

if __name__ == "__main__":
    demo.queue().launch(theme=gr.themes.Soft())