backend/app/api/upload.py

"""
upload.py — POST /upload

Receives a raw image file (multipart/form-data), normalises EXIF rotation,
saves it as steps/0/original.jpg in a new job folder, creates the job record
with step 0 already completed, and returns the job_id.

Step 0 design
-------------
Step 0 is created at upload time with status="completed".  Its original_image
and result_image both point to the same uploaded file.  This allows /generate
to always read `current_step.result_image` as the base image for the next step
without special-casing step 0.

Why EXIF normalisation at upload time?
---------------------------------------
JPEG files from phones often store pixels in a raw (unrotated) orientation and
record the intended display rotation in EXIF tag 0x0112.  Viewers that honour
EXIF show the image "right way up"; OpenCV does NOT.  If we skip normalisation:
  1. Frontend shows the correctly-oriented image (browser respects EXIF).
  2. Backend loads raw unrotated pixels via cv2.imread().
  3. Brush coordinates (drawn against the oriented preview) are misaligned.

Fix: apply PIL's ImageOps.exif_transpose() at upload time, which physically
rotates the pixel data to match EXIF orientation and strips the tag.
"""

import uuid
from datetime import datetime, timezone
from pathlib import Path

import cv2
import numpy as np
from fastapi import APIRouter, HTTPException, UploadFile, File
from PIL import Image, ImageOps
import io

from app.config import STORAGE_DIR
from app.jobs_store import create_job, add_step
from app.models.schemas import Job, StepRecord
from app.utils import job_folder_name

router = APIRouter()


@router.post("")
async def upload_image(file: UploadFile = File(...)):
    """Accept an image upload, normalise EXIF orientation, persist to storage.

    Returns
    -------
    {"job_id": str, "width": int, "height": int}
    """
    # --- Basic content-type guard -------------------------------------------
    if not file.content_type.startswith("image/"):
        raise HTTPException(
            status_code=415,
            detail=f"Expected an image file, got content-type '{file.content_type}'.",
        )

    # --- Read upload into memory --------------------------------------------
    raw_bytes = await file.read()
    if not raw_bytes:
        raise HTTPException(status_code=400, detail="Uploaded file is empty.")

    # --- EXIF-correct orientation via Pillow --------------------------------
    try:
        pil_img = Image.open(io.BytesIO(raw_bytes))
        pil_img = ImageOps.exif_transpose(pil_img)
        pil_img = pil_img.convert("RGB")
    except Exception as exc:
        raise HTTPException(status_code=422, detail=f"Could not decode image: {exc}")

    # --- Convert PIL RGB → OpenCV BGR numpy array ---------------------------
    img_bgr: np.ndarray = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
    height, width = img_bgr.shape[:2]

    # --- Create job folder and step-0 subfolder -----------------------------
    job_id     = str(uuid.uuid4())
    created_at = datetime.now(timezone.utc)
    folder     = job_folder_name(job_id, created_at)
    step0_dir  = STORAGE_DIR / "watermark" / folder / "steps" / "0"
    step0_dir.mkdir(parents=True, exist_ok=True)

    original_path = step0_dir / "original.jpg"
    ok = cv2.imwrite(str(original_path), img_bgr)
    if not ok:
        raise HTTPException(
            status_code=500,
            detail="Failed to write uploaded image to storage.",
        )

    # --- Create job record (empty steps list, current_step=0) ---------------
    # Pass created_at explicitly so create_job() doesn't generate a different
    # timestamp — we need the exact same value to recompute the folder name later.
    job = Job(id=job_id, current_step=0, job_type="watermark", created_at=created_at)
    create_job(job)

    # --- Create step 0: status="completed", result_image = original ----------
    # result_image = original_image for step 0 (no inference, upload IS the result).
    step0 = StepRecord(
        step_number=0,
        original_image=str(original_path),
        result_image=str(original_path),
        status="completed",
    )
    add_step(job_id, step0)

    return {"job_id": job_id, "width": width, "height": height}