import os
import shutil
import zipfile
import tempfile
import gradio as gr

def process(phase, template_zip, object_file):
    try:
        if template_zip is None:
            return "Error: Please upload the template directory as a ZIP file."
        if object_file is None:
            return "Error: Please upload the object names TXT file."
        
        # Create a temporary directory for processing
        with tempfile.TemporaryDirectory() as tmpdir:
            # Unzip the template directory
            template_dir = os.path.join(tmpdir, 'template')
            with zipfile.ZipFile(template_zip.name, 'r') as zip_ref:
                zip_ref.extractall(template_dir)
            
            # Read object names from the txt file
            object_names = object_file.read().decode('utf-8').splitlines()
            
            # Create a directory to store the results
            result_dir = os.path.join(tmpdir, 'results')
            os.makedirs(result_dir, exist_ok=True)
            
            for object_name in object_names:
                if not object_name.strip():
                    continue  # Skip empty lines
                # Create a directory for the object under result_dir
                object_dir = os.path.join(result_dir, f"{phase}_{object_name}")
                os.makedirs(object_dir, exist_ok=True)
                
                # Copy files from template_dir to object_dir
                for root, dirs, files in os.walk(template_dir):
                    rel_path = os.path.relpath(root, template_dir)
                    dest_dir = os.path.join(object_dir, rel_path)
                    os.makedirs(dest_dir, exist_ok=True)
                    for filename in files:
                        template_file = os.path.join(root, filename)
                        if os.path.isfile(template_file):
                            # Get the file extension
                            name, ext = os.path.splitext(filename)
                            # Create new filename with object name appended
                            new_filename = f"{name}_{object_name}{ext}"
                            new_file = os.path.join(dest_dir, new_filename)
                            shutil.copy(template_file, new_file)
            
            # Zip the result_dir
            result_zip_path = os.path.join(tmpdir, 'result.zip')
            shutil.make_archive(os.path.splitext(result_zip_path)[0], 'zip', result_dir)
            
            # Return the zip file for download
            return result_zip_path
    except Exception as e:
        return f"Error: {str(e)}"

phase_dropdown = gr.Dropdown(choices=['FAT', 'iFAT', 'iSAT'], label='Select Phase')

template_zip_input = gr.File(label='Template Directory (ZIP file)', file_types=['.zip'])

object_file_input = gr.File(label='Object Names File (TXT)', file_types=['.txt'])

output_file = gr.File(label='Download Result ZIP')

gr.Interface(
    fn=process,
    inputs=[phase_dropdown, template_zip_input, object_file_input],
    outputs=output_file,
    title='Directory and File Creator',
    description='Upload the template directory as a ZIP file and the object names TXT file. The processed files will be provided as a downloadable ZIP.'
).launch()