api_server.py

# from fastapi import FastAPI, HTTPException, UploadFile, File, Form
# from pydantic import BaseModel
# import numpy as np
# from PIL import Image
# import io, uuid, os, shutil, timeit
# from datetime import datetime
# from fastapi.staticfiles import StaticFiles
# from fastapi.middleware.cors import CORSMiddleware

# # import your three wrappers
# from app import predict_simple, predict_middle, predict_full

# app = FastAPI()

# # allow CORS if needed
# app.add_middleware(
#     CORSMiddleware,
#     allow_origins=["*"],
#     allow_methods=["*"],
#     allow_headers=["*"],
# )

# BASE_URL = "https://snapanddtraceapp-988917236820.us-central1.run.app"
# OUTPUT_DIR = os.path.abspath("./outputs")
# os.makedirs(OUTPUT_DIR, exist_ok=True)
# app.mount("/outputs", StaticFiles(directory=OUTPUT_DIR), name="outputs")

# UPDATES_DIR = os.path.abspath("./updates")
# os.makedirs(UPDATES_DIR, exist_ok=True)
# app.mount("/updates", StaticFiles(directory=UPDATES_DIR), name="updates")


# def save_and_build_urls(
#     session_id: str,
#     output_image: np.ndarray,
#     outlines: np.ndarray,
#     dxf_path: str,
#     mask: np.ndarray
# ):
#     """Helper to save all four artifacts and return public URLs."""
#     request_dir = os.path.join(OUTPUT_DIR, session_id)
#     os.makedirs(request_dir, exist_ok=True)

#     # filenames
#     out_fn = "overlay.jpg"
#     outlines_fn = "outlines.jpg"
#     mask_fn = "mask.jpg"
#     current_date = datetime.now().strftime("%d-%m-%Y")
#     dxf_fn = f"out_{current_date}_{session_id}.dxf"

#     # full paths
#     out_path = os.path.join(request_dir, out_fn)
#     outlines_path = os.path.join(request_dir, outlines_fn)
#     mask_path = os.path.join(request_dir, mask_fn)
#     new_dxf_path = os.path.join(request_dir, dxf_fn)

#     # save images
#     Image.fromarray(output_image).save(out_path)
#     Image.fromarray(outlines).save(outlines_path)
#     Image.fromarray(mask).save(mask_path)

#     # copy dx file
#     if os.path.exists(dxf_path):
#         shutil.copy(dxf_path, new_dxf_path)
#     else:
#         # fallback if your DXF generator returns bytes or string
#         with open(new_dxf_path, "wb") as f:
#             if isinstance(dxf_path, (bytes, bytearray)):
#                 f.write(dxf_path)
#             else:
#                 f.write(str(dxf_path).encode("utf-8"))

#     # build URLs
#     return {
#         "output_image_url": f"{BASE_URL}/outputs/{session_id}/{out_fn}",
#         "outlines_url":      f"{BASE_URL}/outputs/{session_id}/{outlines_fn}",
#         "mask_url":          f"{BASE_URL}/outputs/{session_id}/{mask_fn}",
#         "dxf_url":           f"{BASE_URL}/outputs/{session_id}/{dxf_fn}",
#     }


# @app.post("/predict1")
# async def predict1_api(
#     file: UploadFile = File(...)
# ):
#     """
#     Simple predict: only image → overlay, outlines, mask, DXF
#     """
#     session_id = str(uuid.uuid4())
#     try:
#         img_bytes = await file.read()
#         image = np.array(Image.open(io.BytesIO(img_bytes)).convert("RGB"))
#     except Exception:
#         raise HTTPException(400, "Invalid image upload")

#     try:
#         start = timeit.default_timer()
#         out_img, outlines, dxf_path, mask = predict_simple(image)
#         elapsed = timeit.default_timer() - start
#         print(f"[{session_id}] predict1 in {elapsed:.2f}s")

#         return save_and_build_urls(session_id, out_img, outlines, dxf_path, mask)

#     except Exception as e:
#         raise HTTPException(500, f"predict1 failed: {e}")
#     except ReferenceBoxNotDetectedError:
#         raise HTTPException(status_code=400, detail="Error detecting reference battery! Please try again with a clearer image.")
#     except FingerCutOverlapError:
#         raise HTTPException(status_code=400, detail="There was an overlap with fingercuts!s Please try again to generate dxf.")


# @app.post("/predict2")
# async def predict2_api(
#     file: UploadFile = File(...),
#     enable_fillet: str = Form(..., regex="^(On|Off)$"),
#     fillet_value_mm: float = Form(...)
# ):
#     """
#     Middle predict: image + fillet toggle + fillet value → overlay, outlines, mask, DXF
#     """
#     session_id = str(uuid.uuid4())
#     try:
#         img_bytes = await file.read()
#         image = np.array(Image.open(io.BytesIO(img_bytes)).convert("RGB"))
#     except Exception:
#         raise HTTPException(400, "Invalid image upload")

