import streamlit as st
import streamlit.components.v1 as components
import subprocess
import os
import glob

# Function to run the model and generate HTML
def generate_html(text_input, length):
    command = [
        "python", "generate.py",
        "--resume-pth", "output/vq/2023-07-19-04-17-17_12_VQVAE_20batchResetNRandom_8192_32/net_last.pth",
        "--resume-trans", "output/t2m/2023-10-10-03-17-01_HML3D_44_crsAtt2lyr_mask0.5-1/net_last.pth",
        "--text", text_input,
        "--length", str(length)
    ]
    try:
        # Capture the standard error output
        result = subprocess.run(command, check=True, text=True, capture_output=True)
        return find_latest_html_file('output')
    except subprocess.CalledProcessError as e:
        print("Error:", e.stderr)  # Print the error message from stderr
        return None

# Function to find the latest HTML file in the specified directory
def find_latest_html_file(base_path):
    # List all HTML files in the output directory and find the latest one
    list_of_files = glob.glob(f'{base_path}/*.html')
    if not list_of_files:
        return None
    latest_file = max(list_of_files, key=os.path.getctime)
    return latest_file

# Streamlit app
components.html("""
    <h1 style='text-align: center; color: white;'>MMM Model Demo</h1>
    """,
    height=100)

# Initialize session state variable for text input
text_input = st.text_area("Enter text here:", value="", key="text_input")
length = st.number_input("Length of the generated text:", value=156, key="length")

# Button trigger to generate HTML
if st.button("Generate HTML"):
    # Check if the input is not empty
    if text_input and length:
        html_file_path = generate_html(text_input, length)
        if html_file_path and os.path.exists(html_file_path):
            with open(html_file_path, 'r') as file:
                html_content = file.read()
            components.html(html_content, height=800, width=800)
        else:
            st.error("Error generating HTML file. Please try again.")