backend/app/api/upscale.py

"""
upscale.py — POST /upscale/upload  and  GET /upscale/result/{job_id}

INDEPENDENT of the LaMa/watermark-removal pipeline.
Do not import from jobs_store.py, schemas.py, generate.py, inpaint_engine.py,
or any other watermark-removal module.

Design
------
Unlike watermark removal (which separates upload from inference because it needs
a user-drawn mask between the two steps), upscaling has no user input after the
image is chosen — so upload and inference kick-off are combined into a single
POST /upscale/upload endpoint.

Flow
----
1. POST /upscale/upload (multipart file)
   a. Validate content-type.
   b. EXIF-correct and convert to RGB via Pillow (same reason as /upload).
   c. Assign UUID job_id; create storage/upscale/{job_id}/.
   d. Save EXIF-corrected image as original.jpg.
   e. Persist UpscaleJob(status="pending") to upscale_jobs.json.
   f. Immediately advance status to "processing" and enqueue BackgroundTask.
   g. Return {"job_id": job_id} immediately (inference runs in background).

2. GET /upscale/result/{job_id}  (polling)
   Returns {"job_id", "status", "result_image_url", "error_message"}.
   result_image_url is a /storage/... URL served by the existing StaticFiles
   mount in main.py — no extra route needed.

3. run_upscale_job (background worker)
   Runs after the POST response is sent; has no access to the Request object.
   Receives the engine as an explicit argument (same pattern as run_inpaint_job).
   Updates job status to "completed" or "failed"; never raises.

Storage layout (kept entirely separate from watermark-removal storage)
----------------------------------------------------------------------
  storage/
    upscale/
      {job_id}/
        original.jpg   — EXIF-corrected upload
        result.jpg     — 4× upscaled output (written on completion)

These are served as:
  /storage/upscale/{job_id}/original.jpg
  /storage/upscale/{job_id}/result.jpg
via the existing app.mount("/storage", StaticFiles(...)) in main.py.
"""

from __future__ import annotations

import io
import logging
import math
import time
import traceback
import uuid
from concurrent.futures import ThreadPoolExecutor, TimeoutError as _FuturesTimeout
from datetime import datetime, timezone
from pathlib import Path

import cv2
import numpy as np
from fastapi import APIRouter, BackgroundTasks, File, Form, HTTPException, Request, UploadFile
from PIL import Image, ImageOps

from app.config import STORAGE_DIR, UPSCALE_TILE_THRESHOLD, UPSCALE_TILE_SIZE
from app.utils import job_folder_name
from app.upscale_jobs_store import UpscaleJob, create_job, get_job, update_job
from app.services.upscale_engine import BaseUpscaleEngine

logger = logging.getLogger(__name__)

router = APIRouter()

# ---------------------------------------------------------------------------
# Job-level timeout
# ---------------------------------------------------------------------------
# Hard ceiling on how long engine.upscale() is allowed to run.
#
# Justification (33 s/tile measured baseline on this machine):
#   A 2048×1367 image → 4×3 = 12 tiles → expected ~6.6 min at baseline.
#   The reported runaway job hit 36+ min (≥5× baseline) before manual kill.
#   30 minutes gives 4.5× headroom over the 12-tile baseline — generous
#   enough for moderate thermal throttling — while cutting off a stuck job
#   before it can run indefinitely.  At normal speed a practical worst-case
#   upload (~20 tiles for a very large image) takes ~11 min baseline;
#   30 min = 2.7× headroom, still safe.
#
#   Note: on timeout the *upscale thread* becomes a zombie — Python cannot
#   forcibly terminate it.  It will continue consuming CPU in the background
#   until the computation finishes or the server restarts.  The job record is
#   immediately marked "failed" so the user is not left waiting indefinitely.
_UPSCALE_TIMEOUT_SECS: int = 30 * 60   # 30 minutes

# Upscale jobs live in a dedicated sub-directory of storage/ so they can
# never collide with watermark-removal job folders (which use storage/{uuid}/steps/).
_UPSCALE_STORAGE = STORAGE_DIR / "upscale"


# ---------------------------------------------------------------------------
# Post-processing helpers
# ---------------------------------------------------------------------------

def _denoise(img: np.ndarray) -> np.ndarray:
    """Non-local means denoising (BGR uint8 in → BGR uint8 out).

    Averages similar patches across the image to reduce random noise and JPEG
    compression artifacts while preserving edges better than a simple blur.
    Parameters (h=10, hColor=10, templateWindowSize=7, searchWindowSize=21)
    are the established standard defaults for general-purpose color denoising.

    Applied BEFORE upscaling so we process the smaller original image — cheaper
    and removes artifacts before the model can amplify them 4×.
    """
    return cv2.fastNlMeansDenoisingColored(
        img, None,
        h=10, hColor=10,
        templateWindowSize=7, searchWindowSize=21,
    )


