#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
DepthAnything V2 – einfacher HTTP-Dienst (FastAPI) für Depthmap-Inferenz.

- Endpunkt: POST /depth
  Request (JSON):
    {
      "image_url": "https://…/bild.jpg",   # optional
      "image_base64": "<BASE64>",          # optional, wird priorisiert
      "encoder": "vitb" | "vits" | "vitl", # optional, default: "vitb"
      "grayscale": true | false,           # optional, default: true
      "input_size": 518                    # optionale Kantenlänge (Skalierung vor Inferenz)
    }
  Response:
    - PNG Binary (image/png) bei Erfolg
    - JSON {"error": "..."} bei Fehler

Voraussetzungen:
- Repo: https://github.com/DepthAnything/Depth-Anything-V2
- Checkpoints liegen in DEPTH_REPO_DIR/checkpoints/
- venv aktiv, uvicorn installiert

Start (Entwicklung):
  uvicorn depth_server:app --host 0.0.0.0 --port 5005

Systemd siehe bereitgestellte Unit.
"""

import io
import os
import sys
import base64
import json
import cv2
import numpy as np
import torch
import requests
from typing import Optional
from PIL import Image
from fastapi import FastAPI, Response
from pydantic import BaseModel

# Pfad zu Depth-Anything-V2 einbinden
DEPTH_REPO_DIR = "/opt/depthanything/Depth-Anything-V2"
sys.path.append(DEPTH_REPO_DIR)

from depth_anything_v2.dpt import DepthAnythingV2  # noqa: E402

app = FastAPI()

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

MODEL_CFG = {
    "vits": {"encoder": "vits", "features": 64,  "out_channels": [48, 96, 192, 384]},
    "vitb": {"encoder": "vitb", "features": 128, "out_channels": [96, 192, 384, 768]},
    "vitl": {"encoder": "vitl", "features": 256, "out_channels": [256, 512, 1024, 1024]},
}

# Passe ggf. die Checkpoint-Pfade an deine Downloads an
CHECKPOINTS = {
    "vits": os.path.join(DEPTH_REPO_DIR, "checkpoints", "depth_anything_v2_vits.pth"),
    "vitb": os.path.join(DEPTH_REPO_DIR, "checkpoints", "depth_anything_v2_vitb.pth"),
    "vitl": os.path.join(DEPTH_REPO_DIR, "checkpoints", "depth_anything_v2_vitl.pth"),
}

MODELS = {}  # Cache geladener Modelle


def get_model(encoder: str) -> DepthAnythingV2:
    enc = encoder if encoder in MODEL_CFG else "vitb"
    if enc not in MODELS:
        m = DepthAnythingV2(**MODEL_CFG[enc])
        m.load_state_dict(torch.load(CHECKPOINTS[enc], map_location="cpu"))
        MODELS[enc] = m.to(DEVICE).eval()
    return MODELS[enc]


class DepthReq(BaseModel):
    image_url: Optional[str] = None
    image_base64: Optional[str] = None  # VR kann Base64 mitsenden
    encoder: str = "vitb"
    grayscale: bool = True
    input_size: int = 518               # Kantenlänge zum Runterskalieren vor Inferenz


@app.get("/health")
def health():
    return {"status": "ok", "device": DEVICE}


@app.post("/depth")
def depth(req: DepthReq):
    # 1) Bild laden: priorisiere Base64, sonst URL
    try:
        if req.image_base64:
            image_bytes = base64.b64decode(req.image_base64)
            image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        elif req.image_url:
            r = requests.get(req.image_url, timeout=25)
            r.raise_for_status()
            image = Image.open(io.BytesIO(r.content)).convert("RGB")
        else:
            return Response(
                status_code=400,
                content=json.dumps({"error": "no image provided (image_url or image_base64)"}).encode(),
                media_type="application/json",
            )
    except Exception as e:
        return Response(
            status_code=400,
            content=json.dumps({"error": f"image load failed: {str(e)}"}).encode(),
            media_type="application/json",
        )

    # PIL RGB -> OpenCV BGR
    arr = np.array(image)[:, :, ::-1]  # HxWx3

    # 2) Vor-Inferenz skalieren, um VRAM zu schonen (Seitenverhältnis erhalten)
    try:
        h0, w0 = arr.shape[:2]
        s = req.input_size / max(h0, w0)
        if s < 1.0:
            new_w = max(1, int(w0 * s))
            new_h = max(1, int(h0 * s))
            arr_resized = cv2.resize(arr, (new_w, new_h), interpolation=cv2.INTER_AREA)
        else:
            arr_resized = arr
    except Exception as e:
        return Response(
            status_code=500,
            content=json.dumps({"error": f"resize failed: {str(e)}"}).encode(),
            media_type="application/json",
        )

    # 3) Inferenz
    try:
        model = get_model(req.encoder)
        with torch.no_grad():
            depth_small = model.infer_image(arr_resized)  # HxW float32 numpy
    except torch.cuda.OutOfMemoryError as e:
        return Response(
            status_code=500,
            content=json.dumps({"error": f"cuda oom: {str(e)}"}).encode(),
            media_type="application/json",
        )
    except Exception as e:
        return Response(
            status_code=500,
            content=json.dumps({"error": f"inference failed: {str(e)}"}).encode(),
            media_type="application/json",
        )

    # 4) Wieder auf Originalgröße hochskalieren (optional, bessere Passung)
    try:
        if depth_small.shape[0] != h0 or depth_small.shape[1] != w0:
            depth = cv2.resize(depth_small, (w0, h0), interpolation=cv2.INTER_CUBIC)
        else:
            depth = depth_small
    except Exception as e:
        return Response(
            status_code=500,
            content=json.dumps({"error": f"upsample failed: {str(e)}"}).encode(),
            media_type="application/json",
        )

    # 5) Normalisieren auf [0,255] und ggf. einfärben
    try:
        d = depth.astype(np.float32)
        d -= d.min()
        if d.max() > 0:
            d /= d.max()
        d = (d * 255.0).clip(0, 255).astype(np.uint8)
        if not req.grayscale:
            d = cv2.applyColorMap(d, cv2.COLORMAP_TURBO)
    except Exception as e:
        return Response(
            status_code=500,
            content=json.dumps({"error": f"normalize failed: {str(e)}"}).encode(),
            media_type="application/json",
        )

    # 6) PNG encodieren und binär zurückgeben
    ok, png = cv2.imencode(".png", d)
    if not ok:
        return Response(
            status_code=500,
            content=json.dumps({"error": "png encode failed"}).encode(),
            media_type="application/json",
        )

    return Response(content=png.tobytes(), media_type="image/png")