import gradio as gr
import os
import zipfile
from pathlib import Path

# The base directory that will be explored
BASE_DIR = Path(__file__).resolve().parent / "your_target_folder"

# Ensure BASE_DIR exists
BASE_DIR.mkdir(parents=True, exist_ok=True)

def resolve_path(current_rel_path: str) -> Path:
    """Secure path resolution within BASE_DIR"""
    resolved_path = (BASE_DIR / current_rel_path).resolve()
    if BASE_DIR not in resolved_path.parents and resolved_path != BASE_DIR:
        raise ValueError("Access outside base directory is not allowed.")
    return resolved_path

def list_dir(current_rel_path: str = ""):
    current_path = resolve_path(current_rel_path)

    # Parent folder navigation
    parent_rel = str(Path(current_rel_path).parent) if current_rel_path else ""
    entries = []

    if current_path != BASE_DIR:
        entries.append(("..", "⬆️ Parent Folder"))

    # List directories and files
    for item in sorted(current_path.iterdir()):
        rel_item = os.path.relpath(item, BASE_DIR)
        label = f"📁 {item.name}" if item.is_dir() else f"📄 {item.name}"
        entries.append((rel_item, label))

    return gr.update(choices=entries, value=None), f"Currently in: /{current_rel_path}"

def download_entry(selected_rel_path: str):
    selected_path = resolve_path(selected_rel_path)

    if selected_path.is_file():
        return selected_path
    elif selected_path.is_dir():
        zip_path = f"/tmp/{selected_path.name}.zip"
        with zipfile.ZipFile(zip_path, "w") as zipf:
            for root, dirs, files in os.walk(selected_path):
                for file in files:
                    abs_file = os.path.join(root, file)
                    arc_file = os.path.relpath(abs_file, selected_path)
                    zipf.write(abs_file, arc_file)
        return zip_path
    else:
        return None

with gr.Blocks() as demo:
    current_path = gr.State("")

    gr.Markdown("# 📁 Folder Browser")
    
    with gr.Row():
        folder_dropdown = gr.Dropdown(label="Select File or Folder", choices=[])
        refresh_btn = gr.Button("🔄 Refresh")
    
    status_text = gr.Textbox(label="Current Path", interactive=False)
    download_btn = gr.Button("⬇️ Download Selected")
    file_output = gr.File(label="Download Result")

    # Events
    refresh_btn.click(fn=list_dir, inputs=current_path, outputs=[folder_dropdown, status_text])
    
    folder_dropdown.change(
        fn=lambda x: (x, *list_dir(x)),  # update path, refresh list
        inputs=folder_dropdown,
        outputs=[current_path, folder_dropdown, status_text],
    )

    download_btn.click(fn=download_entry, inputs=folder_dropdown, outputs=file_output)

    # Initial trigger
    demo.load(fn=list_dir, inputs=current_path, outputs=[folder_dropdown, status_text])

demo.launch()