# -------------------------------------------------------------------
# This source file is available under the terms of the
# Pimcore Open Core License (POCL)
# Full copyright and license information is available in
# LICENSE.md which is distributed with this source code.
#
#  @copyright  Copyright (c) Pimcore GmbH (https://www.pimcore.com)
#  @license    Pimcore Open Core License (POCL)
# -------------------------------------------------------------------

"""Who may be loaded, and whose code may run.

Callers name the model in the request, so `model` is untrusted input reaching
`from_pretrained`. Two separate decisions follow from it, and fusing them is what makes the
service dangerous:

1. **May this id be loaded at all?** Anything that is not a well-formed Hub repo id is refused —
   `from_pretrained` also accepts a *local filesystem path*, so an unvalidated id lets a caller
   point the loader at the container's disk.
2. **May loading it execute the author's Python?** `trust_remote_code=True` means a repo whose
   `config.json` declares an `auto_map` gets its `modeling_*.py` imported, and importing is
   executing. Enabled for everyone, that is remote code execution by design.

Splitting them means any standard-architecture model still works for any caller with no operator
involvement (decision 1 says yes, decision 2 says no code) while custom-architecture models —
Nomic being the one this service actually needs — require the operator to trust that namespace
first.

Note the trust anchor is a *name*, and Hub namespaces can be renamed, released and re-registered,
so an entry keeps trusting whoever holds the name later. That is true of any name-based allowlist,
per-model ones included; only pinning a revision SHA avoids it. Review the list periodically.
"""

import os
import re
from typing import Optional

from fastapi import HTTPException

# Anchored, whole-string match — a prefix check is not enough. `from_pretrained` treats an
# existing local path as a directory to load from, so ids like `trusted-org/../../etc` must fail
# structurally rather than be caught by comparing the leading segment.
_REPO_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*(/[A-Za-z0-9][A-Za-z0-9._-]*)?$")

# Only namespaces whose models genuinely need custom code belong here — everything with an
# architecture transformers supports natively (CLIP, SigLIP, BERT, ViT, BLIP …) loads fine with
# remote code disabled, so adding an org buys nothing and costs trust. `nomic-ai` is the one this
# service depends on today. Override with the TRUSTED_CODE_ORGS env var (comma-separated), set
# either in docker-compose.yaml or as a Variable on the Hugging Face Space.
_DEFAULT_TRUSTED_CODE_ORGS = "nomic-ai,pimcore"


def _configured_orgs() -> frozenset:
    """Read the allowlist from the environment on each call so a test (or a restart) picks up a
    change without re-importing the module."""
    raw = os.environ.get("TRUSTED_CODE_ORGS", _DEFAULT_TRUSTED_CODE_ORGS)
    # Compared lowercase: Hub namespaces are unique case-insensitively, so folding here closes a
    # trivial `Nomic-AI` vs `nomic-ai` bypass rather than creating one.
    return frozenset(part.strip().lower() for part in raw.split(",") if part.strip())


def validate_model_id(model_name: str) -> None:
    """Refuse anything that is not a well-formed Hub repo id, before it reaches `from_pretrained`.

    Raises:
        HTTPException: 400, if the id is malformed or names something on the local filesystem.
    """
    if not _REPO_ID.match(model_name):
        raise HTTPException(
            status_code=400,
            detail=(
                f"Invalid model id '{model_name}'. Expected a Hugging Face repo id such as "
                "'namespace/model' or a canonical model name."
            ),
        )

    # Belt and braces on top of the regex: a syntactically valid id that happens to name a real
    # path (`src`, `tests`) would still be loaded from disk rather than the Hub.
    if os.path.exists(model_name):
        raise HTTPException(status_code=400, detail=f"Invalid model id '{model_name}'.")


def namespace_of(model_name: str) -> Optional[str]:
    """The owning namespace, or None for a canonical model id (`bert-base-uncased`, `gpt2`) which
    has no owner segment."""
    if "/" not in model_name:
        return None
    return model_name.split("/", 1)[0]


def allows_remote_code(model_name: str) -> bool:
    """Whether loading this model may execute code shipped in its repo.

    Canonical (namespace-less) ids are always False: they are Hub-owned and carry no custom code,
    so there is no namespace to have trusted.
    """
    namespace = namespace_of(model_name)
    if namespace is None:
        return False
    return namespace.lower() in _configured_orgs()