import os
import requests
from flask import Flask, request, jsonify
import gdown

app = Flask(__name__)

HF_TOKEN = os.environ.get("HF_TOKEN")
BUCKET_NAME = "gopallikehack/Hitek-storage"

@app.route("/", methods=["GET", "POST"])
def upload():
    if request.method == "GET":
        return "<h1>🚀 Send POST with JSON: {\"file_id\": \"YOUR_FILE_ID\"}</h1>"
    
    data = request.get_json()
    file_id = data.get("file_id")
    if not file_id:
        return jsonify({"error": "Missing file_id"}), 400
    
    # Download from Google Drive
    url = f"https://drive.google.com/uc?id={file_id}"
    filename = "data_part_001.zip"
    print(f"⏳ Downloading from Google Drive...")
    gdown.download(url, filename, quiet=False)
    
    if not os.path.exists(filename):
        return jsonify({"error": "Download failed"}), 500
    
    # ✅ SAHI HF BUCKET UPLOAD (API v2)
    print(f"⏳ Uploading to Hugging Face Bucket...")
    with open(filename, "rb") as f:
        content = f.read()
    
    # Correct endpoint for buckets (without /storage/)
    upload_url = f"https://huggingface.co/api/buckets/{BUCKET_NAME}/upload/{filename}"
    headers = {
        "Authorization": f"Bearer {HF_TOKEN}",
        "Content-Type": "application/octet-stream"
    }
    resp = requests.put(upload_url, headers=headers, data=content)
    
    os.remove(filename)
    
    if resp.status_code in [200, 201]:
        link = f"https://huggingface.co/buckets/{BUCKET_NAME}/blob/main/{filename}"
        print(f"✅ Uploaded Successfully!")
        return jsonify({"success": True, "link": link})
    else:
        print(f"❌ Upload Failed: {resp.status_code}")
        print(f"Response: {resp.text}")
        return jsonify({"error": f"Upload failed: {resp.status_code}"}), 500

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)