"""Hugging Face Inference Endpoints handler: a description (or an image, when a vision
model is reachable) in, ranked Vienna codes out.

Copied to the root of the model repository by ``i2vienna hf-export``. The repository
also carries ``scheme/`` and ``index/`` (parquet), ``image2vienna.json`` (which index
to serve) and a copy of the ``image2vienna`` package, so no PyPI release is required.

Request, text (the description of the figurative elements, stage two alone)::

    {"inputs": "Three stars above a crescent moon. Gold and blue.",
     "parameters": {"level": "section", "top_k": 5}}

Request, image (needs a describer the endpoint can reach, named by the environment
variable ``IMAGE2VIENNA_DESCRIBER``, e.g. ``hf:florence-community/Florence-2-base-ft``
to run Florence-2 on the endpoint itself; without the variable an image is refused)::

    {"inputs": {"image": "<base64 bytes>"}, "parameters": {"level": "division"}}

Response (one list per input; a single input yields a single list)::

    [{"code": "1.1.4", "padded": "01.01.04", "level": "section", "auxiliary": true,
      "score": 0.83, "similarity": 0.83, "title": "Three stars", "path": "..."}, ...]
"""

from __future__ import annotations

import base64
import json
import os
import sys
from pathlib import Path
from typing import Any

ALLOWED_PARAMS = {
    "level",
    "top_k",
    "gap",
    "auto_margin",
    "normalize",
    "chunking",
    "principal_only",
    "exclude_codes",
}


class EndpointHandler:
    def __init__(self, path: str = ""):
        root = Path(path or ".").resolve()
        if str(root) not in sys.path:
            sys.path.insert(0, str(root))  # vendored image2vienna package
        from image2vienna.classifier import ViennaClassifier

        cfg = json.loads((root / "image2vienna.json").read_text())
        self.clf = ViennaClassifier(
            cfg["edition"],
            lang=cfg["lang"],
            model=cfg["model"],
            notes=bool(cfg.get("notes", False)),
            describer=os.environ.get("IMAGE2VIENNA_DESCRIBER") or "fixed:/nonexistent",
            root=root,
        )
        self.has_describer = bool(os.environ.get("IMAGE2VIENNA_DESCRIBER"))
        _ = self.clf.index  # load eagerly so the first request is fast
        _ = self.clf.embedder

    def __call__(self, data: dict[str, Any]) -> list[dict] | list[list[dict]]:
        inputs = data.get("inputs", "")
        params = {k: v for k, v in (data.get("parameters") or {}).items() if k in ALLOWED_PARAMS}
        if "exclude_codes" in params:
            params["exclude_codes"] = tuple(params["exclude_codes"] or ())
        if isinstance(inputs, list):
            return [self._one(item, params) for item in inputs]
        return self._one(inputs, params)

    def _one(self, item: Any, params: dict[str, Any]) -> list[dict]:
        if isinstance(item, dict) and "image" in item:
            if not self.has_describer:
                raise ValueError(
                    "This endpoint has no vision model: send the description of the figurative "
                    "elements as text, or run the image2vienna package locally for images"
                )
            matches = self.clf.classify(base64.b64decode(item["image"]), **params)
        else:
            text = item["text"] if isinstance(item, dict) else str(item)
            matches = self.clf.classify_text(text, **params)
        return [
            {
                "code": m.code,
                "padded": m.padded,
                "level": m.level,
                "auxiliary": m.auxiliary,
                "score": round(m.score, 4),
                "similarity": round(m.similarity, 4),
                "title": m.title,
                "path": m.text,
            }
            for m in matches
        ]