morpho.py

"""Neutone Morpho client — hosted audio-to-audio transform.

Wraps `https://dev.neutone.ai:32475` — the cloud equivalent of the Neutone
Max/AU plugin. Each call is one full file in / one full file out (no
streaming endpoint). 30 transform models exposed via `/models/info`.

Auth: `X-API-Key` header. Key resolution order:
  1. `session_key` arg (HF Spaces — held in gr.State)
  2. `NEUTONE_API_KEY` env var
  3. `~/.neutone_key` (chmod 600, desktop only — None on HF Spaces)

Pricing: per-input-second per-model. Each model has a 1000-hour
(3,600,000-second) budget tracked server-side and returned in
`X-Inference-Seconds-Used` / `X-Inference-Seconds-Limit` response
headers on every `/process_audio` call.

Reference: memory note `pj-neutone-morpho.md`.
"""
from __future__ import annotations

import json
import os
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import requests


BASE_URL = os.environ.get("NEUTONE_BASE_URL", "https://dev.neutone.ai:32475")

# HF Spaces detection — same rule as wallet.py. On Spaces the key MUST NOT
# be read from a disk file (the filesystem is multi-user). It must arrive
# via the NEUTONE_API_KEY environment variable set on the Space, or be
# threaded through gr.State for per-session user-supplied keys.
IS_HF_SPACE = bool(os.environ.get("SPACE_ID"))

KEY_PATH = Path.home() / ".neutone_key"


class MorphoError(RuntimeError):
    """Raised on any Morpho-API call failure. Keeps the user-facing message
    short — the full diagnostic goes to stderr in the call site."""


def _stored_key() -> str | None:
    """Reads ~/.neutone_key on desktop. Returns None on HF Spaces — the
    key must come from env or the session-scoped UI value there."""
    if IS_HF_SPACE:
        return None
    try:
        return KEY_PATH.read_text().strip() or None
    except (FileNotFoundError, PermissionError):
        return None


def get_key(session_key: str | None = None) -> str | None:
    """Resolve a key for outbound /models/info or /process_audio calls.

    HF-Spaces hardening (Codex P1): on Spaces (`SPACE_ID` env set) we
    **skip the NEUTONE_API_KEY env-var fallback entirely**, even if the
    deploy admin set it. The key was issued to us under confidence; we
    must not auto-connect every visitor on the shared deploy key.
    Spaces visitors paste their own key into the connect band; the value
    lives in gr.State for the session only. Desktop and self-hosted
    deploys still resolve in the documented env / file order."""
    if session_key and session_key.strip():
        return session_key.strip()
    if IS_HF_SPACE:
        return None
    env = os.environ.get("NEUTONE_API_KEY", "").strip()
    if env:
        return env
    return _stored_key()


def _auth_headers(session_key: str | None = None) -> dict[str, str]:
    key = get_key(session_key)
    if not key:
        raise MorphoError(
            "Neutone key not configured. Set NEUTONE_API_KEY env var, "
            "or write a key to ~/.neutone_key (chmod 600)."
        )
    return {"X-API-Key": key}


# ---------------------------------------------------------------------------
# /models/info — fetched once per process and cached. 30 models, ~55 kB JSON.
# Cache invalidates on key change so a different session can refresh.
# ---------------------------------------------------------------------------

_MODELS_CACHE: dict[str, list[dict[str, Any]]] = {}


def models(session_key: str | None = None, *, timeout: float = 15.0,
           force_refresh: bool = False) -> list[dict[str, Any]]:
    """GET /models/info → list of model dicts (see pj-neutone-morpho.md
    for the per-model schema). Cached per-key for the lifetime of the
    Python process."""
    key = get_key(session_key) or ""
    if not force_refresh and key in _MODELS_CACHE:
        return _MODELS_CACHE[key]
    r = requests.get(
        f"{BASE_URL}/models/info",
        headers=_auth_headers(session_key),
        timeout=timeout,
    )
    if r.status_code == 401:
        raise MorphoError("Neutone API rejected the key (401). Verify "
                          "NEUTONE_API_KEY / ~/.neutone_key matches the "
                          "value the Neutone team issued.")
    if not r.ok:
        raise MorphoError(f"HTTP {r.status_code} from /models/info: "
                          f"{r.text[:200]}")
    try:
        data = r.json()
    except Exception as e:
        raise MorphoError(f"/models/info returned non-JSON: {e}") from e
    ms = data.get("models") if isinstance(data, dict) else data
    if not isinstance(ms, list):
        raise MorphoError(f"/models/info returned unexpected shape: {type(data).__name__}")
    _MODELS_CACHE[key] = ms
    return ms


def model_by_name(name: str, session_key: str | None = None) -> dict[str, Any] | None:
    for m in models(session_key=session_key):
        if m.get("model_name") == name or m.get("internal_name") == name:
            return m
    return None


# ---------------------------------------------------------------------------
# /process_audio — multipart upload, audio back. Wall-clock is 3-5×
# realtime warm; 7-10× cold. Output is stereo 44.1 kHz WAV regardless of
# input channel count.
# ---------------------------------------------------------------------------

@dataclass
class TransformResult:
    path: str
    bytes: int
    wall_s: float
    model_name: str
    p1: float
    p2: float
    p3: float
    p4: float
    input_seconds: float | None       # x-audio-duration-seconds header
    seconds_used: float | None        # x-inference-seconds-used header
    seconds_limit: float | None       # x-inference-seconds-limit header
    content_type: str

    def quota_pct(self) -> float | None:
        # Codex P1: original `is None or self.seconds_limit` was inverted —
        # returned None when both values were present, raised TypeError when
        # used was None. Now: only compute when BOTH are real numbers.
        if self.seconds_used is None or not self.seconds_limit:
            return None
        try:
            return 100.0 * self.seconds_used / self.seconds_limit
        except (TypeError, ZeroDivisionError):
            return None


