"""Modal deployment of Apple SHARP (https://huggingface.co/apple/Sharp).

Single photo in -> 3D Gaussian splat PLY out, for the "upload your own image"
flow. SHARP is a ~1B-parameter feedforward model (2.8 GB checkpoint) that
predicts the splat in well under a second on a GPU, so this runs comfortably
on an L4.

Endpoints:
  - predict (POST): JSON {image_base64, filename?, build_lod?, prune_keep?, ...}
              -> JSON {ply_base64, rad_base64?, meta}. By default it returns the
              full ~1.18M-splat PLY AND its precomputed .rad LoD tree (built in
              the container via the bundled static build-lod binary), so callers
              never need a separate LoD step. prune_keep defaults to 0 (no
              prune): the .rad handles AR performance and the full PLY keeps
              fine detail. The meta carries counts, intrinsics, suggested plane
              depth, and timings.
  - warmup  (POST): fire-and-forget container boot.

License note: the SHARP model weights are under Apple's AMLR (research)
license — fine for a hackathon, check before commercial use.

Deploy: modal deploy modal_sharp.py
"""

import base64
import json
import os
import time
import traceback
from pathlib import Path
from tempfile import TemporaryDirectory

import modal
from pydantic import BaseModel

try:  # deploy-time convenience: pick up MODAL_REQUIRE_PROXY_AUTH etc. from .env
    from dotenv import load_dotenv

    load_dotenv()
except ImportError:
    pass

MODEL_REPO = "apple/Sharp"
CHECKPOINT_FILE = "sharp_2572gikvuh.pt"
GPU_CONFIG = os.environ.get("SHARP_GPU", "L4")
MIN_CONTAINERS = int(os.environ.get("SHARP_MIN_CONTAINERS", "0"))
SCALEDOWN_WINDOW = int(os.environ.get("SHARP_SCALEDOWN_WINDOW", "240"))
SERVICE_VERSION = "sharp-2026-06-13a-lod"
DEFAULT_PRUNE_KEEP = 0  # no prune: the .rad LoD tree handles AR perf, full PLY keeps detail
# See modal_minicpmo.py: leaked/guessed URLs must not be able to burn GPU time.
REQUIRES_PROXY_AUTH = os.environ.get("MODAL_REQUIRE_PROXY_AUTH", "1").lower() not in {"0", "false", ""}

BUILD_LOD_BIN = "/root/build-lod"

image = (
    modal.Image.debian_slim(python_version="3.13")
    .apt_install("git", "libgl1", "libglib2.0-0")
    .pip_install("torch", "torchvision")
    # gsplat (a sharp dependency) is only needed for --render; we never import
    # it, so its CUDA kernels are never JIT-compiled.
    .pip_install("git+https://github.com/apple/ml-sharp.git")
    .pip_install("huggingface_hub[hf_transfer]", "fastapi[standard]>=0.115")
    .env({"HF_HUB_ENABLE_HF_TRANSFER": "1"})
    .add_local_file("prune_ply.py", "/root/prune_ply.py")
    # Spark's LoD builder (the `build-lod` tool), compiled as a static musl
    # binary so it runs on the Debian image with no glibc/toolchain in the
    # image. Lets the worker produce a .rad LoD tree right after generating the
    # splat — no separate manual step. To refresh it, rebuild from the spark
    # repo: cargo build --release --target x86_64-unknown-linux-musl
    #   --manifest-path rust/build-lod/Cargo.toml
    .add_local_file("build-lod-linux", BUILD_LOD_BIN)
)

app = modal.App("myapp-sharp", image=image)
hf_cache = modal.Volume.from_name("hf-cache", create_if_missing=True)


class PredictRequest(BaseModel):
    image_base64: str
    filename: str = "photo.jpg"  # extension matters: EXIF focal + HEIC support
    prune_keep: int = DEFAULT_PRUNE_KEEP  # 0 disables pruning
    area_exponent: float = 0.0
    scale_boost: float = 1.05
    build_lod: bool = True  # also build the .rad LoD tree (the recommended path)
    return_ply: bool = True  # include the (large) raw PLY in the response; callers
    #                          that only need the .rad set this False to avoid
    #                          downloading hundreds of MB over the HTTP endpoint.


