import os
import sys
import argparse
import torch
import numpy as np
import gradio as gr

# Setup paths for indextts
current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(current_dir)
sys.path.append(os.path.join(current_dir, "indextts"))

# Setup env variables to avoid clutter
hf_cache_dir = os.path.join(current_dir, "checkpoints", "hf_cache")
torch_cache_dir = os.path.join(current_dir, "checkpoints", "torch_cache")
os.environ.setdefault("INDEXTTS_USE_DEEPSPEED", "0")
os.environ.setdefault("HF_HOME", hf_cache_dir)
os.environ.setdefault("HF_HUB_CACHE", hf_cache_dir)
os.environ.setdefault("TRANSFORMERS_CACHE", hf_cache_dir)
os.environ.setdefault("TORCH_HOME", torch_cache_dir)
os.makedirs(hf_cache_dir, exist_ok=True)
os.makedirs(torch_cache_dir, exist_ok=True)

from indextts.infer_v2_thai import IndexTTS2

MODEL_DIR = os.path.join(current_dir, "checkpoints")
GPT_PATH = os.path.join(current_dir, "models", "thaiseperate2.pth")
BPE_PATH = os.path.join(MODEL_DIR, "thai_segmented_bpe.model")

print(">> Loading IndexTTS2 for Native Streaming Test...")
tts = IndexTTS2(
    model_dir=MODEL_DIR,
    cfg_path=os.path.join(MODEL_DIR, "config.yaml"),
    is_fp16=False,
    use_cuda_kernel=False,
    use_accel=True,
    use_torch_compile=False,
    gpt_checkpoint_path=GPT_PATH,
    bpe_model_path=BPE_PATH
)
print(">> Models loaded successfully!")

def generate_stream(text, prompt_path):
    print(f"\n[STREAM] Starting generation for text length {len(text)}")
    print(f"[STREAM] Prompt: {prompt_path}")
    
    # Yield None first to clear the audio player from previous generations
    yield None
    
    # Initialize generator
    generator = tts.infer_generator(
        spk_audio_prompt=prompt_path,
        text=text,
        output_path=None,
        stream_return=True,
        do_sample=True,
        top_p=0.8,
        temperature=0.8
    )
    
    sr = 22050
    for i, chunk in enumerate(generator):
        audio_tensor = chunk
        if audio_tensor is None:
            continue
        
        # Audio is returned as a PyTorch tensor, e.g., shape [1, N].
        # Convert to a 1D numpy float array for Gradio.
        audio_np = audio_tensor.squeeze().cpu().numpy()
        dur = len(audio_np) / sr
        
        print(f"[STREAM] Yielding chunk {i+1} (duration: {dur:.2f}s)")
        
        # Gradio 4+ native streaming format: tuple(sample_rate, numpy_array)
        yield (sr, audio_np)
        
    print("[STREAM] Generation finished!")

DEFAULT_TEXT = "ไม่ค่อย เห็น ใคร ทำ ข่าว หรือ พูดถึง เรื่อง นี้ แต่ ตอนนี้ กำลัง มี กระแส แบน ผลิตภัณฑ์ ที่ เป็น ของ เอเชีย หรือ เกี่ยวกับ เอเชีย ซึ่ง มัน แรง มาก จุด กระแส เหยียด เชื้อชาติ เอเชีย แรง ขึ้น อีก บาง ร้าน ปิดกิจการ ไป แล้ว ใน ไทย ก็ มี คน บอยคอต แต่ ประชากร พี่ แก น้อย เลย ไม่ ได้รับ ผลกระทบ ต้นเหตุ เกิด จา กร้าน คนจีน คน นึง มี คน ดำ มา ใช้ บริการ แล้ว คิด ว่า คน นั้น เป็น ขโมย เลย ยิง สรุป ขิต แล้วก็ มา รู้ ทีหลัง ว่าไม่ได้ ขโมย อะไร คิด ว่า พวก พี่ พี่ น่าจะ อยาก แบน หมด แต่ มัน ทำ ไม่ ได้ เลย กลายเป็น แบน แค่ ธุรกิจ เล็ก เล็ก แทน จุด กระแส แบน เอเชีย"
DEFAULT_PROMPT = r"C:\Users\User\Downloads\audio (51).wav"

with gr.Blocks(title="Native Stream Test") as demo:
    gr.Markdown("<h1 style='text-align: center'>🚀 IndexTTS Native Streaming Test</h1>")
    gr.Markdown("<p style='text-align: center'>สคริปต์นี้ถูกเขียนแยกมาเพื่อใช้ <b>Gradio Native Audio Streaming</b> แบบไม่ต้องง้อ Javascript แล้วครับ</p>")
    
    with gr.Row():
        input_text = gr.Textbox(value=DEFAULT_TEXT, lines=8, label="Input Text")
    
    with gr.Row():
        prompt_audio = gr.Textbox(value=DEFAULT_PROMPT, label="Prompt Audio Path (.wav)")
        
    with gr.Row():
        generate_btn = gr.Button("▶ Generate Audio Stream", variant="primary", size="lg")
        
    with gr.Row():
        # streaming=True allows Gradio to playback chunks immediately as they are yielded
        output_audio = gr.Audio(label="Streaming Output (Plays instantly)", autoplay=True, streaming=True)
        
    generate_btn.click(
        fn=generate_stream,
        inputs=[input_text, prompt_audio],
        outputs=output_audio,
        show_progress=True
    )

if __name__ == "__main__":
    demo.queue()
    print(">> Launching server on http://127.0.0.1:7865")
    demo.launch(server_name="127.0.0.1", server_port=7865)