def transform(
    in_wav: str | Path,
    *,
    model_name: str,
    p1: float = 50.0,
    p2: float = 50.0,
    p3: float = 50.0,
    p4: float = 50.0,
    out_path: str | Path,
    session_key: str | None = None,
    timeout: float = 240.0,
) -> TransformResult:
    """POST /process_audio. Clamps p1..p4 to [0, 100] (server-side range).

    Writes the response body to `out_path` and returns a TransformResult
    with the quota headers so the UI can update an inference-seconds pill.

    Raises MorphoError on auth / quota / network / non-audio response.
    """
    in_path = Path(in_wav)
    if not in_path.is_file():
        raise MorphoError(f"input file not found: {in_path}")
    # 20 MB cap (per Neutone docs).
    size = in_path.stat().st_size
    if size > 20 * 1024 * 1024:
        raise MorphoError(
            f"input is {size/1024/1024:.1f} MB; Neutone cap is 20 MB. "
            "Shorten the file or downsample before re-trying."
        )

    def _clamp(v: float) -> float:
        try:
            f = float(v)
        except (TypeError, ValueError):
            return 50.0
        return max(0.0, min(100.0, f))

    params = {
        "model_name": model_name,
        "p1": _clamp(p1), "p2": _clamp(p2), "p3": _clamp(p3), "p4": _clamp(p4),
    }
    headers = _auth_headers(session_key)
    t0 = time.perf_counter()
    with in_path.open("rb") as fh:
        files = {"file": (in_path.name, fh, "audio/wav")}
        data = {"params": json.dumps(params)}
        r = requests.post(
            f"{BASE_URL}/process_audio",
            headers=headers, files=files, data=data,
            timeout=timeout,
        )
    wall = time.perf_counter() - t0

    if r.status_code == 401:
        raise MorphoError("Neutone API rejected the key (401).")
    if r.status_code == 429:
        # Per-model inference-seconds quota exhausted (1000h budget).
        raise MorphoError(
            f"Neutone model '{model_name}' is out of inference seconds "
            "(429). Try a different model or wait for quota to reset."
        )
    if not r.ok:
        snippet = r.text[:200].replace("\n", " ")
        raise MorphoError(f"HTTP {r.status_code} from /process_audio: {snippet}")

    ctype = (r.headers.get("content-type") or "").lower()
    if not ctype.startswith("audio/"):
        # Server should always return audio/* on a 200. If we got something
        # else (e.g. JSON error wrapped in 200), bail loudly so we don't
        # write garbage to the crate.
        raise MorphoError(
            f"expected audio/* response, got '{ctype}' · "
            f"first 120 bytes: {r.content[:120]!r}"
        )

    out_path = Path(out_path)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_bytes(r.content)

    def _f(h: str) -> float | None:
        v = r.headers.get(h)
        try:
            return float(v) if v is not None else None
        except (TypeError, ValueError):
            return None

    return TransformResult(
        path=str(out_path),
        bytes=len(r.content),
        wall_s=round(wall, 2),
        model_name=model_name,
        p1=params["p1"], p2=params["p2"], p3=params["p3"], p4=params["p4"],
        input_seconds=_f("x-audio-duration-seconds"),
        seconds_used=_f("x-inference-seconds-used"),
        seconds_limit=_f("x-inference-seconds-limit"),
        content_type=ctype,
    )


# ---------------------------------------------------------------------------
# UI helpers — turn the per-model schema into the bits a Gradio panel needs.
# ---------------------------------------------------------------------------

def model_choices_grouped(session_key: str | None = None) -> list[tuple[str, str]]:
    """Return `[(display_label, model_name), …]` ordered by category then
    name. Display label is `{CATEGORY} · {model_name}` so a flat dropdown
    is browseable without nested grouping (Gradio v6 Dropdown can't nest)."""
    try:
        ms = models(session_key=session_key)
    except MorphoError:
        return []
    rows = sorted(ms, key=lambda m: (m.get("model_category") or "zzz",
                                      m.get("model_name") or ""))
    return [
        (f"{(m.get('model_category') or '?').upper():<10} · {m['model_name']}",
         m["model_name"])
        for m in rows
    ]


def param_labels(model_name: str, session_key: str | None = None) -> dict[str, dict[str, Any]]:
    """`{p1: {name, description, default}, …}` for a model. Falls back to
    generic labels if the model is missing or the params dict is empty."""
    m = model_by_name(model_name, session_key=session_key)
    fallback = {f"p{i}": {"name": f"Macro {i}", "description": "",
                           "default": 50.0} for i in range(1, 5)}
    if not m:
        return fallback
    params = m.get("parameters") or {}
    out = {}
    for k in ("p1", "p2", "p3", "p4"):
        p = params.get(k) or {}
        out[k] = {
            "name": p.get("name") or f"Macro {k[1]}",
            "description": p.get("description") or "",
            "default": float(p.get("default_value", 50.0)),
        }
    return out


def preset_choices(model_name: str, session_key: str | None = None) -> list[dict[str, Any]]:
    """List of presets for a model: `[{name, p1, p2, p3, p4, tag}, …]`.
    Empty list if the model has no presets configured."""
    m = model_by_name(model_name, session_key=session_key)
    if not m:
        return []
    return list(m.get("presets_minimal") or [])