@app.cls(
    gpu=GPU_CONFIG,
    timeout=300,
    min_containers=MIN_CONTAINERS,
    scaledown_window=SCALEDOWN_WINDOW,
    volumes={"/root/.cache/huggingface": hf_cache},
    # Same GPU-snapshot setup as the MiniCPM-o worker: cold starts restore the
    # loaded model in seconds instead of re-loading the checkpoint.
    enable_memory_snapshot=True,
    experimental_options={"enable_gpu_snapshot": True},
)
class SharpWorker:
    @modal.enter(snap=True)
    def load(self):
        import sys

        import torch
        from huggingface_hub import hf_hub_download
        from sharp.models import PredictorParams, create_predictor

        sys.path.insert(0, "/root")  # for prune_ply

        started = time.perf_counter()
        self.container_started_at = time.time()
        checkpoint_path = hf_hub_download(repo_id=MODEL_REPO, filename=CHECKPOINT_FILE)
        state_dict = torch.load(checkpoint_path, weights_only=True)
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.predictor = create_predictor(PredictorParams())
        self.predictor.load_state_dict(state_dict)
        self.predictor.eval().to(self.device)
        self.model_load_seconds = time.perf_counter() - started
        print(
            f"SHARP loaded: service_version={SERVICE_VERSION}, gpu={GPU_CONFIG}, "
            f"device={self.device}, load_seconds={self.model_load_seconds:.3f}"
        )

    def _timings(self, request_started: float, extra: dict | None = None):
        timings = {
            "request_seconds": round(time.perf_counter() - request_started, 3),
            "model_load_seconds": round(getattr(self, "model_load_seconds", 0), 3),
            "model_source": MODEL_REPO,
            "service_version": SERVICE_VERSION,
            "container_age_seconds": round(time.time() - getattr(self, "container_started_at", time.time()), 3),
            "gpu": GPU_CONFIG,
        }
        if extra:
            timings.update(extra)
        return timings

    def _predict_gaussians(self, image, f_px: float):
        """Predict metric-space gaussians from an RGB array.

        Reimplements sharp.cli.predict.predict_image (Apple ml-sharp, AMLR
        license) so we never import sharp.cli — importing it pulls in the
        gsplat renderer, whose CUDA kernels would JIT-compile at import time.
        """
        import torch
        import torch.nn.functional as F  # noqa: N812 (torch convention)
        from sharp.utils.gaussians import unproject_gaussians

        internal_shape = (1536, 1536)
        image_pt = torch.from_numpy(image.copy()).float().to(self.device).permute(2, 0, 1) / 255.0
        _, height, width = image_pt.shape
        disparity_factor = torch.tensor([f_px / width]).float().to(self.device)

        image_resized = F.interpolate(
            image_pt[None],
            size=(internal_shape[1], internal_shape[0]),
            mode="bilinear",
            align_corners=True,
        )

        with torch.no_grad():
            gaussians_ndc = self.predictor(image_resized, disparity_factor)

        intrinsics = (
            torch.tensor(
                [
                    [f_px, 0, width / 2, 0],
                    [0, f_px, height / 2, 0],
                    [0, 0, 1, 0],
                    [0, 0, 0, 1],
                ]
            )
            .float()
            .to(self.device)
        )
        intrinsics_resized = intrinsics.clone()
        intrinsics_resized[0] *= internal_shape[0] / width
        intrinsics_resized[1] *= internal_shape[1] / height

        return unproject_gaussians(
            gaussians_ndc, torch.eye(4).to(self.device), intrinsics_resized, internal_shape
        )

    @modal.method()
    def ping(self) -> dict:
        """No-op used by /warmup to boot a container ahead of the first predict."""
        return {
            "service_version": SERVICE_VERSION,
            "gpu": GPU_CONFIG,
            "model_load_seconds": getattr(self, "model_load_seconds", None),
        }

    def _build_lod(self, ply_bytes: bytes) -> bytes:
        """Run Spark's build-lod on a PLY, returning the .rad bytes."""
        import os
        import stat
        import subprocess

        # add_local_file doesn't preserve the exec bit; set it once.
        if not os.access(BUILD_LOD_BIN, os.X_OK):
            os.chmod(BUILD_LOD_BIN, os.stat(BUILD_LOD_BIN).st_mode | stat.S_IEXEC)

        with TemporaryDirectory() as tmp:
            ply_path = Path(tmp) / "scene.ply"
            ply_path.write_bytes(ply_bytes)
            # build-lod writes "<stem>-lod.rad" next to the input.
            subprocess.run([BUILD_LOD_BIN, str(ply_path)], check=True, capture_output=True)
            return (Path(tmp) / "scene-lod.rad").read_bytes()

    @modal.method()
    def predict(self, request: PredictRequest) -> dict:
        """Photo -> {ply_base64, rad_base64?, meta}. Raises ValueError on bad input."""
        import numpy as np
        from sharp.utils import io as sharp_io
        from sharp.utils.gaussians import save_ply

        import prune_ply

        request_started = time.perf_counter()
        raw = base64.b64decode(request.image_base64)
        if not raw:
            raise ValueError("image_base64 is empty")

        with TemporaryDirectory() as tmp:
            # Keep the original extension: sharp_io.load_rgb reads the EXIF
            # focal length (and handles HEIC) based on the actual file.
            suffix = Path(request.filename or "photo.jpg").suffix or ".jpg"
            image_path = Path(tmp) / f"photo{suffix}"
            image_path.write_bytes(raw)

            preprocess_started = time.perf_counter()
            image, _, f_px = sharp_io.load_rgb(image_path)
            height, width = image.shape[:2]
            preprocess_seconds = time.perf_counter() - preprocess_started

            inference_started = time.perf_counter()
            gaussians = self._predict_gaussians(image, f_px)
            inference_seconds = time.perf_counter() - inference_started

            ply_path = Path(tmp) / "out.ply"
            save_ply(gaussians, f_px, (height, width), ply_path)
            ply_bytes = ply_path.read_bytes()

        prune_info = None
        if request.prune_keep and request.prune_keep > 0:
            prune_started = time.perf_counter()
            ply_bytes, prune_info = prune_ply.prune_bytes(
                ply_bytes,
                keep=request.prune_keep,
                area_exponent=request.area_exponent,
                scale_boost=request.scale_boost,
            )
            prune_info["prune_seconds"] = round(time.perf_counter() - prune_started, 3)

        # Build the LoD tree (the recommended runtime path: smooth AR +
        # streaming + full detail). Done here so callers never need a separate
        # build-lod step.
        rad_bytes = None
        lod_seconds = None
        if request.build_lod:
            lod_started = time.perf_counter()
            rad_bytes = self._build_lod(ply_bytes)
            lod_seconds = round(time.perf_counter() - lod_started, 3)

        # Median gaussian depth: what the viewer uses as the image-plane
        # distance (SOURCE_CAMERA.depth) when mapping selections back to pixels.
        header_lines, body_offset = prune_ply.read_header(ply_bytes)
        count = int(
            next(line for line in header_lines if line.startswith("element vertex ")).split()[-1]
        )
        z = np.frombuffer(
            ply_bytes, dtype=np.float32, count=count * len(prune_ply.VERTEX_PROPS), offset=body_offset
        ).reshape(count, len(prune_ply.VERTEX_PROPS))[:, 2]

        meta = {
            "splat_count": count,
            "image_width": width,
            "image_height": height,
            "intrinsics": {
                "fx": float(f_px),
                "fy": float(f_px),
                "cx": (width - 1) / 2.0,
                "cy": (height - 1) / 2.0,
            },
            "suggested_plane_depth": float(np.median(z)),
            "prune": prune_info,
            "timings": self._timings(
                request_started,
                {
                    "preprocess_seconds": round(preprocess_seconds, 3),
                    "inference_seconds": round(inference_seconds, 3),
                    "lod_seconds": lod_seconds,
                    "ply_bytes": len(ply_bytes),
                    "rad_bytes": len(rad_bytes) if rad_bytes else 0,
                },
            ),
        }
        # Keep the PLY only when asked for it, or as a fallback when no .rad was
        # built — otherwise the raw PLY (hundreds of MB at full resolution) bloats
        # the HTTP response and can stall callers downloading it over the gateway.
        include_ply = request.return_ply or not rad_bytes
        return {
            "ply_base64": base64.b64encode(ply_bytes).decode("ascii") if include_ply else None,
            "rad_base64": base64.b64encode(rad_bytes).decode("ascii") if rad_bytes else None,
            "meta": meta,
        }


