screencoder/main.py

import subprocess
import sys
import os
import json


# A simple placeholder for prompt injection
# In a real scenario, this should be a more robust mechanism
def inject_prompt_to_generator(prompt_text):
    if not prompt_text:
        return

    # In this example, we assume the prompt is a simple string for the "main content"
    # A more complex implementation would parse a structured prompt
    user_instruction = {
        "sidebar": "Make all icons look better; fill in relevant English text; beautify the layout.",
        "header": "Make the Google logo look better; change the avatar color to be more appealing.",
        "navigation": "Please beautify the layout.",
        "main content": prompt_text
    }
    
    generator_path = os.path.join(os.path.dirname(__file__), 'html_generator.py')
    with open(generator_path, 'r', encoding='utf-8') as f:
        lines = f.readlines()
    
    # Find the user_instruction dictionary and replace it
    new_lines = []
    in_dict = False
    for line in lines:
        if 'user_instruction = {' in line:
            in_dict = True
            new_lines.append(f"user_instruction = {json.dumps(user_instruction, indent=4)}\n")
        elif in_dict and '}' in line:
            in_dict = False
            continue # Skip the closing brace of the old dict
        elif not in_dict:
            new_lines.append(line)

    with open(generator_path, 'w', encoding='utf-8') as f:
        f.writelines(new_lines)


def run_script(script_path):
    script_path = os.path.normpath(script_path)
    
    print(f"\n{'='*20}")
    print(f"Executing: python {script_path}")
    print(f"{'='*20}")
    
    try:
        result = subprocess.run(
            [sys.executable, script_path],
            check=True,
            capture_output=True,
            text=True
        )
        print("Success!")
        print("Output:")
        print(result.stdout)
        if result.stderr:
            print("Stderr:")
            print(result.stderr)
    except FileNotFoundError:
        print(f"ERROR: Script not found at '{script_path}'")
        sys.exit(1)
    except subprocess.CalledProcessError as e:
        print(f"ERROR: Script '{script_path}' failed with exit code {e.returncode}")
        print("Stdout:")
        print(e.stdout)
        print("Stderr:")
        print(e.stderr)
        sys.exit(1)
    except Exception as e:
        print(f"An unexpected error occurred while running '{script_path}': {e}")
        sys.exit(1)


def generate_html_for_demo(image_path, prompt, output_dir="screencoder/data/output"):
    """
    A modified workflow for the Gradio demo.
    It takes an image path and a prompt, and returns the path to the final HTML file.
    """
    print("Starting the Screencoder demo workflow...")

    # Setup paths
    project_root = os.path.dirname(__file__)
    # The block_parsor script expects a specific input file name, so we must place our image there.
    # IMPORTANT: This assumes a single-user-at-a-time workflow.
    # For multi-user, you'd need isolated temp directories.
    target_input_path = os.path.join(project_root, "data/input/test1.png")
    
    # Ensure the input directory exists
    os.makedirs(os.path.dirname(target_input_path), exist_ok=True)

    # Copy the user-uploaded image to the location the script expects
    import shutil
    shutil.copy(image_path, target_input_path)

    # --- Part 1: Initial Generation with Placeholders ---
    print("\n--- Part 1: Initial Generation with Placeholders ---")
    inject_prompt_to_generator(prompt)
    run_script(os.path.join(project_root, "block_parsor.py"))
    run_script(os.path.join(project_root, "html_generator.py"))

    # --- Part 2: Final HTML Code Generation ---
    print("\n--- Part 2: Final HTML Code Generation ---")
    run_script(os.path.join(project_root, "image_box_detection.py"))
    run_script(os.path.join(project_root, "UIED/run_single.py"))
    run_script(os.path.join(project_root, "mapping.py"))
    run_script(os.path.join(project_root, "image_replacer.py"))

    final_html_path = os.path.join(output_dir, "test1_layout_final.html")
    print(f"\nScreencoder demo workflow completed! Final HTML at: {final_html_path}")

    # Check if the final file exists
    if os.path.exists(final_html_path):
        with open(final_html_path, 'r', encoding='utf-8') as f:
            return f.read()
    else:
        return "<html><body><h1>Error: Final HTML not generated.</h1></body></html>"


def main():
    """Main function to run the entire Screencoder workflow."""
    print("Starting the Screencoder full workflow...")

    # --- Part 1: Initial Generation with Placeholders ---
    print("\n--- Part 1: Initial Generation with Placeholders ---")
    run_script("block_parsor.py")
    run_script("html_generator.py")

    # --- Part 2: Final HTML Code Generation ---
    print("\n--- Part 2: Final HTML Code Generation ---")
    run_script("image_box_detection.py")
    run_script("UIED/run_single.py")
    run_script("mapping.py")
    run_script("image_replacer.py")

    print("\nScreencoder workflow completed successfully!")

if __name__ == "__main__":
    main()