def _sharpen(img: np.ndarray) -> np.ndarray:
    """Mild unsharp mask (RGB uint8 in → RGB uint8 out).

    Intentionally conservative (amount=0.4, sigma=1.0): the 4× Real-ESRGAN
    output already has an AI-oversharpened look on organic textures (observed
    during project testing — grass/foliage artifacts), so an aggressive sharpen
    would compound that problem. GaussianBlur and the arithmetic are
    channel-independent so no BGR conversion is needed.

    Applied AFTER upscaling, BEFORE any output_scale downscale step.
    """
    blurred = cv2.GaussianBlur(img, (0, 0), sigmaX=1.0)
    img_f   = img.astype(np.float32)
    blur_f  = blurred.astype(np.float32)
    return np.clip(img_f + 0.4 * (img_f - blur_f), 0, 255).astype(np.uint8)


# ---------------------------------------------------------------------------
# Background worker
# ---------------------------------------------------------------------------

def run_upscale_job(
    job_id:        str,
    folder_name:   str,
    original_path: Path,
    engine:        BaseUpscaleEngine,
    output_scale:  str  = "4x",
    denoise:       bool = False,
    sharpen:       bool = False,
) -> None:
    """Run Real-ESRGAN upscaling and save result; update job status.

    Called by FastAPI BackgroundTasks after the /upscale/upload response is sent.
    Has no access to the Request object — engine is passed as an explicit arg.
    Updates status to "completed" or "failed"; never raises.

    Pipeline order
    --------------
    original → [denoise] → engine.upscale() → [sharpen] → [output_scale resize] → save

    denoise is applied on the small original image (cheaper; removes JPEG
    artifacts before the model can amplify them 4×).  sharpen is applied on the
    full 4× output before any downscale, so the mild enhancement survives
    LANCZOS resampling.

    output_scale
    ------------
    "4x"  — save the full 4x output from the engine (current default behaviour).
    "2x"  — downsample the 4x output by exactly 50% in each dimension using
            LANCZOS before saving.
    "1.5x"— downsample to 1.5× the original (resize factor 3/8 from 4× output).
    """
    job_dir = _UPSCALE_STORAGE / folder_name
    result_path = job_dir / "result.jpg"
    t_job_start = time.monotonic()

    try:
        pil_img = Image.open(original_path).convert("RGB")
        img_w, img_h = pil_img.size

        # Log job dimensions and expected tile count so we can rule out a
        # tile-count bug when diagnosing slow jobs.
        if max(img_w, img_h) >= UPSCALE_TILE_THRESHOLD:
            n_tiles = math.ceil(img_w / UPSCALE_TILE_SIZE) * math.ceil(img_h / UPSCALE_TILE_SIZE)
            logger.info(
                "Upscale job %s: %dx%d → tiled mode (%d tiles expected).  "
                "scale=%s denoise=%s sharpen=%s  timeout=%ds",
                job_id[:8], img_w, img_h, n_tiles,
                output_scale, denoise, sharpen, _UPSCALE_TIMEOUT_SECS,
            )
        else:
            logger.info(
                "Upscale job %s: %dx%d → whole-image mode.  "
                "scale=%s denoise=%s sharpen=%s",
                job_id[:8], img_w, img_h, output_scale, denoise, sharpen,
            )

        # --- Optional: denoise BEFORE upscale (on the small original) ----------
        if denoise:
            img_bgr = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
            img_bgr = _denoise(img_bgr)
            pil_img = Image.fromarray(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB))

        # --- Upscale with job-level timeout ------------------------------------
        # Run engine.upscale() in a worker thread so we can enforce a hard time
        # ceiling.  If the timeout fires, the job is immediately marked "failed"
        # so the user gets feedback.  The underlying thread becomes a daemon and
        # will be collected when the server restarts.
        with ThreadPoolExecutor(max_workers=1) as _pool:
            _future = _pool.submit(engine.upscale, pil_img)
            try:
                sr_img = _future.result(timeout=_UPSCALE_TIMEOUT_SECS)
            except _FuturesTimeout:
                elapsed_min = (time.monotonic() - t_job_start) / 60
                logger.error(
                    "Upscale job %s TIMED OUT after %.1f min (limit %d min). "
                    "Background thread will continue until computation finishes.",
                    job_id[:8], elapsed_min, _UPSCALE_TIMEOUT_SECS // 60,
                )
                update_job(
                    job_id,
                    status        = "failed",
                    error_message = (
                        f"처리 시간이 너무 오래 걸려 중단되었습니다 "
                        f"({_UPSCALE_TIMEOUT_SECS // 60}분 초과). "
                        f"이미지 크기를 줄이거나 나중에 다시 시도해 주세요."
                    ),
                )
                return

        # --- Optional: sharpen AFTER upscale, BEFORE downscale -----------------
        if sharpen:
            img_arr = np.array(sr_img)       # PIL RGB → numpy RGB uint8
            img_arr = _sharpen(img_arr)       # unsharp mask (channel-independent)
            sr_img  = Image.fromarray(img_arr)

        if output_scale == "2x":
            w, h   = sr_img.size
            sr_img = sr_img.resize((w // 2, h // 2), Image.Resampling.LANCZOS)
        elif output_scale == "1.5x":
            # Target: 1.5× the original.  Engine always outputs 4×, so the
            # resize factor is 1.5/4 = 3/8 = 0.375.
            # round() gives the nearest integer; for the 671×406 test image:
            #   round(2684 × 3/8) = round(1006.5) = 1006
            #   round(1624 × 3/8) = round(609.0)  = 609
            # matching 671×1.5≈1006 and 406×1.5=609 as closely as integers allow.
            w, h   = sr_img.size
            sr_img = sr_img.resize(
                (round(w * 3 / 8), round(h * 3 / 8)),
                Image.Resampling.LANCZOS,
            )

        sr_img.save(str(result_path), format="JPEG", quality=95)

        # Build the URL path relative to the /storage mount.
        # result_path is an absolute path under STORAGE_DIR, so we compute the
        # portion after STORAGE_DIR to construct the URL fragment.
        rel = result_path.relative_to(STORAGE_DIR)
        elapsed = time.monotonic() - t_job_start
        logger.info("Upscale job %s completed in %.1fs", job_id[:8], elapsed)
        update_job(job_id, status="completed", result_image=f"/{rel.as_posix()}")

    except Exception:
        update_job(
            job_id,
            status="failed",
            error_message=traceback.format_exc(),
        )


# ---------------------------------------------------------------------------
# POST /upscale/upload
# ---------------------------------------------------------------------------

@router.post("/upload")
async def upscale_upload(
    file:             UploadFile,
    background_tasks: BackgroundTasks,
    request:          Request,
    output_scale:     str  = Form("4x"),
    denoise:          bool = Form(False),
    sharpen:          bool = Form(False),
) -> dict:
    """Accept an image, EXIF-correct it, persist it, and kick off upscaling.

    Form fields
    -----------
    file         : image file (multipart)
    output_scale : "4x" (default) | "2x" | "1.5x"
    denoise      : bool (default False) — apply NLM denoising before upscale
    sharpen      : bool (default False) — apply mild unsharp mask after upscale

    Returns ``{"job_id": <str>}`` immediately; upscaling runs in the background.
    Poll GET /upscale/result/{job_id} for completion.
    """
    # --- output_scale validation ---------------------------------------------
    if output_scale not in ("4x", "2x", "1.5x"):
        raise HTTPException(
            status_code=422,
            detail=f"output_scale must be '4x', '2x', or '1.5x', got '{output_scale}'.",
        )

    # --- Content-type guard --------------------------------------------------
    if not (file.content_type or "").startswith("image/"):
        raise HTTPException(
            status_code=415,
            detail=f"Expected an image file, got content-type '{file.content_type}'.",
        )

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

    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}")

    # --- Persist original image ----------------------------------------------
    job_id     = str(uuid.uuid4())
    created_at = datetime.now(timezone.utc)
    folder     = job_folder_name(job_id, created_at)
    job_dir    = _UPSCALE_STORAGE / folder
    job_dir.mkdir(parents=True, exist_ok=True)

    original_path = job_dir / "original.jpg"
    pil_img.save(str(original_path), format="JPEG", quality=95)

    # --- Create job record ---------------------------------------------------
    # Pass created_at explicitly (as ISO string, matching UpscaleJob.created_at: str)
    # so upscale_jobs_store.create_job() doesn't generate a different timestamp.
    orig_rel = original_path.relative_to(STORAGE_DIR)
    job = UpscaleJob(
        id             = job_id,
        original_image = f"/{orig_rel.as_posix()}",
        status         = "processing",
        output_scale   = output_scale,
        denoise        = denoise,
        sharpen        = sharpen,
        created_at     = created_at.isoformat(),
    )
    create_job(job)

    # --- Enqueue background upscaling ----------------------------------------
    engine: BaseUpscaleEngine = request.app.state.upscale_engine

    background_tasks.add_task(
        run_upscale_job,
        job_id        = job_id,
        folder_name   = folder,
        original_path = original_path,
        engine        = engine,
        output_scale  = output_scale,
        denoise       = denoise,
        sharpen       = sharpen,
    )

    return {"job_id": job_id}


# ---------------------------------------------------------------------------
# GET /upscale/result/{job_id}
# ---------------------------------------------------------------------------

@router.get("/result/{job_id}")
def upscale_result(job_id: str) -> dict:
    """Return the current status of an upscale job.

    Response fields:
        job_id           : str
        status           : "pending" | "processing" | "completed" | "failed"
        result_image_url : str | null  — /storage/upscale/{job_id}/result.jpg
        error_message    : str | null
    """
    job = get_job(job_id)
    if job is None:
        raise HTTPException(status_code=404, detail=f"Upscale job '{job_id}' not found.")

    return {
        "job_id":           job.id,
        "status":           job.status,
        "result_image_url": job.result_image,
        "error_message":    job.error_message,
    }