import gradio as gr
import os
import pandas as pd
from huggingface_hub import HfFileSystem
import tempfile

# --- CONFIGURATION ---
HF_BUCKET = "hf://buckets/Krixter/storage"

def get_fs():
    """Initialize the Hugging Face File System. Uses a token if available, otherwise runs anonymously."""
    token = os.environ.get("HF_TOKEN")
    # If there's no token, HfFileSystem will just run in public/anonymous mode
    return HfFileSystem(token=token)

def format_size(size_bytes):
    """Convert bytes to a human-readable format."""
    for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
        if size_bytes < 1024.0:
            return f"{size_bytes:.2f} {unit}"
        size_bytes /= 1024.0

def list_files():
    """Fetch files from the bucket and return a formatted DataFrame."""
    try:
        fs = get_fs()
        # Fetch file list with details (size, modified date)
        files = fs.ls(HF_BUCKET, detail=True)
        
        data = []
        for f in files:
            if f["type"] == "file":
                name = os.path.basename(f["name"])
                size_str = format_size(f["size"])
                last_modified = f.get("last_modified", "Unknown")
                
                data.append({
                    "File Name": name,
                    "Last Modified": last_modified,
                    "Size": size_str
                })
        
        df = pd.DataFrame(data)
        if not df.empty:
            # Convert 'Last Modified' to datetime for accurate chronological sorting
            df["Last Modified"] = pd.to_datetime(df["Last Modified"]).dt.tz_localize(None)
            
            # Default sort by File Name (alphabetical)
            df = df.sort_values(by="File Name", ascending=True)
            
        file_names = df["File Name"].tolist() if not df.empty else []
        return df, gr.update(choices=file_names)
        
    except Exception as e:
        print(f"Error fetching files: {e}")
        return pd.DataFrame(columns=["File Name", "Last Modified", "Size"]), gr.update(choices=[])

def download_file(filename):
    """Download the selected file from the bucket to a temporary local file."""
    if not filename:
        return None
    
    try:
        fs = get_fs()
        src_path = f"{HF_BUCKET}/{filename}"
        
        # Create a temporary directory to hold the downloaded file
        temp_dir = tempfile.mkdtemp()
        dest_path = os.path.join(temp_dir, filename)
        
        # Pull the file down from the bucket
        print(f"📥 Downloading {filename}...")
        fs.get(src_path, dest_path)
        print("✅ Download complete!")
        
        return dest_path
    except Exception as e:
        print(f"⚠️ Error downloading {filename}: {e}")
        return None

# --- UI LAYOUT ---
with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo")) as app:
    gr.Markdown("# 🗄️ AudioNovel Bucket Explorer")
    gr.Markdown("View, sort, and download your processed audio files. **Click the column headers below to sort by Name, Date, or Size.**")
    
    with gr.Row():
        refresh_btn = gr.Button("🔄 Refresh Bucket contents", variant="secondary")
        
    # Interactive table
    file_table = gr.Dataframe(
        headers=["File Name", "Last Modified", "Size"],
        interactive=False,
        label="Bucket Contents"
    )
    
    gr.Markdown("---")
    gr.Markdown("### 📥 Download a File")
    
    with gr.Row():
        with gr.Column(scale=2):
            file_dropdown = gr.Dropdown(label="Select File to Download", choices=[])
        with gr.Column(scale=1):
            download_btn = gr.Button("🚀 Prepare Download", variant="primary")
            
    file_output = gr.File(label="Ready for Download")
    
    # --- EVENT TRIGGERS ---
    # Trigger refresh
    refresh_btn.click(fn=list_files, inputs=[], outputs=[file_table, file_dropdown])
    # Trigger download
    download_btn.click(fn=download_file, inputs=[file_dropdown], outputs=[file_output])
    # Load data automatically on startup
    app.load(fn=list_files, inputs=[], outputs=[file_table, file_dropdown])

if __name__ == "__main__":
    app.launch()