#     try:
#         start = timeit.default_timer()
#         out_img, outlines, dxf_path, mask = predict_middle(
#             image, enable_fillet, fillet_value_mm
#         )
#         elapsed = timeit.default_timer() - start
#         print(f"[{session_id}] predict2 in {elapsed:.2f}s")

#         return save_and_build_urls(session_id, out_img, outlines, dxf_path, mask)

#     except Exception as e:
#         raise HTTPException(500, f"predict2 failed: {e}")
#     except ReferenceBoxNotDetectedError:
#         raise HTTPException(status_code=400, detail="Error detecting reference battery! Please try again with a clearer image.")
#     except FingerCutOverlapError:
#         raise HTTPException(status_code=400, detail="There was an overlap with fingercuts!s Please try again to generate dxf.")

# @app.post("/predict3")
# async def predict3_api(
#     file: UploadFile = File(...),
#     enable_fillet: str = Form(..., regex="^(On|Off)$"),
#     fillet_value_mm: float = Form(...),
#     enable_finger_cut: str = Form(..., regex="^(On|Off)$")
# ):
#     """
#     Full predict: image + fillet toggle/value + finger-cut toggle → overlay, outlines, mask, DXF
#     """
#     session_id = str(uuid.uuid4())
#     try:
#         img_bytes = await file.read()
#         image = np.array(Image.open(io.BytesIO(img_bytes)).convert("RGB"))
#     except Exception:
#         raise HTTPException(400, "Invalid image upload")

#     try:
#         start = timeit.default_timer()
#         out_img, outlines, dxf_path, mask = predict_full(
#             image, enable_fillet, fillet_value_mm, enable_finger_cut
#         )
#         elapsed = timeit.default_timer() - start
#         print(f"[{session_id}] predict3 in {elapsed:.2f}s")

#         return save_and_build_urls(session_id, out_img, outlines, dxf_path, mask)

#     except Exception as e:
#         raise HTTPException(500, f"predict3 failed: {e}")
#     except ReferenceBoxNotDetectedError:
#         raise HTTPException(status_code=400, detail="Error detecting reference battery! Please try again with a clearer image.")
#     except FingerCutOverlapError:
#         raise HTTPException(status_code=400, detail="There was an overlap with fingercuts!s Please try again to generate dxf.")

# @app.post("/update")
# async def update_files(
#     output_image: UploadFile = File(...),
#     outlines_image: UploadFile = File(...),
#     mask_image: UploadFile = File(...),
#     dxf_file: UploadFile = File(...)
# ):
#     session_id = str(uuid.uuid4())
#     update_dir = os.path.join(UPDATES_DIR, session_id)
#     os.makedirs(update_dir, exist_ok=True)

#     try:
#         upload_map = {
#             "output_image":  output_image,
#             "outlines_image": outlines_image,
#             "mask_image":     mask_image,
#             "dxf_file":       dxf_file,
#         }
#         urls = {}
#         for key, up in upload_map.items():
#             fn = up.filename
#             path = os.path.join(update_dir, fn)
#             with open(path, "wb") as f:
#                 shutil.copyfileobj(up.file, f)
#             urls[key] = f"{BASE_URL}/updates/{session_id}/{fn}"

#         return {"session_id": session_id, "uploaded": urls}

#     except Exception as e:
#         raise HTTPException(500, f"Update failed: {e}")


# if __name__ == "__main__":
#     import uvicorn
#     port = int(os.environ.get("PORT", 8082))
#     print(f"Starting FastAPI server on 0.0.0.0:{port}...")
#     uvicorn.run(app, host="0.0.0.0", port=port)










from fastapi import FastAPI, HTTPException, UploadFile, File, Form
from pydantic import BaseModel
import numpy as np
from PIL import Image
import io, uuid, os, shutil, timeit
from datetime import datetime
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse

# import your three wrappers
from app import predict_simple, predict_middle, predict_full

from app import (
    predict_simple, predict_middle, predict_full,
    ReferenceBoxNotDetectedError,
    FingerCutOverlapError
)


app = FastAPI()

# allow CORS if needed
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

BASE_URL = "https://snapanddtraceapp-988917236820.us-central1.run.app"

OUTPUT_DIR = os.path.abspath("./outputs")
os.makedirs(OUTPUT_DIR, exist_ok=True)

UPDATES_DIR = os.path.abspath("./updates")
os.makedirs(UPDATES_DIR, exist_ok=True)

