import gradio as gr
import subprocess
import os
import shutil
import sys

target_paths = {
    "video": "/home/user/app/upload/source_and_edits/source.mp4",
    "image": "/home/user/app/upload/source_and_edits/ref.jpg",
    "config": "/home/user/app/upload/config/customize_train.yaml",
    "lora": "/homw/user/app/upload/lora/lora.pt",
    "output_l": "/home/user/app/outputs/train_motion_lora",
    "output_r": "/home/user/app/outputs/ref.mp4",
    "zip": "/home/user/app/outputs/train_motion_lora.zip",
}


def zip_outputs():
    if os.path.exists(target_paths["zip"]):
        os.remove(target_paths["zip"])
    shutil.make_archive(target_paths["zip"].replace(".zip", ""), 'zip', root_dir=target_paths["output_l"])
    return target_paths["zip"]

def output_video():
    if os.path.exists(target_paths["output_r"]):
        return target_paths["output_r"]
    return None 


def start_training_stream():
    process = subprocess.Popen(
        ["python", "main.py", "--config=" + target_paths["config"]],
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        bufsize=1,
        universal_newlines=True
    )
    
    output = []
    for line in process.stdout:
        output.append(line)
        yield "".join(output)

def install_i2vedit():
    try:
        import i2vedit
        print("i2vedit already installed")
    except ImportError:
        print("Installing i2vedit...")
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-e", "./i2vedit"])
        print("i2vedit installed")


def install_package(package_name):
    try:
        result = subprocess.run(
            [sys.executable, "-m", "pip", "install", package_name],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )
        output = result.stdout + "\n" + result.stderr
        return output
    except Exception as e:
        return f"Error: {str(e)}"
    

def show_package(pkg_name):
    try:
        result = subprocess.run(
            [sys.executable, "-m", "pip", "show", pkg_name],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )
        return result.stdout if result.stdout else result.stderr
    except Exception as e:
        return str(e)
    

def uninstall_package(package_name):
    try:
        result = subprocess.run(
            [sys.executable, "-m", "pip", "uninstall", package_name, "-y"],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )
        output = result.stdout + "\n" + result.stderr
        return output
    except Exception as e:
        return f"Error: {str(e)}"



def save_files(video_file, image_file, config_file, lora_file=None):
    os.makedirs(os.path.dirname(target_paths["video"]), exist_ok=True)
    os.makedirs(os.path.dirname(target_paths["config"]), exist_ok=True)

    shutil.copy(video_file.name, target_paths["video"])
    shutil.copy(image_file.name, target_paths["image"])
    shutil.copy(config_file.name, target_paths["config"])
    if lora_file:
        os.makedirs(os.path.dirname(target_paths["lora"]), exist_ok=True)
        shutil.copy(lora_file.name, target_paths["lora"])
    return "檔案已成功上傳並儲存！"


install_i2vedit()
install_package("huggingface_hub==0.25.1")
install_package("diffusers==0.25.1")
install_package("gradio==5.0.0")
uninstall_package("datasets")
print("package version set complete")


with gr.Blocks(theme=gr.themes.Origin()) as demo:
    gr.Markdown("## 請先上傳檔案")
    with gr.Row():
        video_input = gr.File(label="原始影片", file_types=[".mp4"])
        image_input = gr.File(label="編輯圖像", file_types=[".jpg", ".jpeg", ".png"])
        config_input = gr.File(label="Config 檔", file_types=[".yaml", ".yml"])
        lora_input = gr.File(label="LoRA 檔案", file_types=[".pt"])

    upload_button = gr.Button("上傳並儲存")
    output = gr.Textbox(label="狀態")


    gr.Markdown("## Training")
    with gr.Column():
        log_output = gr.Textbox(label="Training Log", lines=20)
        train_btn = gr.Button("Start Training")
 
    gr.Markdown("## Pip Installer")
    with gr.Column():
        with gr.Row():
            pkg_input = gr.Textbox(lines=1, placeholder="輸入想安裝的套件名稱，例如 diffusers 或 numpy==1.2.0")
            install_output = gr.Textbox(label="Install Output", lines=10)
        install_btn = gr.Button("Install Package")
    
    gr.Markdown("## Pip Uninstaller")
    with gr.Column():
        with gr.Row():
            pkg_input2 = gr.Textbox(lines=1, placeholder="輸入想解除安裝的套件名稱，例如 diffusers 或 numpy")
            uninstall_output = gr.Textbox(label="Uninstall Output", lines=10)
        uninstall_btn = gr.Button("Uninstall Package")
    
    gr.Markdown("## Pip show")
    with gr.Column():
        with gr.Row():
            show_input = gr.Textbox(label="輸入套件名稱（如 diffusers）")
            show_output = gr.Textbox(label="套件資訊", lines=10)
        show_btn = gr.Button("pip show")

    gr.Markdown("## Download lora")
    with gr.Column():
        file_output = gr.File(label="點擊下載", interactive=True)
        download_btn = gr.Button("下載lora")
    
    gr.Markdown("## Download results")
    with gr.Column():
        file_output2 = gr.File(label="點擊下載", interactive=True)
        download_btn2 = gr.Button("下載結果")

    show_btn.click(fn=show_package, inputs=show_input, outputs=show_output)
    download_btn.click(fn=zip_outputs, outputs=file_output)
    download_btn2.click(fn=output_video, outputs=file_output2)
    install_btn.click(fn=install_package, inputs=pkg_input, outputs=install_output)
    train_btn.click(fn=start_training_stream, outputs=log_output)
    uninstall_btn.click(fn=uninstall_package, inputs=pkg_input2, outputs=uninstall_output)
    upload_button.click(fn=save_files,inputs=[video_input, image_input, config_input, lora_input],outputs=output)
demo.launch()