from fastapi import FastAPI, File, UploadFile
from huggingface_hub import HfApi
import os

# Initialize FastAPI app
app = FastAPI()
api = HfApi()

# Get repo_id and token from environment variables (secrets)
repo_id = os.getenv("HF_REPO_ID")  # Set in Space secrets (e.g., "username/my-model")
token = os.getenv("HF_TOKEN")      # Set in Space secrets

# --- FastAPI Endpoint ---
@app.post("/upload/")
async def upload_file(folder_name: str, file: UploadFile = File(...)):
    file_path = f"/tmp/temp_{file.filename}"
    with open(file_path, "wb") as f:
        f.write(await file.read())
    
    # Upload to Hugging Face with dynamic folder name provided by the user
    api.upload_file(
        path_or_fileobj=file_path,
        path_in_repo=f"{folder_name}/{file.filename}",  # Dynamic folder name
        repo_id=repo_id,
        repo_type="model",
        token=token
    )
    os.remove(file_path)
    return {"message": f"File uploaded as {file.filename} to {folder_name}/{file.filename}"}