"""Dedicated ML "amenity" endpoints (`/ml/*`).

These are named, *preordained* routes for the FIXED set of media-intelligence
tasks (ASR, summarize, keywords, embeddings, similarity, vision, text-to-image).
They're a static part of hugpy's default amenities — necessarily distinct from
the open-ended LLM chat surface — so they get stable named endpoints AND a
reserved worker pool, instead of riding the generic `/prompt {task}` verb (which
stays for general/back-compat use).

DESIGN — deliberately THIN. Each route forces its `task` and a default ML worker
pool, then hands off to `execute_prompt`. No heavy/contrived ML code is imported
here: the dispatch lazy-imports each task's implementation only when it actually
runs — ideally on the reserved `ml`-pool worker (and on a CLEAN seam, NOT the
contrived `utils` helpers, which drag unhospitable imports). If no `ml`-pool
worker is registered, the pool reservation falls back to LOCAL (central
in-process) — never the general LLM pool.

Pool: `HUGPY_ML_POOL` (default ``ml``; empty disables the default). An API-key
bound pool or an explicit per-request `pool` still wins over the default.
"""
from __future__ import annotations

import inspect
import os

from ..functions import *  # get_bp, request, jsonify (via abstract_flask), …

ml_bp, logger = get_bp("ml_bp", __name__)

# The fixed media-intelligence amenity set: route name -> dispatch task key.
ML_TASKS = {
    "transcribe": "automatic-speech-recognition",
    "summarize":  "text-summarization",
    "keywords":   "keyword-extraction",
    "embed":      "feature-extraction",
    "similarity": "sentence-similarity",
    "vision":     "image-text-to-text",
    "imagine":    "text-to-image",
}


# Per-amenity dependency probe -> the extra that provides it. Reported by GET /ml
# so the UI shows enabled tools vs a clean "enable with abstract_hugpy_dev[X]"
# affordance — WITHOUT importing the heavy dep (find_spec only; phone-clean).
ML_DEP = {
    "transcribe": ("whisper", "audio"),
    "summarize":  ("transformers", "transformers"),
    "keywords":   ("keybert", "keywords"),
    "embed":      ("sentence_transformers", "embed"),
    "similarity": ("sentence_transformers", "embed"),
    "vision":     ("llama_cpp", "engine"),
    "imagine":    ("diffusers", "imagegen"),
}


def _have(mod: str) -> bool:
    import importlib.util
    try:
        return importlib.util.find_spec(mod) is not None
    except (ImportError, ValueError):
        return False


def _ml_pool() -> str:
    """The reserved worker pool these endpoints route to (env-overridable).
    Empty string disables the default (normal pool resolution then applies)."""
    return (os.environ.get("HUGPY_ML_POOL", "ml") or "").strip()


def _await_sync(value):
    """Drive execute_prompt's (possibly) awaitable result on the shared loop."""
    if not inspect.isawaitable(value):
        return value
    from abstract_hugpy_dev._platform import async_runtime
    return async_runtime.run(value)


def _result_payload(result) -> dict:
    for attr in ("model_dump", "to_dict", "dict"):
        fn = getattr(result, attr, None)
        if callable(fn):
            try:
                return fn()
            except TypeError:
                continue
    return {"text": str(result)}


def _run_ml(task: str):
    body = request.get_json(silent=True) or {}
    # Underscore-prefixed keys are internal routing controls; never client-set.
    kwargs = {k: v for k, v in body.items()
              if not k.startswith("_") and v is not None}
    kwargs["task"] = task  # the route fixes the task (any client `task` ignored)

    # Effective pool: API-key-bound pool / explicit override wins; otherwise the
    # reserved ML pool — so ML work never competes on the general LLM pool.
    try:
        from ..functions.chat.streaming import _resolve_request_pool
        eff = _resolve_request_pool(kwargs.get("pool")) or _ml_pool()
    except Exception:
        eff = (kwargs.get("pool") or _ml_pool())
    if eff:
        kwargs["pool"] = eff
    else:
        kwargs.pop("pool", None)

    from ..functions.imports import execute_prompt  # late import dodges circulars
    try:
        result = _await_sync(execute_prompt(**kwargs))
    except (KeyError, ValueError, TypeError, FileNotFoundError) as exc:
        # builder/validation errors — caller's to fix.
        return jsonify({"ok": False, "error": str(exc).strip("'\"")}), 400
    except Exception as exc:
        logger.exception("ml task %s failed", task)
        return jsonify({"ok": False, "error": f"{type(exc).__name__}: {exc}"}), 500

    payload = _result_payload(result)
    payload.setdefault("ok", not payload.get("error"))
    return jsonify(payload)


def _make_view(task: str):
    def _view():
        return _run_ml(task)
    return _view


# Register one POST route per named amenity (closure-factory avoids late-binding).
for _name, _task in ML_TASKS.items():
    ml_bp.add_url_rule(
        f"/ml/{_name}",
        endpoint=f"ml_{_name}",
        view_func=_make_view(_task),
        methods=["POST"],
    )


@ml_bp.route("/ml", methods=["GET"])
def ml_index():
    """Discovery + capability: each amenity's task, whether its dependency is
    installed, and the extra that enables it. Uses find_spec only — it never
    imports the heavy dep, so this stays cheap and phone-clean."""
    endpoints = {}
    for name, task in ML_TASKS.items():
        mod, extra = ML_DEP.get(name, (None, None))
        endpoints[f"/ml/{name}"] = {
            "task": task,
            "ready": _have(mod) if mod else True,
            "extra": extra,
        }
    return jsonify({"endpoints": endpoints, "pool": _ml_pool() or None})
