###############################################################################
# Dockerfile — HF Spaces (Docker SDK, CPU-only, port 7860)
#
# Two-stage build:
#   Stage A  — node:22-slim  — compile React/Vite frontend (VITE_API_BASE="")
#   Stage B  — python:3.11-slim — runtime with all ML models baked in
#
# HF Spaces rules enforced:
#   • Non-root user uid=1000  (useradd -m -u 1000 user)
#   • All writable runtime dirs under /tmp  (STORAGE_DIR, JOBS_DB_DIR)
#   • Model cache baked at build time under /home/user (readable by uid 1000)
#   • Port 7860
#
# Layer-size note
# ---------------
# Model downloads and COPY instructions run AFTER "USER user" so files are
# created owned by uid 1000 directly.  The only root-owned step after mkdir is
# facexlib (must write to site-packages).  This eliminates the old
# "RUN chown -R 1000:1000 /home/user" layer that duplicated ~1.5 GB of model
# weights via Docker's copy-on-write mechanism.
###############################################################################


###############################################################################
# Stage A — Build the React / Vite frontend
###############################################################################
FROM node:22-slim AS frontend-build

WORKDIR /build/frontend

# Install npm deps (cached until package*.json changes)
COPY frontend/package.json frontend/package-lock.json ./
RUN npm ci

# Build with VITE_API_BASE="" so every fetch() in the compiled JS uses a
# relative path (same-origin), not hardcoded http://localhost:8000.
COPY frontend/ .
RUN VITE_API_BASE="" npm run build
# Output lives at /build/frontend/dist


###############################################################################
# Stage B — Python 3.11 runtime
###############################################################################
FROM python:3.11-slim