# Mount static directories with normal StaticFiles
app.mount("/outputs", StaticFiles(directory=OUTPUT_DIR), name="outputs")
app.mount("/updates", StaticFiles(directory=UPDATES_DIR), name="updates")


def save_and_build_urls(
    session_id: str,
    output_image: np.ndarray,
    outlines: np.ndarray,
    dxf_path: str,
    mask: np.ndarray,
    endpoint_type: str,
    fillet_value: float = None,
    finger_cut: str = None
):
    """Helper to save all four artifacts and return public URLs."""
    request_dir = os.path.join(OUTPUT_DIR, session_id)
    os.makedirs(request_dir, exist_ok=True)

    # filenames
    out_fn = "overlay.jpg"
    outlines_fn = "outlines.jpg"
    mask_fn = "mask.jpg"
    
    # Get current date
    current_date = datetime.utcnow().strftime("%d-%m-%Y")


    # Format fillet value with underscore instead of dot
    fillet_str = f"{fillet_value:.2f}".replace(".", "_") if fillet_value is not None else None
    
    # Determine DXF filename based on endpoint type
    if endpoint_type == "predict1":
        dxf_fn = f"DXF_{current_date}.dxf"
    elif endpoint_type == "predict2":
        dxf_fn = f"DXF_{current_date}.dxf"
    elif endpoint_type == "predict3":
        dxf_fn = f"DXF_{current_date}.dxf"

    # full paths
    out_path = os.path.join(request_dir, out_fn)
    outlines_path = os.path.join(request_dir, outlines_fn)
    mask_path = os.path.join(request_dir, mask_fn)
    new_dxf_path = os.path.join(request_dir, dxf_fn)

    # save images
    Image.fromarray(output_image).save(out_path)
    Image.fromarray(outlines).save(outlines_path)
    Image.fromarray(mask).save(mask_path)

    # copy dxf file
    if os.path.exists(dxf_path):
        shutil.copy(dxf_path, new_dxf_path)
    else:
        # fallback if your DXF generator returns bytes or string
        with open(new_dxf_path, "wb") as f:
            if isinstance(dxf_path, (bytes, bytearray)):
                f.write(dxf_path)
            else:
                f.write(str(dxf_path).encode("utf-8"))

    # build URLs with /download prefix for DXF
    return {
        "output_image_url": f"{BASE_URL}/outputs/{session_id}/{out_fn}",
        "outlines_url":     f"{BASE_URL}/outputs/{session_id}/{outlines_fn}",
        "mask_url":         f"{BASE_URL}/outputs/{session_id}/{mask_fn}",
        "dxf_url":          f"{BASE_URL}/download/{session_id}/{dxf_fn}",  # Changed to use download endpoint
    }

# Add new endpoint for downloading DXF files
@app.get("/download/{session_id}/{filename}")
async def download_file(session_id: str, filename: str):
    file_path = os.path.join(OUTPUT_DIR, session_id, filename)
    if not os.path.exists(file_path):
        raise HTTPException(status_code=404, detail="File not found")
    
    return FileResponse(
        path=file_path,
        filename=filename,
        media_type="application/x-dxf",
        headers={"Content-Disposition": f"attachment; filename={filename}"}
    )


@app.post("/predict1")
async def predict1_api(
    file: UploadFile = File(...)
):
    """
    Simple predict: only image → overlay, outlines, mask, DXF
    DXF naming format: DXF_DD-MM-YYYY.dxf
    """
    session_id = str(uuid.uuid4())
    try:
        img_bytes = await file.read()
        image = np.array(Image.open(io.BytesIO(img_bytes)).convert("RGB"))
    except Exception:
        raise HTTPException(400, "Invalid image upload")

    try:
        start = timeit.default_timer()
        out_img, outlines, dxf_path, mask = predict_simple(image)
        elapsed = timeit.default_timer() - start
        print(f"[{session_id}] predict1 in {elapsed:.2f}s")

        return save_and_build_urls(
            session_id=session_id,
            output_image=out_img,
            outlines=outlines,
            dxf_path=dxf_path,
            mask=mask,
            endpoint_type="predict1"
        )

    except ReferenceBoxNotDetectedError:
        raise HTTPException(status_code=400, detail="Error detecting reference battery! Please try again with a clearer image.")
    except FingerCutOverlapError:
        raise HTTPException(status_code=400, detail="There was an overlap with fingercuts! Please try again to generate dxf.")
    except HTTPException as e:
        raise e
    except Exception as e:
        raise HTTPException(status_code=500, detail="Error detecting reference battery! Please try again with a clearer image.")

