import gradio as gr
import os
import shutil

# --- Configuration ---
INITIAL_PATH = os.getcwd()  # Start in the current working directory
ALLOWED_EXTENSIONS_UPLOAD = None  # Allow all file types for upload initially, or set a list like ['.txt', '.pdf']
ALLOWED_EXTENSIONS_DOWNLOAD = None # Allow all for download, or restrict if needed

# --- Helper Functions ---

def list_files_and_folders(path):
    """Lists files and folders in a given directory."""
    try:
        items = os.listdir(path)
        files = []
        folders = []
        for item in items:
            item_path = os.path.join(path, item)
            if os.path.isfile(item_path):
                files.append(item)
            elif os.path.isdir(item_path):
                folders.append(item)
        return folders, files, path, None  # Return folders, files, current path, and no error
    except Exception as e:
        return [], [], path, str(e) # Return empty lists and error message

def change_directory(path, current_path, direction):
    """Changes the current directory (up or into a folder)."""
    new_path = current_path
    if direction == "up":
        new_path = os.path.dirname(current_path)
    elif direction == "enter":
        new_path = os.path.join(current_path, path) # Path here is the selected folder
    elif direction == "refresh":
        new_path = current_path # Just refresh current

    if not os.path.exists(new_path) or not os.path.isdir(new_path):
        return [], [], current_path, "Invalid directory path."

    return list_files_and_folders(new_path)

def upload_file(files, current_path):
    """Uploads files to the current directory."""
    if not files:
        return [], [], current_path, "No files uploaded."

    uploaded_filenames = []
    for file in files:
        try:
            filename = os.path.basename(file.name) # Get original filename
            destination_path = os.path.join(current_path, filename)

            if ALLOWED_EXTENSIONS_UPLOAD: # Check allowed extensions if specified
                file_extension = os.path.splitext(filename)[1].lower()
                if file_extension not in ALLOWED_EXTENSIONS_UPLOAD:
                    return [], [], current_path, f"File type '{file_extension}' not allowed for upload."

            shutil.copyfile(file.name, destination_path) # Copy uploaded file
            uploaded_filenames.append(filename)
        except Exception as e:
            return [], [], current_path, f"Error uploading file '{filename}': {str(e)}"

    # Refresh file list after upload
    return list_files_and_folders(current_path)


def download_file(file_path, current_path):
    """Returns the file path for download."""
    full_file_path = os.path.join(current_path, file_path) # Construct full path

    if not os.path.isfile(full_file_path):
        return None, "File not found for download."

    if ALLOWED_EXTENSIONS_DOWNLOAD: # Check allowed extensions if specified
        file_extension = os.path.splitext(full_file_path)[1].lower()
        if file_extension not in ALLOWED_EXTENSIONS_DOWNLOAD:
            return None, f"File type '{file_extension}' not allowed for download."

    return full_file_path, None


# --- Gradio Interface ---

with gr.Blocks() as demo:
    current_directory = gr.State(INITIAL_PATH)

    with gr.Row():
        current_path_display = gr.Textbox(value=INITIAL_PATH, label="Current Path", interactive=False, scale=4)
        refresh_button = gr.Button("Refresh", scale=1)
        up_button = gr.Button("Up", scale=1)

    with gr.Row():
        with gr.Column(scale=1):
            folder_output = gr.List([], headers=["Folders"], label="Folders",  elem_id="folder-list", interactive=False) # Removed type="value"
            file_output = gr.List([], headers=["Files"], label="Files", elem_id="file-list", interactive=False) # Removed type="value"
        with gr.Column(scale=2):
            upload_component = gr.File(label="Upload Files", file_count="multiple") # Allows multiple file uploads
            download_button = gr.File(label="Download Selected File")
            error_output = gr.Markdown(label="Status")


    # --- Functionality ---

    def update_file_list(current_path):
        folders, files, updated_path, error = list_files_and_folders(current_path)
        return folders, files, updated_path, updated_path, error

    def on_folder_select(folder_name, current_path_state):
        return change_directory(folder_name, folder_output.value, "enter") # Corrected: use folder_output.value not folder_name directly

    def on_up_button_click(current_path_state):
        return change_directory("", current_path_state, "up") # Path is empty for "up"

    def on_refresh_button_click(current_path_state):
        return change_directory("", current_path_state, "refresh") # Path is empty for "refresh"

    def on_file_upload(uploaded_files, current_path_state):
        return upload_file(uploaded_files, current_path_state)

    def on_file_select_for_download(file_name, current_path_state):
        return download_file(file_name, current_path_state)


    # --- Event Handlers ---

    demo.load(update_file_list, inputs=current_directory, outputs=[folder_output, file_output, current_directory, current_path_display, error_output])

    folder_output.select(on_folder_select, [folder_output, current_directory], [folder_output, file_output, current_directory, current_path_display, error_output])
    up_button.click(on_up_button_click, inputs=current_directory, outputs=[folder_output, file_output, current_directory, current_path_display, error_output])
    refresh_button.click(on_refresh_button_click, inputs=current_directory, outputs=[folder_output, file_output, current_directory, current_path_display, error_output])
    upload_component.upload(on_file_upload, [upload_component, current_directory], [folder_output, file_output, current_directory, current_path_display, error_output] ) # clear_value=True to reset upload after submission
    file_output.select(on_file_select_for_download, [file_output, current_directory], [download_button, error_output])


if __name__ == "__main__":
    demo.launch(share=True,server_port=1337)