import gradio as gr
import requests
from huggingface_hub import HfFileSystem

def download_to_bucket(url, bucket_path):
    # The HfFileSystem automatically picks up the HF_TOKEN secret
    fs = HfFileSystem()
    
    try:
        # Stream the file from the URL
        with requests.get(url, stream=True) as response:
            response.raise_for_status()
            
            # Open the Hugging Face bucket destination for writing
            with fs.open(bucket_path, "wb") as f:
                # Write the file in 8MB chunks to prevent memory overloads
                for chunk in response.iter_content(chunk_size=8 * 1024 * 1024):
                    if chunk:
                        f.write(chunk)
                        
        return f"✅ Successfully streamed into {bucket_path}"
    
    except Exception as e:
        return f"❌ Error: {str(e)}"

# Create a simple Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# 🚀 Stream URL directly to Hugging Face Bucket")
    gr.Markdown("This tool streams files directly into an S3-like HF Storage Bucket without filling up this Space's disk.")
    
    with gr.Row():
        url_input = gr.Textbox(
            label="Source URL", 
            placeholder="https://example.com/huge-dataset.zip"
        )
        path_input = gr.Textbox(
            label="Destination Path", 
            placeholder="hf://buckets/vish85521/videos"
        )
        
    download_btn = gr.Button("Download to Bucket", variant="primary")
    output_text = gr.Textbox(label="Status")
    
    download_btn.click(
        fn=download_to_bucket,
        inputs=[url_input, path_input],
        outputs=output_text
    )

demo.launch()