@app.post("/predict2")
async def predict2_api(
    file: UploadFile = File(...),
    enable_fillet: str = Form(..., regex="^(On|Off)$"),
    fillet_value_mm: float = Form(...)
):
    """
    Middle predict: image + fillet toggle + fillet value → overlay, outlines, mask, DXF
    DXF naming format: DXF_DD-MM-YYYY_fillet-value_mm.dxf
    """
    session_id = str(uuid.uuid4())
    try:
        img_bytes = await file.read()
        image = np.array(Image.open(io.BytesIO(img_bytes)).convert("RGB"))
    except Exception:
        raise HTTPException(400, "Invalid image upload")

    try:
        start = timeit.default_timer()
        out_img, outlines, dxf_path, mask = predict_middle(
            image, enable_fillet, fillet_value_mm
        )
        elapsed = timeit.default_timer() - start
        print(f"[{session_id}] predict2 in {elapsed:.2f}s")

        return save_and_build_urls(
            session_id=session_id,
            output_image=out_img,
            outlines=outlines,
            dxf_path=dxf_path,
            mask=mask,
            endpoint_type="predict2",
            fillet_value=fillet_value_mm
        )

    except ReferenceBoxNotDetectedError:
        raise HTTPException(status_code=400, detail="Error detecting reference battery! Please try again with a clearer image.")
    except FingerCutOverlapError:
        raise HTTPException(status_code=400, detail="There was an overlap with fingercuts! Please try again to generate dxf.")
    except HTTPException as e:
        raise e
    except Exception as e:
        raise HTTPException(status_code=500, detail="Error detecting reference battery! Please try again with a clearer image.")


@app.post("/predict3")
async def predict3_api(
    file: UploadFile = File(...),
    enable_fillet: str = Form(..., regex="^(On|Off)$"),
    fillet_value_mm: float = Form(...),
    enable_finger_cut: str = Form(..., regex="^(On|Off)$")
):
    """
    Full predict: image + fillet toggle/value + finger-cut toggle → overlay, outlines, mask, DXF
    DXF naming format: DXF_DD-MM-YYYY_fillet-value_mm_fingercut-On|Off.dxf
    """
    session_id = str(uuid.uuid4())
    try:
        img_bytes = await file.read()
        image = np.array(Image.open(io.BytesIO(img_bytes)).convert("RGB"))
    except Exception:
        raise HTTPException(400, "Invalid image upload")

    try:
        start = timeit.default_timer()
        out_img, outlines, dxf_path, mask = predict_full(
            image, enable_fillet, fillet_value_mm, enable_finger_cut
        )
        elapsed = timeit.default_timer() - start
        print(f"[{session_id}] predict3 in {elapsed:.2f}s")

        return save_and_build_urls(
            session_id=session_id,
            output_image=out_img,
            outlines=outlines,
            dxf_path=dxf_path,
            mask=mask,
            endpoint_type="predict3",
            fillet_value=fillet_value_mm,
            finger_cut=enable_finger_cut
        )

    except ReferenceBoxNotDetectedError:
        raise HTTPException(status_code=400, detail="Error detecting reference battery! Please try again with a clearer image.")
    except FingerCutOverlapError:
        raise HTTPException(status_code=400, detail="There was an overlap with fingercuts! Please try again to generate dxf.")
    except HTTPException as e:
        raise e
    except Exception as e:
        raise HTTPException(status_code=500, detail="Error detecting reference battery! Please try again with a clearer image.")


@app.post("/update")
async def update_files(
    output_image: UploadFile = File(...),
    outlines_image: UploadFile = File(...),
    mask_image: UploadFile = File(...),
    dxf_file: UploadFile = File(...)
):
    session_id = str(uuid.uuid4())
    update_dir = os.path.join(UPDATES_DIR, session_id)
    os.makedirs(update_dir, exist_ok=True)

    try:
        upload_map = {
            "output_image":  output_image,
            "outlines_image": outlines_image,
            "mask_image":     mask_image,
            "dxf_file":       dxf_file,
        }
        urls = {}
        for key, up in upload_map.items():
            fn = up.filename
            path = os.path.join(update_dir, fn)
            with open(path, "wb") as f:
                shutil.copyfileobj(up.file, f)
            urls[key] = f"{BASE_URL}/updates/{session_id}/{fn}"

        return {"session_id": session_id, "uploaded": urls}

    except Exception as e:
        raise HTTPException(500, f"Update failed: {e}")


from fastapi import Response

@app.get("/health")
def health():
    return Response(content="OK", status_code=200)


if __name__ == "__main__":
    import uvicorn
    port = int(os.environ.get("PORT", 8080))
    print(f"Starting FastAPI server on 0.0.0.0:{port}...")
    uvicorn.run(app, host="0.0.0.0", port=port)