# ── OS dependencies ───────────────────────────────────────────────────────────
# libgl1 + libglib2.0-0           : opencv-python requirement
# libsm6 libxext6 libxrender1     : X11 stubs required by cv2 on headless Linux
# wget ca-certificates            : download model weights in subsequent RUN steps
RUN apt-get update \
    && apt-get install -y --no-install-recommends \
        libgl1 \
        libglib2.0-0 \
        libsm6 \
        libxext6 \
        libxrender1 \
        wget \
        ca-certificates \
    && rm -rf /var/lib/apt/lists/*

# ── Non-root user (HF Spaces hard requirement) ────────────────────────────────
RUN useradd -m -u 1000 user

# ── Python packages ───────────────────────────────────────────────────────────
# Step 1: CPU-only torch + torchvision from the dedicated wheel index.
#   Versions pinned to match the local dev environment (torch 2.8.0, tv 0.23.0).
#   Installing from pytorch.org/whl/cpu guarantees no CUDA binary is pulled.
# Step 2: rest of requirements.txt.  pip finds torch>=2.2.0 already satisfied
#   by the CPU wheel above and skips it; all other deps resolve normally.
COPY requirements.txt /tmp/requirements.txt
RUN pip install --no-cache-dir \
        "torch==2.8.0" "torchvision==0.23.0" \
        --index-url https://download.pytorch.org/whl/cpu \
    && pip install --no-cache-dir -r /tmp/requirements.txt

# ── facexlib weights (must run as root — writes to site-packages) ─────────────
# URLs verified against facexlib 0.3.0 source:
#   recognition/__init__.py : recognition_arcface_ir_se50.pth  → v0.1.0
#   __init__.py             : detection_Resnet50_Final.pth     → v0.1.0
#   __init__.py             : parsing_parsenet.pth             → v0.2.2
#
# FaceRestoreHelper.__init__ triggers the download of detection + parsing models
# automatically via load_file_from_url.  ArcFace is called separately because
# facerestore_engine.py loads it independently (raises FileNotFoundError if absent).
RUN python -c "\
import torch; \
from facexlib.utils.face_restoration_helper import FaceRestoreHelper; \
from facexlib.utils.misc import load_file_from_url; \
import facexlib; \
from pathlib import Path; \
weights_dir = str(Path(facexlib.__file__).parent / 'weights'); \
print('facexlib weights dir:', weights_dir); \
print('Pre-downloading RetinaFace (detection_Resnet50_Final.pth, 104 MB) and ParseNet (parsing_parsenet.pth, 81 MB)...'); \
FaceRestoreHelper(upscale_factor=1, face_size=512, crop_ratio=(1, 1), det_model='retinaface_resnet50', save_ext='png', use_parse=True, device=torch.device('cpu')); \
print('RetinaFace + ParseNet ready.'); \
print('Downloading ArcFace IR-SE50 (recognition_arcface_ir_se50.pth, 167 MB)...'); \
load_file_from_url(url='https://github.com/xinntao/facexlib/releases/download/v0.1.0/recognition_arcface_ir_se50.pth', model_dir='facexlib/weights', progress=True, save_dir=weights_dir); \
print('All facexlib weights downloaded.')"

# ── Pre-create model/cache dirs and chown while they are EMPTY (cheap) ────────
# chown on empty directories has zero copy-on-write cost.
# All subsequent writes (model downloads, COPY) happen as uid 1000 directly.
RUN mkdir -p \
        /home/user/app/ml_models/birefnet \
        /home/user/app/backend/ml_models/codeformer \
        /home/user/.cache/huggingface \
        /home/user/.paddleocr \
    && chown -R 1000:1000 \
        /home/user/app \
        /home/user/.cache \
        /home/user/.paddleocr

# ── Switch to uid 1000 — all subsequent RUN steps write as the target user ────
USER user

# ─────────────────────────────────────────────────────────────────────────────
# MODEL DOWNLOADS  (all happen at build time — zero network I/O at startup)
# Running as uid 1000 so files are created with correct ownership directly —
# no chown layer needed at the end.
# Layer order: stable downloads first so code changes don't invalidate them.
# ─────────────────────────────────────────────────────────────────────────────

# ── RealESRGAN_x4plus.pth  (64 MB) ───────────────────────────────────────────
# Official release: https://github.com/xinntao/Real-ESRGAN/releases/tag/v0.1.0
# Destination matches UPSCALE_MODEL_PATH = BASE_DIR / "ml_models" / "RealESRGAN_x4plus.pth"
RUN wget --tries=3 --no-verbose --show-progress \
        -O /home/user/app/ml_models/RealESRGAN_x4plus.pth \
        "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth"

# ── codeformer.pth  (360 MB) ─────────────────────────────────────────────────
# Official release: https://github.com/sczhou/CodeFormer/releases/tag/v0.1.0
# Destination matches CODEFORMER_MODEL_PATH = backend/ml_models/codeformer/codeformer.pth
RUN wget --tries=3 --no-verbose --show-progress \
        -O /home/user/app/backend/ml_models/codeformer/codeformer.pth \
        "https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth"

# ── BiRefNet  (~424 MB, HuggingFace Hub) ─────────────────────────────────────
# Uses snapshot_download (no model instantiation) so it is reliable in a
# headless build environment.  The cache_dir matches BIREFNET_MODEL_CACHE in
# config.py (BASE_DIR / "ml_models" / "birefnet").  trust_remote_code=True is
# NOT needed here — snapshot_download just copies files.  At runtime,
# from_pretrained(..., trust_remote_code=True) reads the same cache.
RUN python -c "\
from huggingface_hub import snapshot_download; \
print('Downloading BiRefNet from HuggingFace Hub (~424 MB)...'); \
snapshot_download(repo_id='ZhengPeng7/BiRefNet', cache_dir='/home/user/app/ml_models/birefnet'); \
print('BiRefNet download complete.')"

# ── PaddleOCR models  (~200 MB, text detection/recognition/orientation) ───────
# HOME is set to /home/user so models land in /home/user/.paddleocr, which is
# exactly where the runtime process will look for them.
# Same PaddleOCR constructor args as ocr_service.py: use_textline_orientation,
# lang="en", enable_mkldnn=False (safe on Linux CPU, avoids OneDNN issues).
RUN HOME=/home/user python -c "\
from paddleocr import PaddleOCR; \
print('Pre-warming PaddleOCR (first-run model download ~200 MB)...'); \
PaddleOCR(use_textline_orientation=True, lang='en', enable_mkldnn=False); \
print('PaddleOCR pre-warm complete.')"

# ─────────────────────────────────────────────────────────────────────────────
# SOURCE & COMMITTED WEIGHTS  (change more often than model downloads → last)
# --chown=user:user ensures uid 1000 owns all copied files without a separate
# chown layer (avoids copy-on-write duplication of large files).
# ─────────────────────────────────────────────────────────────────────────────

# ── Backend source (.dockerignore excludes codeformer.pth — already downloaded)
COPY --chown=user:user backend/ /home/user/app/backend/

# ── lama.pth  (197 MB) ───────────────────────────────────────────────────────
# No reliable public download URL exists for the big-lama TorchScript export.
# This file MUST be committed to the repository via Git LFS.
# Register with:  git lfs track "ml_models/lama.pth"
# Destination matches MODEL_PATH = BASE_DIR / "ml_models" / "lama.pth"
COPY --chown=user:user ml_models/lama.pth /home/user/app/ml_models/lama.pth

# ── Built React app from Stage A ─────────────────────────────────────────────
# Destination matches the path computed in main.py:
#   Path(__file__).resolve().parent.parent.parent / "frontend" / "dist"
#   = /home/user/app/frontend/dist   (when __file__ is backend/app/main.py)
COPY --chown=user:user --from=frontend-build /build/frontend/dist /home/user/app/frontend/dist/

# ── Runtime environment variables ─────────────────────────────────────────────
# Every variable here has a local-safe default in the Python code; these
# overrides are ONLY active inside the container (start.bat sets none of them).
ENV HOME=/home/user \
    PORT=7860 \
    SERVE_FRONTEND=true \
    PUBLIC_BASE_URL="" \
    ALLOWED_ORIGINS="*" \
    STORAGE_DIR=/tmp/storage \
    JOBS_DB_DIR=/tmp/data \
    HF_HOME=/home/user/.cache/huggingface \
    HF_HUB_DISABLE_PROGRESS_BARS=1

EXPOSE 7860

# Working directory mirrors start.bat: `cd backend && uvicorn app.main:app`
# "app.main:app" resolves to /home/user/app/backend/app/main.py
WORKDIR /home/user/app/backend

# No --reload (dev-only).  --host 0.0.0.0 required for HF Spaces to route traffic.
CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]