@app.function()
@modal.fastapi_endpoint(method="POST", requires_proxy_auth=REQUIRES_PROXY_AUTH)
def warmup():
    """Fire-and-forget container boot (mirrors the MiniCPM-o warmup)."""
    call = SharpWorker().ping.spawn()
    return {"status": "warming", "call_id": call.object_id, "service_version": SERVICE_VERSION}


@app.function()
@modal.fastapi_endpoint(method="POST", requires_proxy_auth=REQUIRES_PROXY_AUTH)
def predict(request: PredictRequest):
    # JSON: {ply_base64, rad_base64?, meta}. (Base64 in JSON is convenient for a
    # not-yet-wired upload flow; if/when payload size matters, switch to a
    # binary multipart response.)
    from fastapi import HTTPException

    try:
        return SharpWorker().predict.remote(request)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    except Exception as exc:  # surface worker tracebacks to the caller
        raise HTTPException(status_code=500, detail=f"{exc}\n{traceback.format_exc()}") from exc


@app.local_entrypoint()
def test_local(
    image_id: str = "mybook",
    image_path: str = "best1.jpg",
    prune_keep: int = DEFAULT_PRUNE_KEEP,
    area_exponent: float = 0.0,
    scale_boost: float = 1.05,
    keep_ply: bool = False,
):
    """Generate a picture book from a photo.

    Writes the files the server needs, named <image_id>.*:
      <id>-lod.rad    the splat the viewer loads (LoD: smooth AR + streaming)
      <id>.meta.json  camera intrinsics / image size / depth for selections
      <id>.ply        only with --keep-ply (the .rad makes it redundant)

    Still need a <id>.caption and an <id>.jpg/.webp next to these for the book
    to appear. Example:
      modal run modal_sharp.py --image-id vatican --image-path vatican.jpg
    """
    data = Path(image_path).read_bytes()
    request = PredictRequest(
        image_base64=base64.b64encode(data).decode("ascii"),
        filename=Path(image_path).name,
        prune_keep=prune_keep,
        area_exponent=area_exponent,
        scale_boost=scale_boost,
        build_lod=True,
    )
    result = SharpWorker().predict.remote(request)
    meta = result["meta"]
    print(json.dumps(meta, indent=2))

    # Sidecar the server reads (replaces parsing the .ply for camera params).
    sidecar = {
        "width": meta["image_width"],
        "height": meta["image_height"],
        "fx": meta["intrinsics"]["fx"],
        "fy": meta["intrinsics"]["fy"],
        "cx": meta["intrinsics"]["cx"],
        "cy": meta["intrinsics"]["cy"],
        "plane_depth": meta["suggested_plane_depth"],
        "splat_count": meta["splat_count"],
    }
    Path(f"{image_id}.meta.json").write_text(json.dumps(sidecar, indent=2))
    print(f"wrote {image_id}.meta.json")

    if result["rad_base64"]:
        rad_bytes = base64.b64decode(result["rad_base64"])
        Path(f"{image_id}-lod.rad").write_bytes(rad_bytes)
        print(f"wrote {image_id}-lod.rad ({len(rad_bytes) / 1e6:.1f} MB)")
    else:
        print("WARNING: no .rad returned (build_lod failed?); pass --keep-ply to keep the splat")

    if keep_ply or not result["rad_base64"]:
        ply_bytes = base64.b64decode(result["ply_base64"])
        Path(f"{image_id}.ply").write_bytes(ply_bytes)
        print(f"wrote {image_id}.ply ({len(ply_bytes) / 1e6:.1f} MB)")