import streamlit as st
import subprocess
import shutil
from io import BytesIO
import os


# Create 'temp' directory if it doesn't exist
os.makedirs('temp', exist_ok=True)



def run_scripts(target, source, mode):
    if target is None or (source is None):
        return None
    with st.spinner("Processing..."):
        target_extension = os.path.splitext(target.name)[-1]
        output_path = "output" + target_extension

        target_bytes = target.read()  # Read target file as bytes
        source_bytes = source.read()  # Read source file as bytes

        target_io = BytesIO(target_bytes)  # Convert bytes to BytesIO
        source_io = BytesIO(source_bytes)  # Convert bytes to BytesIO

        # Save files to temporary directory
        with open(f'temp/target{target_extension}', 'wb') as f:
            f.write(target_bytes)
        with open(f'temp/source{target_extension}', 'wb') as f:
            f.write(source_bytes)

        if mode in ["Face Swapper", "Both"]:
            cmd1 = ["python3", "run.py", "-s", f'temp/source{target_extension}', "-t", f'temp/target{target_extension}', "-o", output_path, "--frame-processor", "face_swapper"]
            subprocess.run(cmd1)

        if mode in ["Face Enhancer", "Both"]:
            cmd2 = ["python3", "run.py", "-t", f'temp/target{target_extension}', "-o", output_path, "--frame-processor", "face_enhancer"]
            subprocess.run(cmd2)

        os.remove(f'temp/source{target_extension}')
        os.remove(f'temp/target{target_extension}')

        return output_path

# Streamlit UI
st.markdown('<p style="color:#191970;text-align:center;font-size:30px;">Aiconvert.online</p>', unsafe_allow_html=True)
st.title("AIconvert Face swapper / Face Enhancer")
st.markdown('<style>h1{color: Crimson; text-align: center;}</style>', unsafe_allow_html=True)

source_file = st.file_uploader("Upload source image")
target_file = st.file_uploader("Upload target image")
mode = st.radio("Choose Mode", ("Face Swapper", "Face Enhancer"))

if source_file and target_file and st.button("Run"):
    result = run_scripts(target_file, source_file, mode)
    if result:
        st.image(result)  # Display the result as an image