import gradio as gr
import zipfile
import io
import re
import traceback

def process_files(phase, objecten_file_data, templates_zip_file_data):
    try:
        # Read and process the 'objecten.txt' file
        content = objecten_file_data.decode('utf-8')
        t = content.strip().split('\n\n')
        phases = ["iFAT", "(i)SAT"]
        objs = {p: [] for p in phases}
        if phase in phases:
            objs = {phase: []}
        else:
            objs = {p: [] for p in phases}

        regObject = r"\d{4}[a-zA-Z]{2}"
        for g in t:
            ls = [line.strip() for line in g.strip().split('\n') if line.strip()]
            k = ls[0].strip()
            if k in objs:
                objs[k] = ls[1:]
            else:
                error_msg = f"Key [{k}] is not recognized."
                return None, error_msg

        # Extract object codes
        for k in objs:
            objs[k] = [re.search(regObject, o).group(0) for o in objs[k] if re.search(regObject, o)]

        # Read template files from the uploaded ZIP file
        templates_zip_data = templates_zip_file_data
        with zipfile.ZipFile(io.BytesIO(templates_zip_data), 'r') as zip_ref:
            template_files = {info.filename: zip_ref.read(info.filename) for info in zip_ref.infolist()}

        # Create an in-memory ZIP file for the output
        output_zip_buffer = io.BytesIO()
        with zipfile.ZipFile(output_zip_buffer, 'w') as zf:
            for k in objs:
                regPhase = ''
                if k == '(i)SAT':
                    regPhase = r'sat'
                elif k == 'iFAT':
                    regPhase = r'fat'

                # Filter template files for this phase
                phase_templates = []
                for filename, file_data in template_files.items():
                    if re.search(regPhase, filename, re.IGNORECASE):
                        phase_templates.append((filename, file_data))

                if not phase_templates:
                    error_msg = f"Phase {k} has no templates."
                    return None, error_msg

                for o in objs[k]:
                    folder_path = f"{o}/"
                    for template_filename, file_content in phase_templates:
                        # Adjust filename if needed
                        if re.search(r"hut_\d{4}[a-zA-Z]{2}", template_filename, re.IGNORECASE):
                            adjusted_filename = template_filename[:4] + o + template_filename[10:]
                        else:
                            adjusted_filename = template_filename

                        file_path = folder_path + adjusted_filename
                        zf.writestr(file_path, file_content)
        
        # Finalize the output
        output_zip_buffer.seek(0)
        return output_zip_buffer.getvalue(), ""  # Return bytes and an empty error log on success
    
    except Exception as e:
        # Capture the full traceback
        error_msg = traceback.format_exc()
        return None, error_msg

phase_options = ['iFAT', '(i)SAT', 'All']

interface = gr.Interface(
    fn=process_files,
    inputs=[
        gr.Dropdown(choices=phase_options, label="Select Phase"),
        gr.File(label="Upload 'objecten.txt' File", type='binary'),
        gr.File(label="Upload Templates ZIP File", type='binary')
    ],
    outputs=[
        gr.File(label="Download ZIP File"),  # This now expects bytes or path
        gr.Textbox(label="Error Log")        # Textbox to display the error log
    ],
    title="Template Folder Generator",
    description="Upload 'objecten.txt' and a ZIP file containing template files to generate folders and files."
)

interface.launch()