import os
import requests
import gradio as gr
from huggingface_hub import HfApi
from pathlib import Path

# Get token from Space Secrets
HF_TOKEN = os.environ.get("HF_TOKEN")
api = HfApi(token=HF_TOKEN)

def get_repo_info(url):
    """Extracts repo_id and repo_type from a standard HF URL."""
    try:
        parts = url.strip().split("huggingface.co/")[1].split("/")
        if parts[0] == "datasets":
            return f"{parts[1]}/{parts[2]}", "dataset"
        else:
            return f"{parts[0]}/{parts[1]}", "model"
    except:
        return None, None

def download_and_upload(dataset_url, download_url, progress=gr.Progress()):
    if not HF_TOKEN:
        return "❌ Error: HF_TOKEN not found in Secrets."
    
    repo_id, repo_type = get_repo_info(dataset_url)
    if not repo_id:
        return "❌ Error: Invalid Dataset/Repo URL format."

    filename = download_url.split("/")[-1].split("?")[0]
    local_path = Path(filename)

    try:
        # Step 1: Download with status updates
        progress(0, desc="Initializing stream...")
        with requests.get(download_url, stream=True, timeout=60) as r:
            r.raise_for_status()
            total_size = int(r.headers.get('content-length', 0))
            downloaded = 0
            
            with open(local_path, 'wb') as f:
                for chunk in r.iter_content(chunk_size=1024*1024): # 1MB chunks
                    if chunk:
                        f.write(chunk)
                        downloaded += len(chunk)
                        if total_size > 0:
                            done = downloaded / total_size
                            # Update every few MBs to keep the connection alive
                            progress(done * 0.5, desc=f"Downloading: {downloaded/(1024**3):.2f}GB / {total_size/(1024**3):.2f}GB")
        
        # Step 2: Upload to HF
        progress(0.6, desc="Uploading to Hugging Face...")
        api.upload_file(
            path_or_fileobj=str(local_path),
            path_in_repo=filename,
            repo_id=repo_id,
            repo_type=repo_type,
            commit_message=f"Uploaded {filename} via UpDownUrl",
            # This is key for large files:
            run_as_future=False 
        )
        
        return f"✅ Done! '{filename}' is now in {repo_id}"

    except Exception as e:
        return f"❌ Error: {str(e)}"
    finally:
        if local_path.exists():
            os.remove(local_path)

# --- UI Setup ---
with gr.Blocks() as demo:
    gr.Markdown("# 🚀 UpDownUrl v1.1")
    
    with gr.Row():
        with gr.Column():
            dataset_link = gr.Textbox(label="Target Repo Link", placeholder="https://huggingface.co/datasets/AIDev07/AIModelsLoaded")
            download_link = gr.Textbox(label="Download URL", placeholder="Direct link to file")
            upload_btn = gr.Button("Boom! Upload", variant="primary")
        
        with gr.Column():
            output_log = gr.Textbox(label="Status", interactive=False)

    upload_btn.click(
        fn=download_and_upload,
        inputs=[dataset_link, download_link],
        outputs=output_log
    )

# Fixed Gradio 6.0 syntax: Theme goes here
if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft())