import gradio as gr
import subprocess
import os 
import shutil
import tempfile

from huggingface_hub import snapshot_download 

# Create xcodec_mini_infer folder
folder_path = './inference/xcodec_mini_infer'

# Create the folder if it doesn't exist
if not os.path.exists(folder_path):
    os.mkdir(folder_path)
    print(f"Folder created at: {folder_path}")
else:
    print(f"Folder already exists at: {folder_path}")

snapshot_download(
    repo_id = "m-a-p/xcodec_mini_infer",
    local_dir = "./inference/xcodec_mini_infer"
)

# Install required package
def install_flash_attn():
    try:
        print("Installing flash-attn...")
        subprocess.run(
            ["pip", "install", "flash-attn", "--no-build-isolation"], 
            check=True
        )
        print("flash-attn installed successfully!")
    except subprocess.CalledProcessError as e:
        print(f"Failed to install flash-attn: {e}")
        exit(1)

# Change to the "inference" directory
inference_dir = "./inference"
try:
    os.chdir(inference_dir)
    print(f"Changed working directory to: {os.getcwd()}")
except FileNotFoundError:
    print(f"Directory not found: {inference_dir}")
    exit(1)

# Function to create a temporary file with string content
def create_temp_file(content, prefix, suffix=".txt"):
    temp_file = tempfile.NamedTemporaryFile(delete=False, mode="w", prefix=prefix, suffix=suffix)
    temp_file.write(content)
    temp_file.close()
    return temp_file.name

def infer(genre_txt_content, lyrics_txt_content):
    # Create temporary files
    genre_txt_path = create_temp_file(genre_txt_content, prefix="genre_")
    lyrics_txt_path = create_temp_file(lyrics_txt_content, prefix="lyrics_")

    # Ensure the output folder exists
    output_dir = "./output"
    os.makedirs(output_dir, exist_ok=True)
    print(f"Output folder ensured at: {output_dir}")

    
    # Command and arguments
    command = [
        "python", "infer.py",
        "--stage1_model", "m-a-p/YuE-s1-7B-anneal-en-cot",
        "--stage2_model", "m-a-p/YuE-s2-1B-general",
        "--genre_txt", f"{genre_txt_path}",
        "--lyrics_txt", f"{lyrics_txt_path}",
        "--run_n_segments", "2",
        "--stage2_batch_size", "4",
        "--output_dir", f"{output_dir}",
        "--cuda_idx", "0",
        "--max_new_tokens", "3000"
    ]
    
    # Execute the command
    try:
        subprocess.run(command, check=True)
        print("Command executed successfully!")
        
        # Check and print the contents of the output folder
        output_files = os.listdir(output_dir)
        if output_files:
            print("Output folder contents:")
            for file in output_files:
                print(f"- {file}")
        else:
            print("Output folder is empty.")
        return None
    except subprocess.CalledProcessError as e:
        print(f"Error occurred: {e}")
        return None
    finally:
        # Clean up temporary files
        os.remove(genre_txt_path)
        os.remove(lyrics_txt_path)
        print("Temporary files deleted.")

# Gradio 

with gr.Blocks() as demo:
    with gr.Column():
        gr.Markdown("# YuE")
        with gr.Row():
            with gr.Column():
                genre_txt = gr.Textbox(label="Genre")
                lyrics_txt = gr.Textbox(label="Lyrics")
                submit_btn = gr.Button("Submit")
            with gr.Column():
                music_out = gr.Audio(label="Audio Result")
    
    submit_btn.click(
        fn = infer, 
        inputs = [genre_txt, lyrics_txt],
        outputs = [music_out]
    )
demo.queue().launch(show_api=False, show_error=True)