from fastapi import FastAPI, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
from huggingface_hub import upload_file
import os

app = FastAPI()

# CORS setup to allow requests from your frontend
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Replace "*" with your frontend domain in production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.get("/")
def health_check():
    return {"status": "✅ FastAPI running on Hugging Face Spaces!"}





@app.post("/upload-zip")
async def upload_zip(file: UploadFile = File(...)):
    if not file.filename.endswith(".zip"):
        return {"error": "Please upload a .zip file"}

    # Save the ZIP to /tmp
    temp_zip_path = f"/tmp/{file.filename}"
    with open(temp_zip_path, "wb") as f:
        f.write(await file.read())

    try:
        with tempfile.TemporaryDirectory() as extract_dir:
            # Extract all contents of the zip file
            with zipfile.ZipFile(temp_zip_path, 'r') as zip_ref:
                zip_ref.extractall(extract_dir)

            uploaded_files = []

            # Walk through the extracted files and upload
            for root_dir, _, files in os.walk(extract_dir):
                for name in files:
                    file_path = os.path.join(root_dir, name)
                    relative_path = os.path.relpath(file_path, extract_dir)
                    repo_path = f"demo/{relative_path}".replace("\\", "/")  # handle Windows paths

                    upload_file(
                        path_or_fileobj=file_path,
                        path_in_repo=repo_path,
                        repo_id="rahul7star/ohamlab",
                        repo_type="model",
                        commit_message=f"Upload {relative_path} from zip",
                        token=True,
                    )
                    uploaded_files.append(repo_path)

        return {"message": f"✅ Uploaded {len(uploaded_files)} files", "files": uploaded_files}

    except Exception as e:
        return {"error": f"❌ Failed to process zip: {str(e)}"}

    

@app.post("/upload")
async def upload_image(file: UploadFile = File(...)):
    filename = file.filename
    contents = await file.read()

    # Save temporarily
    temp_path = f"/tmp/{filename}"
    with open(temp_path, "wb") as f:
        f.write(contents)

    try:
        upload_file(
            path_or_fileobj=temp_path,
            path_in_repo=f"demo/{filename}",
            repo_id="rahul7star/ohamlab",
            repo_type="model",
            commit_message=f"Upload {filename} to demo folder",
            token=True,  # uses HF_TOKEN from Space secrets
        )
        return {"message": f"✅ {filename} uploaded successfully!"}
    except Exception as e:
        return {"error": f"❌ Upload failed: {str(e)}"}