import gradio as gr
import torch
import torchaudio
import subprocess
import os
import sys
from huggingface_hub import snapshot_download

# Clone chunkformer nếu chưa có
if not os.path.exists("chunkformer"):
    os.system("git clone https://github.com/khanhld/chunkformer.git")

# Tải checkpoint
snapshot_download(
    repo_id="khanhld/chunkformer-large-vie",
    local_dir="./chunkformer-large-vie",
    local_dir_use_symlinks=False
)

# Thêm đường dẫn nếu cần
sys.path.append("./chunkformer")

# Giả sử bạn có model tiền xử lý custom
class DummyPreprocModel:
    def __call__(self, audio_tensor):
        return audio_tensor * 0.5

preproc_model = DummyPreprocModel()

# Thay thế bằng đường dẫn checkpoint chunkformer của bạn
MODEL_CHECKPOINT = "./chunkformer-large-vie"

def process(audio_file, transcript, option):
    waveform, sr = torchaudio.load(audio_file)

    if option == "Không xử lý":
        audio_in = audio_file  # Dùng audio gốc
    else:
        # Tiền xử lý
        processed_waveform = preproc_model(waveform)
        output_audio = "processed_audio.wav"
        torchaudio.save(output_audio, processed_waveform, sr)
        audio_in = output_audio

    # Gọi decode.py bằng subprocess
    output_txt = "asr_result.txt"

    decode_cmd = [
        sys.executable,  # Đảm bảo chạy đúng Python env
        "chunkformer/decode.py",
        "--model_checkpoint", MODEL_CHECKPOINT,
        "--long_form_audio", audio_in,
        "--total_batch_duration", "1800",
        "--chunk_size", "64",
        "--left_context_size", "128",
        "--right_context_size", "128"
    ]

    print("Running:", " ".join(decode_cmd))
    # Ghi output ra file
    with open(output_txt, "w") as f:
        subprocess.run(decode_cmd, stdout=f, stderr=subprocess.STDOUT)

    # Đọc kết quả từ file txt
    with open(output_txt) as f:
        transcript_out = f.read()

    if option == "Không xử lý":
        audio_out = None
    else:
        audio_out = audio_in

    return audio_out, transcript_out


iface = gr.Interface(
    fn=process,
    inputs=[
        gr.Audio(type="filepath", label="Tải file audio"),
        gr.Textbox(label="Transcript gốc"),
        gr.Radio(choices=["Không xử lý", "Có xử lý"], label="Chọn phương án"),
    ],
    outputs=[
        gr.Audio(type="filepath", label="File audio đã xử lý (nếu có)"),
        gr.Textbox(label="Transcript ASR"),
    ],
    title="Demo ASR & Tiền Xử Lý (Chunkformer)",
    description="Tải file audio, chọn tiền xử lý hoặc không, chạy Chunkformer."
)

if __name__ == "__main__":
    iface.launch()