app.py

import os
import uuid
from typing import List
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse
import dagshub
from dagshub import upload_files
from dagshub.data_engine import datasources

# =====================================================
# CONFIG
# =====================================================

REPO_OWNER = "Rizwanali324"
REPO_NAME = "my-first-repo"
FULL_REPO = f"{REPO_OWNER}/{REPO_NAME}"

REMOTE_FOLDER = "data"
BRANCH = "main"

# =====================================================
# AUTHENTICATION
# =====================================================

if "DAGSHUB_USER_TOKEN" not in os.environ:
    raise RuntimeError(
        "DAGSHUB_USER_TOKEN secret not found. Add it in Space Settings → Secrets."
    )

dagshub.init(
    repo_owner=REPO_OWNER,
    repo_name=REPO_NAME,
    mlflow=False
)

# =====================================================
# FASTAPI INIT
# =====================================================

app = FastAPI(title="Wildlife Detection Image API")

# local temp folder
os.makedirs("temp_uploads", exist_ok=True)


# =====================================================
# IMAGE UPLOAD ENDPOINT
# =====================================================

@app.post("/upload")
async def upload_images(
    files: List[UploadFile] = File(..., description="Upload image files")
):

    if not files:
        raise HTTPException(status_code=400, detail="No files uploaded")

    uploaded_files = []

    for file in files:

        if not file.content_type.startswith("image/"):
            raise HTTPException(
                status_code=400,
                detail=f"{file.filename} is not a valid image"
            )

        original_name = file.filename

        # generate unique filename
        unique_id = uuid.uuid4().hex
        stored_name = f"{unique_id}_{original_name}"

        local_path = os.path.join("temp_uploads", stored_name)

        try:
            content = await file.read()

            # save locally
            with open(local_path, "wb") as f:
                f.write(content)

        except Exception as e:
            raise HTTPException(
                status_code=500,
                detail=f"Failed saving {original_name}: {str(e)}"
            )

        # upload to dagshub
        try:
            upload_files(
                repo=FULL_REPO,
                local_path=local_path,
                remote_path=f"{REMOTE_FOLDER}/{stored_name}"
            )

        except Exception as e:
            raise HTTPException(
                status_code=500,
                detail=f"Upload failed for {original_name}: {str(e)}"
            )

        # public url
        file_url = (
            f"https://dagshub.com/{REPO_OWNER}/{REPO_NAME}/raw/"
            f"{BRANCH}/{REMOTE_FOLDER}/{stored_name}"
        )

        uploaded_files.append({
            "original_filename": original_name,
            "stored_filename": stored_name,
            "url": file_url
        })

    return JSONResponse(
        content={
            "message": "Images uploaded successfully",
            "count": len(uploaded_files),
            "files": uploaded_files
        }
    )


# =====================================================
# FETCH DATA FROM DAGSHUB
# =====================================================

@app.get("/data")
def get_uploaded_data(min_size: int = None):

    try:
        ds = datasources.get_datasource(FULL_REPO)

        query = ds

        if min_size:
            query = ds["size"] > min_size

        query = query.select("path", "size")

        df = query.all().dataframe

        records = df.to_dict(orient="records")

        return {
            "message": "Data fetched successfully",
            "count": len(records),
            "data": records
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# =====================================================
# HEALTH CHECK
# =====================================================

@app.get("/")
def root():
    return {"status": "API running successfully"}


# =====================================================
# RUN SERVER
# =====================================================

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