import os
import subprocess
import sys
import fileinput
import gradio as gr

# --- 1. Setup Environment ---
# Install ffmpeg if it's not already installed by packages.txt
print("Updating apt and installing ffmpeg...")
try:
    subprocess.run(["apt-get", "update", "-y"], check=True)
    subprocess.run(["apt-get", "install", "ffmpeg", "-y"], check=True)
    print("ffmpeg installed successfully.")
except subprocess.CalledProcessError as e:
    print(f"Error installing ffmpeg: {e.stderr}")
    # Continue anyway, as it might be pre-installed

# --- 2. Clone the VibeVoice Repository ---
repo_dir = "VibeVoice"
if not os.path.exists(repo_dir):
    print("Cloning the VibeVoice repository...")
    try:
        subprocess.run(
            ["git", "clone", "https://github.com/microsoft/VibeVoice.git"],
            check=True,
            capture_output=True,
            text=True
        )
        print("Repository cloned successfully.")
    except subprocess.CalledProcessError as e:
        print(f"Error cloning repository: {e.stderr}")
        sys.exit(1)
else:
    print("Repository already exists. Skipping clone.")

# --- 3. Install the Package ---
os.chdir(repo_dir)
print(f"Changed directory to: {os.getcwd()}")
print("Installing the VibeVoice package...")
try:
    subprocess.run(
        [sys.executable, "-m", "pip", "install", "-e", "."],
        check=True,
        capture_output=True,
        text=True
    )
    print("Package installed successfully.")
except subprocess.CalledProcessError as e:
    print(f"Error installing package: {e.stderr}")
    sys.exit(1)


# --- 4. Modify the demo script for CPU execution ---
demo_script_path = "demo/gradio_demo.py"
print(f"Modifying {demo_script_path} for CPU execution...")

# Use fileinput to perform in-place replacement
with fileinput.FileInput(demo_script_path, inplace=True) as file:
    for line in file:
        # Change the model loading line to be CPU-compatible
        if 'VibeVoiceForConditionalGenerationInference.from_pretrained(' in line:
            print('        self.model = VibeVoiceForConditionalGenerationInference.from_pretrained(')
            print('            self.model_path,')
            print('            torch_dtype=torch.float32,  # Use float32 for CPU')
            print('            device_map="cpu",')
            print('        )')
            # Skip the next few lines of the original arguments
            next(file, None)
            next(file, None)
            next(file, None)
        else:
            print(line, end='')

print("Script modified successfully.")


# --- 5. Launch the Gradio Demo ---
model_id = "microsoft/VibeVoice-1.5B"

# Construct the command as specified in the README
command = [
    "python",
    demo_script_path,
    "--model_path",
    model_id,
    "--share"
]

print(f"Launching Gradio demo with command: {' '.join(command)}")
# This command will start the Gradio server
subprocess.run(command)