# coding=utf-8
"""
Carte de profondeur d'une image, pour l'extension « Gray Iso-Layers ».

Le Python d'Inkscape ne peut pas charger onnxruntime : le calcul est confie
au script `depth_estimate.py`, lance avec le Python d'un environnement a part
(cree par `deploy.ps1 -Depth`). Le resultat est garde sur disque, car
Inkscape relance l'extension a chaque « Appliquer » : seul le premier calcul
d'une image coute l'inference.

Dossier de donnees (hors du dossier des extensions, qu'Inkscape parcourt en
entier au demarrage) : `venv/`, `models/` et `cache/`.

Sans dependance a inkex. Dependances : numpy et Pillow (fournis avec Inkscape).
"""

import os
import subprocess
import sys
import tempfile

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))

# Modeles proposes : nom (parametre depth_model du .inx) -> fichier de poids.
MODELS = {
    "midas_small": "midas_v21_small_256.onnx",
    "depth_anything_v2_small": "depth_anything_v2_small.onnx",
}

CACHE_LIMIT = 30   # cartes gardees sur disque
TIMEOUT = 120      # secondes laissees au calcul


class DepthUnavailable(Exception):
    """La carte de profondeur ne peut pas etre calculee (environnement ou
    modele absent, echec du calcul). Le message en donne la cause."""


def data_dir():
    """Dossier de l'environnement Python, des modeles et du cache."""
    if os.name == "nt":
        base = os.environ.get("LOCALAPPDATA") or os.path.expanduser("~")
    else:
        base = os.path.join(os.path.expanduser("~"), ".local", "share")
    return os.path.join(base, "GrayIsoLayers")


def find_python():
    """Python qui execute `depth_estimate.py` : celui de la variable
    GRAY_ISO_LAYERS_PYTHON, sinon celui de l'environnement du dossier de
    donnees."""
    chosen = os.environ.get("GRAY_ISO_LAYERS_PYTHON")
    if chosen:
        return chosen
    if os.name == "nt":
        return os.path.join(data_dir(), "venv", "Scripts", "python.exe")
    return os.path.join(data_dir(), "venv", "bin", "python")


def depth_field(picture, digest, model="midas_small"):
    """Profondeur inverse de `picture` (image PIL RGB) : tableau 2D, plus la
    valeur est haute, plus le point est proche. Sa taille est celle de la
    sortie du modele, pas celle de l'image.

    `digest` identifie l'image (empreinte de ses octets) : avec le nom du
    modele, c'est la cle du cache. Leve `DepthUnavailable` si le calcul est
    impossible.
    """
    if model not in MODELS:
        raise DepthUnavailable("unknown model '{}'".format(model))
    cache = os.path.join(data_dir(), "cache")
    target = os.path.join(cache, "{}-{}.npy".format(digest, model))
    if os.path.isfile(target):
        try:
            depth = np.load(target, allow_pickle=False)
            os.utime(target)  # carte resservie : la derniere a etre purgee
            return depth
        except (OSError, ValueError):
            pass  # fichier abime : on recalcule

    python = find_python()
    script = os.path.join(HERE, "depth_estimate.py")
    weights = os.path.join(data_dir(), "models", MODELS[model])
    for path in (python, script, weights):
        if not os.path.isfile(path):
            raise DepthUnavailable("missing file: {}".format(path))

    # Inkscape pointe PYTHONPATH et PYTHONHOME vers son propre Python : ils
    # feraient charger ses modules par celui de l'environnement.
    env = {key: value for key, value in os.environ.items()
           if key.upper() not in ("PYTHONPATH", "PYTHONHOME")}
    with tempfile.TemporaryDirectory() as folder:
        source = os.path.join(folder, "image.png")
        output = os.path.join(folder, "depth.npy")
        picture.save(source)
        try:
            done = subprocess.run(
                [python, script, source, output, weights, "--model", model],
                env=env, capture_output=True, encoding="utf-8", errors="replace",
                timeout=TIMEOUT,
                # Pas de fenetre de console sous Windows.
                creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0))
        except (OSError, subprocess.TimeoutExpired) as error:
            raise DepthUnavailable(str(error))
        if done.returncode != 0 or not os.path.isfile(output):
            lines = (done.stderr or "").strip().splitlines()
            raise DepthUnavailable(lines[-1] if lines else
                                   "exit code {}".format(done.returncode))
        try:
            depth = np.load(output, allow_pickle=False)
        except (OSError, ValueError) as error:
            raise DepthUnavailable(str(error))

    try:
        os.makedirs(cache, exist_ok=True)
        np.save(target, depth, allow_pickle=False)
        _prune(cache)
    except OSError as error:
        # Le cache n'est qu'un confort : la carte est calculee, on la rend.
        print("gray_iso_layers_depth: {}".format(error), file=sys.stderr)
    return depth


def _prune(cache, limit=None):
    """Ne garde que les `limit` cartes les plus recemment servies."""
    limit = CACHE_LIMIT if limit is None else limit
    files = [os.path.join(cache, name) for name in os.listdir(cache)
             if name.endswith(".npy")]
    files.sort(key=os.path.getmtime, reverse=True)
    for path in files[limit:]:
        os.remove(path)