# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors

"""Canonical Function values exchanged with LanceDB Enterprise services.

These immutable models contain client/wire state only. Catalog persistence,
environment bake, secret resolution, and execution are owned by Sophon.
``RefreshColumnResult`` is also the backend-neutral result of a local
expression-backed refresh job.
"""

from __future__ import annotations

import ast
import builtins
import base64
import functools
import hashlib
import importlib
import inspect
import symtable
import json
import linecache
import math
import re
import sys
import textwrap
import types
import warnings
from collections.abc import Mapping, Sequence
from datetime import date, datetime
from typing import (
    Annotated,
    Any,
    Callable,
    Literal,
    Optional,
    Union,
    get_args,
    get_origin,
    get_type_hints,
    overload,
)

import pyarrow as pa
from pydantic import (
    AfterValidator,
    BaseModel,
    ConfigDict,
    Field,
    conint,
    field_validator,
    model_validator,
)

from .schema import blob as _blob_field
from .schema import is_blob_v2_field as _is_blob_v2_field
from ._blob_reader import BlobReader
from .secrets import EnvVarSecret

_Int32 = conint(strict=True, ge=-(2**31), le=2**31 - 1)
_UInt32 = conint(strict=True, ge=0, le=2**32 - 1)
_UInt64 = conint(strict=True, ge=0, le=2**64 - 1)


def _validate_gpu_wire_marker(value: Any) -> bool:
    if value is not True:
        raise ValueError("runtime.gpu must be true")
    return True


def _normalize_gpu_marker(value: bool) -> Optional[bool]:
    if not isinstance(value, bool):
        raise ValueError("gpu must be a boolean")
    return True if value else None


class _FrozenDict(dict):
    def _immutable(self, *args, **kwargs):
        raise TypeError("remote canonical values are immutable")

    __setitem__ = _immutable
    __delitem__ = _immutable
    clear = _immutable
    pop = _immutable
    popitem = _immutable
    setdefault = _immutable
    update = _immutable

    def __ior__(self, other):
        self._immutable()


def _freeze_value(value):
    if isinstance(value, Mapping):
        return _FrozenDict({key: _freeze_value(child) for key, child in value.items()})
    if isinstance(value, (list, tuple)):
        return tuple(_freeze_value(child) for child in value)
    return value


def _validate_literal(value):
    if isinstance(value, float):
        raise ValueError(
            "floating-point Function literals are not part of the Slice 1 "
            "canonical wire contract"
        )
    if isinstance(value, int) and not isinstance(value, bool):
        if not -(2**63) <= value <= 2**64 - 1:
            raise ValueError(
                "Function integer literal is outside the canonical JSON range"
            )
    elif isinstance(value, Mapping):
        for child in value.values():
            _validate_literal(child)
    elif isinstance(value, (list, tuple)):
        for child in value:
            _validate_literal(child)
    return value


def _known_wire_value(value):
    if isinstance(value, _RemoteValue):
        return value._known_dict()
    if isinstance(value, Mapping):
        return {key: _known_wire_value(child) for key, child in value.items()}
    if isinstance(value, (list, tuple)):
        return [_known_wire_value(child) for child in value]
    return value


class _RemoteValue(BaseModel):
    model_config = ConfigDict(extra="ignore", frozen=True)

    @model_validator(mode="after")
    def _freeze_mappings(self):
        for name, value in self.__dict__.items():
            object.__setattr__(self, name, _freeze_value(value))
        return self

    @classmethod
    def from_json(cls, payload: str):
        return cls.model_validate_json(payload)

    def _known_dict(self) -> dict[str, Any]:
        known = {}
        for name, field in self.__class__.model_fields.items():
            value = getattr(self, name)
            if value is None:
                continue
            if not field.is_required():
                default_factory = field.default_factory
                if default_factory is not None and value == default_factory():
                    continue
                if default_factory is None and value == field.default:
                    continue
            known[name] = _known_wire_value(value)
        return known

    def _copy(self, *, update: Mapping[str, Any]):
        update = {name: _freeze_value(value) for name, value in update.items()}
        return self.model_copy(update=update)

    def to_canonical_json(self) -> str:
        return json.dumps(
            self._known_dict(),
            ensure_ascii=False,
            allow_nan=False,
            sort_keys=True,
            separators=(",", ":"),
        )


class _OpenRemoteValue(_RemoteValue):
    """Forward-readable value whose extras stay out of canonical encoding."""

    model_config = ConfigDict(extra="allow", frozen=True)

    def _unknown_field_names(self) -> set[str]:
        return set((self.__pydantic_extra__ or {}).keys())


class FunctionArtifact(_RemoteValue):
    """Content-addressed Python artifact identity."""

    kind: str
    digest: str
    entrypoint: str


class FunctionArtifactContent(_RemoteValue):
    """Encoded artifact bytes uploaded during remote registration."""

    encoding: str
    data: str


class PythonAdapterSpec(_RemoteValue):
    """Internal scalar-callable to Arrow-batch adapter selection."""

    kind: str
    version: _UInt32


class FunctionArtifactRequest(_RemoteValue):
    """Source artifact uploaded while registering a Function."""

    kind: str
    digest: str
    entrypoint: str
    content: FunctionArtifactContent
    adapter: PythonAdapterSpec


class FunctionParameter(_RemoteValue):
    name: str
    arrow_type: str
    nullable: bool
    blob_access: Optional[Literal["bytes", "reader"]] = None


class FunctionResultField(_OpenRemoteValue):
    name: str
    arrow_type: str
    nullable: bool


class FunctionOutput(_OpenRemoteValue):
    """Scalar or ordered named-struct output; unknown kinds remain decodable."""

    kind: str
    arrow_type: Optional[str] = None
    nullable: Optional[bool] = None
    fields: tuple[FunctionResultField, ...] = ()


class SecretReference(_RemoteValue):
    """Where a Secret lives, carried as its parts rather than as one string.

    A joined id would need a delimiter, and a delimiter has to be excluded from
    every name and segment forever, agreed on by both sides, and re-agreed each
    time either grows a new way to be configured. Naming the parts settles all
    of that: nothing here is parsed, so nothing can parse two ways.
    """

    name: str
    namespace_path: tuple[str, ...] = ()


class SecretBinding(_RemoteValue):
    """How a Secret reaches the Function that binds it.

    One list rather than a field per delivery mode: a binding is the concept,
    and how it arrives is a property of one. ``kind`` is open, so a binding a
    newer service introduces decodes here instead of failing the whole
    FunctionVersion.
    """

    kind: str
    variable: Optional[str] = None
    secret_ref: Optional[SecretReference] = None


class FunctionSignature(_RemoteValue):
    """Ordered inputs, output, and initialization fields of a Function.

    ``initialization`` lists the fields of the single row each Function
    instance is created with. Values come from each binding, not from the
    Function version; the tuple is empty when the Function takes none.
    """

    inputs: tuple[FunctionParameter, ...]
    output: FunctionOutput
    initialization: tuple[FunctionParameter, ...] = ()


class PythonEnvironmentSpec(_RemoteValue):
    """One Sophon-managed Python environment source."""

    kind: str
    packages: tuple[str, ...] = ()
    channels: tuple[str, ...] = ()
    path: Optional[str] = None
    modules: tuple[str, ...] = ()
    image: Optional[str] = None


class PythonRuntimeSpec(_RemoteValue):
    """Remote runtime definition with environment values.

    V1 supports ``kind="python"``. Newer runtime kinds remain readable, while
    their unknown payload fields are intentionally not retained by the client.
    """

    kind: str
    python_version: Optional[str] = None
    environment: Optional[PythonEnvironmentSpec] = None
    env: Optional[Mapping[str, str]] = None
    gpu: Optional[bool] = None

    @model_validator(mode="before")
    @classmethod
    def _discard_unknown_runtime_payload(cls, value):
        if isinstance(value, Mapping):
            kind = value.get("kind")
            if isinstance(kind, str) and kind not in {"python", "python_v2"}:
                return {"kind": kind}
        return value

    @field_validator("gpu", mode="before")
    @classmethod
    def _validate_gpu_marker(cls, value):
        if value is None:
            return None
        return _validate_gpu_wire_marker(value)

    @model_validator(mode="after")
    def _validate_runtime_kind(self):
        if self.kind == "python":
            if self.python_version is None:
                raise ValueError("python runtime requires python_version")
            if self.environment is None:
                raise ValueError("python runtime requires environment")
            if self.gpu is not None:
                raise ValueError("python runtime with gpu requires kind='python_v2'")
        elif self.kind == "python_v2":
            if self.python_version is None:
                raise ValueError("python_v2 runtime requires python_version")
            if self.environment is None:
                raise ValueError("python_v2 runtime requires environment")
            if self.gpu is None:
                raise ValueError("python_v2 runtime requires gpu")
        else:
            object.__setattr__(self, "python_version", None)
            object.__setattr__(self, "environment", None)
            object.__setattr__(self, "env", None)
            object.__setattr__(self, "gpu", None)
        return self


class FunctionImage(_RemoteValue):
    """A complete OCI Function image identified by its exact manifest digest."""

    manifest_digest: str
    descriptor: Mapping[str, Any]
    source: bool


def _validate_object_version(value: str) -> str:
    if int(value) > 2**64 - 1:
        raise ValueError("Function version exceeds uint64")
    return value


_ObjectVersion = Annotated[
    str,
    Field(strict=True, pattern=r"^[1-9][0-9]*$"),
    AfterValidator(_validate_object_version),
]


class FunctionVersion(_RemoteValue):
    """A pinned object revision, independent of its executable image digest.

    ``name`` and ``namespace_path`` are the Function's parts, as a
    [SecretReference][lancedb.functions.SecretReference] carries a Secret's;
    they are what [get_function][lancedb.db.DBConnection.get_function] takes.
    An empty ``namespace_path`` is the root namespace.
    """

    name: str
    namespace_path: tuple[str, ...] = ()
    object_id: str
    location: str
    version: _ObjectVersion
    image: FunctionImage
    signature: FunctionSignature
    secret_bindings: tuple[SecretBinding, ...] = ()
    created_at_millis: int
    metadata: Mapping[str, str]
    disabled: bool

    def __call__(self, **arguments: Any) -> FunctionApplication:
        """Bind this exact version to named table columns.

        Every input must be a direct [lancedb.col][lancedb.expr.col]
        reference. A Function that declares initialization fields takes their
        values as further keyword arguments; they are constants of this
        binding, fixed for every row, and every instance the binding runs is
        created with them. A non-nullable field without a default must be
        given; an omitted nullable field is null, and the Function's own
        default applies to a null value.

        The returned application is immutable and retains a named-struct
        output as one binding, so every row's sibling values come from one
        logical Function evaluation. Map result fields to table columns with
        [FunctionApplication.rename][lancedb.functions.FunctionApplication.rename],
        then pass the application to
        [Table.add_columns][lancedb.table.Table.add_columns].

        Examples
        --------
        >>> from lancedb import col
        >>> application = function(  # doctest: +SKIP
        ...     title=col("title"),
        ...     body=col("body"),
        ... ).rename(columns={
        ...     "normalized_text": "search_text",
        ...     "token_count": "search_token_count",
        ... })
        >>> table.add_columns(application)  # doctest: +SKIP
        >>> table.add_columns({  # doctest: +SKIP
        ...     "embedding": embed(text=col("body"), model="small", dimensions=512)
        ... })
        """
        from lancedb.expr import Expr

        parameters = tuple(parameter.name for parameter in self.signature.inputs)
        fields = {field.name: field for field in self.signature.initialization}
        missing = [parameter for parameter in parameters if parameter not in arguments]
        missing_initialization = [
            name
            for name, field in fields.items()
            if not field.nullable and name not in arguments
        ]
        unknown = sorted(set(arguments) - set(parameters) - set(fields))
        if missing or missing_initialization or unknown:
            details = []
            if missing:
                details.append(f"missing inputs: {missing!r}")
            if missing_initialization:
                details.append(f"missing initialization: {missing_initialization!r}")
            if unknown:
                details.append(f"unknown arguments: {unknown!r}")
            raise TypeError("invalid Function arguments (" + "; ".join(details) + ")")

        bindings = []
        for parameter in parameters:
            value = arguments[parameter]
            if not isinstance(value, Expr) or value._column_path is None:
                raise TypeError(
                    f"Function input {parameter!r} must be a direct col(...) reference"
                )
            bindings.append(
                ApplicationInput(
                    parameter=parameter,
                    kind="column",
                    value={"path": value._column_path},
                )
            )
        initialization = {}
        for name in fields:
            if name in arguments:
                initialization[name] = _initialization_value(name, arguments[name])
        return FunctionApplication(
            function=FunctionVersionRef(
                name=self.name,
                namespace_path=self.namespace_path,
                object_id=self.object_id,
                location=self.location,
                version=self.version,
                manifest_digest=self.image.manifest_digest,
            ),
            inputs=tuple(bindings),
            output=self.signature.output,
            initialization=initialization,
        )


def _initialization_value(name: str, value: Any) -> Any:
    """One initialization value as plain JSON; the service checks its type."""
    from lancedb.expr import Expr

    def plain(value):
        if isinstance(value, Expr):
            raise TypeError(
                f"initialization field {name!r} takes a constant, not a column "
                "expression; every row of a binding shares one initialization"
            )
        if value is None or type(value) in (bool, int, str):
            return value
        if type(value) is float:
            if not math.isfinite(value):
                raise ValueError(
                    f"initialization field {name!r} must be finite, got {value!r}"
                )
            return value
        if isinstance(value, (list, tuple)):
            return [plain(child) for child in value]
        if isinstance(value, Mapping) and all(isinstance(key, str) for key in value):
            return {key: plain(child) for key, child in value.items()}
        raise TypeError(
            f"initialization field {name!r} takes booleans, numbers, strings, "
            f"lists, and string-keyed mappings, got {type(value).__name__}"
        )

    return plain(value)


class FunctionRegistrationRequest(_RemoteValue):
    """Stable remote registration envelope produced by :func:`udf`.

    Credential values deliberately have no field here. The only secret-shaped
    thing a client sends is ``secret_bindings``: the name of a Secret the
    database already holds, which the remote service resolves at execution.
    """

    name: str
    artifact: FunctionArtifactRequest
    signature: FunctionSignature
    runtime: PythonRuntimeSpec
    secret_bindings: tuple[SecretBinding, ...] = ()


class FunctionVersionRef(_OpenRemoteValue):
    name: str
    namespace_path: tuple[str, ...] = ()
    object_id: str
    location: str
    version: _ObjectVersion
    manifest_digest: str


class ApplicationInput(_OpenRemoteValue):
    """One parameter value.

    Slice 1 freezes integers, strings, booleans, nulls, arrays, and objects.
    Floating-point literal encoding is deferred until Python authoring is
    introduced with a language-neutral numeric representation.
    """

    parameter: str
    kind: str
    value: Any

    @field_validator("value")
    @classmethod
    def _validate_value(cls, value):
        return _validate_literal(value)


class FunctionApplication(_OpenRemoteValue):
    """Immutable pre-declaration application of an exact Function version.

    A named-struct output remains one application through table
    declaration and execution.
    [FunctionApplication.rename][lancedb.functions.FunctionApplication.rename]
    records the result-field to table-column mapping without splitting sibling
    outputs into separate UDF calls.
    """

    function: FunctionVersionRef
    inputs: tuple[ApplicationInput, ...]
    output: FunctionOutput
    columns: Mapping[str, str] = Field(default_factory=dict)
    # Transport JSON: the service validates it against the Function's
    # initialization schema and persists the Arrow row in the binding, so it is
    # never hashed and floating-point values are allowed.
    initialization: Mapping[str, Any] = Field(default_factory=dict)

    def _known_dict(self) -> dict[str, Any]:
        value = super()._known_dict()
        for name in self._unknown_field_names():
            value.pop(name, None)
        return value

    def _ensure_declarable(self) -> None:
        unknown = {f"application.{name}" for name in self._unknown_field_names()}
        unknown.update(
            f"function.{name}" for name in self.function._unknown_field_names()
        )
        for index, input_value in enumerate(self.inputs):
            unknown.update(
                f"inputs[{index}].{name}" for name in input_value._unknown_field_names()
            )
        unknown.update(f"output.{name}" for name in self.output._unknown_field_names())
        for index, field in enumerate(self.output.fields):
            unknown.update(
                f"output.fields[{index}].{name}"
                for name in field._unknown_field_names()
            )
        if unknown:
            raise ValueError(
                "Function application contains fields from a newer contract: "
                f"{sorted(unknown)!r}"
            )

    def rename(self, *, columns: Mapping[str, str]) -> FunctionApplication:
        """Return a copy with result-field to table-column aliases."""
        if self.output.kind != "named_struct":
            raise ValueError("rename(columns=...) requires a named-struct application")
        result_fields = {field.name for field in self.output.fields}
        unknown = set(columns) - result_fields
        if unknown:
            raise ValueError(f"unknown Function result fields: {sorted(unknown)!r}")
        merged = dict(self.columns)
        merged.update(columns)
        destinations = tuple(
            merged.get(field.name, field.name) for field in self.output.fields
        )
        if len(set(destinations)) != len(destinations):
            raise ValueError("FunctionApplication rename destinations must be unique")
        return self._copy(update={"columns": merged})


class InputBinding(_RemoteValue):
    parameter: str
    field_id: _Int32
    field_path: str
    arrow_type: str
    nullable: bool


class OutputMapping(_RemoteValue):
    """One stable result-field mapping."""

    result_field: str
    output_name: str
    output_field_id: _Int32
    output_ordinal: _UInt32
    arrow_type: str
    nullable: bool


class AssignmentMapping(_RemoteValue):
    """Internal physical column preserving flattened struct validity."""

    output_name: str
    output_field_id: _Int32


class FunctionBinding(_RemoteValue):
    """Immutable Function binding persisted by the Enterprise table service."""

    binding_id: str
    function: FunctionVersionRef
    inputs: tuple[InputBinding, ...]
    outputs: tuple[OutputMapping, ...]
    assignment: Optional[AssignmentMapping] = None
    input_schema: Optional[Mapping[str, Any]] = None
    output_schema: Optional[Mapping[str, Any]] = None
    initialization: Optional[str] = None
    """Standard base64 of the validated one-row Arrow IPC initialization stream."""

    def initialization_row(self) -> Optional[dict[str, Any]]:
        """The initialization values this binding creates instances with."""
        if self.initialization is None:
            return None
        reader = pa.ipc.open_stream(base64.b64decode(self.initialization))
        return reader.read_all().to_pylist()[0]


class RefreshColumnResult(_RemoteValue):
    """Terminal result of an expression-backed or Function-backed refresh Job.

    Local jobs produce this value in process. LanceDB Cloud and Enterprise
    decode the same value from the durable server-job terminal payload.
    """

    rows_assigned: _UInt64
    rows_failed: _UInt64
    rows_remaining: _UInt64
    source_version: _UInt64
    published_version: Optional[_UInt64] = None

    @property
    def rows_filled(self) -> int:
        """Deprecated compatibility alias for :attr:`rows_assigned`."""
        return self.rows_assigned

    @property
    def version(self) -> Optional[int]:
        """Deprecated compatibility alias for :attr:`published_version`."""
        return self.published_version


_FUNCTION_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_.-]*$")

_FUNCTION_BLOB_V2_TYPE = "blob_v2"
_ARROW_EXTENSION_NAME_KEY = "ARROW:extension:name"
_BLOB_V2_EXTENSION_NAME = "lance.blob.v2"
_NESTED_BLOB_COLLECTION_ERROR = (
    "unsupported Arrow type for Function signature: Blob v2 fields nested under "
    "collection types are not supported"
)


_GRAMMAR_PRIMITIVES = (
    (pa.bool_(), "bool"),
    (pa.int8(), "int8"),
    (pa.int16(), "int16"),
    (pa.int32(), "int32"),
    (pa.int64(), "int64"),
    (pa.uint8(), "uint8"),
    (pa.uint16(), "uint16"),
    (pa.uint32(), "uint32"),
    (pa.uint64(), "uint64"),
    (pa.float16(), "float16"),
    (pa.float32(), "float32"),
    (pa.float64(), "float64"),
    (pa.string(), "utf8"),
    (pa.large_string(), "large_utf8"),
    (pa.binary(), "binary"),
    (pa.large_binary(), "large_binary"),
    (pa.date32(), "date32"),
    (pa.date64(), "date64"),
)


def _canonical_arrow_type(data_type: pa.DataType) -> str:
    """The compact Function grammar, or canonical exact JSON for nested types."""
    grammar = _grammar_arrow_type(data_type)
    if grammar is not None:
        return grammar
    exact = _exact_arrow_type(data_type)
    return json.dumps(exact, ensure_ascii=False, sort_keys=True, separators=(",", ":"))


def _grammar_arrow_type(data_type: pa.DataType) -> Optional[str]:
    for candidate, name in _GRAMMAR_PRIMITIVES:
        if data_type == candidate:
            return name
    if pa.types.is_list(data_type) or pa.types.is_large_list(data_type):
        item = _grammar_list_item(data_type)
        if item is None:
            return None
        prefix = "list" if pa.types.is_list(data_type) else "large_list"
        return f"{prefix}<{item}>"
    if pa.types.is_fixed_size_list(data_type) and data_type.list_size > 0:
        item = _grammar_list_item(data_type)
        if item is not None:
            return f"fixed_size_list<{item}, {data_type.list_size}>"
    return None


def _grammar_list_item(data_type: pa.DataType) -> Optional[str]:
    """The grammar names only the item type; it always means a non-nullable
    child called `item`, so other child properties require exact JSON."""
    child = data_type.value_field
    if child.name != "item" or child.nullable or child.metadata:
        return None
    return _grammar_arrow_type(child.type)


def _validate_exact_arrow_field(field: pa.Field) -> None:
    if not field.name:
        raise TypeError(
            "unsupported Arrow type for Function signature: field names "
            "must not be empty"
        )
    if _is_blob_v2_field(field):
        if not _has_supported_blob_v2_layout(field):
            raise TypeError(
                "unsupported Arrow type for Function signature: lance.blob.v2 "
                f"requires a supported Blob storage layout, got {field}"
            )
        metadata = {
            (key.decode() if isinstance(key, bytes) else key): (
                value.decode() if isinstance(value, bytes) else value
            )
            for key, value in (field.metadata or {}).items()
        }
        if metadata and metadata != {
            _ARROW_EXTENSION_NAME_KEY: _BLOB_V2_EXTENSION_NAME
        }:
            raise TypeError(
                "unsupported Arrow type for Function signature: lance.blob.v2 "
                "field metadata must contain only its canonical extension marker"
            )
    elif field.metadata:
        raise TypeError(
            "unsupported Arrow type for Function signature: field metadata "
            f"is not supported, got {field}"
        )


def _has_supported_blob_v2_layout(field: pa.Field) -> bool:
    data_type = field.type
    if isinstance(data_type, pa.ExtensionType):
        data_type = data_type.storage_type
    if not pa.types.is_struct(data_type):
        return False

    fields = tuple(data_type)

    def matches(spec, compare_nullable) -> bool:
        return len(fields) == len(spec) and all(
            actual.name == name
            and actual.type == expected_type
            and (not check_nullable or actual.nullable == nullable)
            for actual, (name, expected_type, nullable), check_nullable in zip(
                fields, spec, compare_nullable
            )
        )

    logical_minimal = (
        ("data", pa.large_binary(), True),
        ("uri", pa.utf8(), True),
    )
    logical_full = logical_minimal + (
        ("position", pa.uint64(), True),
        ("size", pa.uint64(), True),
    )
    prepared = (
        ("kind", pa.uint8(), True),
        ("data", pa.large_binary(), True),
        ("uri", pa.utf8(), True),
        ("blob_id", pa.uint32(), True),
        ("blob_size", pa.uint64(), True),
        ("position", pa.uint64(), True),
    )
    descriptor = (
        ("kind", pa.uint8(), False),
        ("position", pa.uint64(), False),
        ("size", pa.uint64(), False),
        ("blob_id", pa.uint32(), False),
        ("blob_uri", pa.utf8(), False),
    )
    return (
        matches(logical_minimal, (True, True))
        or matches(logical_full, (True, True, False, False))
        or matches(prepared, (True,) * len(prepared))
        or matches(descriptor, (False,) * len(descriptor))
    )


def _canonical_arrow_field(field: pa.Field) -> str:
    _validate_exact_arrow_field(field)
    if _is_blob_v2_field(field):
        return _FUNCTION_BLOB_V2_TYPE
    return _canonical_arrow_type(field.type)


def _blob_storage_type(field: pa.Field) -> pa.DataType:
    data_type = field.type
    if isinstance(data_type, pa.ExtensionType):
        return data_type.storage_type
    return data_type


def _exact_blob_storage_type(field: pa.Field) -> dict[str, Any]:
    storage = _blob_storage_type(field)
    if not pa.types.is_struct(storage):
        raise TypeError(
            "unsupported Arrow type for Function signature: lance.blob.v2 "
            "requires struct storage"
        )
    return {
        "type": "struct",
        "fields": [
            {
                "name": child.name,
                "nullable": child.nullable,
                "type": _exact_arrow_type(child.type),
            }
            for child in storage
        ],
    }


def _data_type_has_blob_v2(data_type: pa.DataType) -> bool:
    if pa.types.is_struct(data_type):
        return any(
            _is_blob_v2_field(field) or _data_type_has_blob_v2(field.type)
            for field in data_type
        )
    if (
        pa.types.is_list(data_type)
        or pa.types.is_large_list(data_type)
        or pa.types.is_fixed_size_list(data_type)
    ):
        field = data_type.value_field
        return _is_blob_v2_field(field) or _data_type_has_blob_v2(field.type)
    if pa.types.is_map(data_type):
        return any(
            _is_blob_v2_field(field) or _data_type_has_blob_v2(field.type)
            for field in (data_type.key_field, data_type.item_field)
        )
    return False


def _exact_arrow_field(
    field: pa.Field, *, inside_collection: bool = False
) -> dict[str, Any]:
    _validate_exact_arrow_field(field)
    if _is_blob_v2_field(field):
        if inside_collection:
            raise TypeError(_NESTED_BLOB_COLLECTION_ERROR)
        return {
            "name": field.name,
            "nullable": field.nullable,
            "type": _exact_blob_storage_type(field),
            "metadata": {
                _ARROW_EXTENSION_NAME_KEY: _BLOB_V2_EXTENSION_NAME,
            },
        }
    value = {
        "name": field.name,
        "nullable": field.nullable,
        "type": _exact_arrow_type(field.type, inside_collection=inside_collection),
    }
    return value


def _exact_arrow_type(
    data_type: pa.DataType, *, inside_collection: bool = False
) -> dict[str, Any]:
    for candidate, name in _GRAMMAR_PRIMITIVES:
        if data_type == candidate:
            return {"type": name}
    if pa.types.is_struct(data_type):
        fields = list(data_type)
        names = [field.name for field in fields]
        if not fields or len(set(names)) != len(names):
            raise TypeError(
                "unsupported Arrow type for Function signature: structs must have "
                "non-empty, uniquely named fields"
            )
        return {
            "type": "struct",
            "fields": [
                _exact_arrow_field(field, inside_collection=inside_collection)
                for field in fields
            ],
        }
    if (
        pa.types.is_list(data_type)
        or pa.types.is_large_list(data_type)
        or pa.types.is_fixed_size_list(data_type)
    ):
        if pa.types.is_fixed_size_list(data_type):
            if data_type.value_field.name != "item":
                raise TypeError(
                    "unsupported Arrow type for Function signature: fixed-size list "
                    "items must be named 'item'"
                )
            if data_type.list_size <= 0:
                raise TypeError(
                    f"unsupported Arrow type for Function signature: {data_type}"
                )
        value: dict[str, Any] = {
            "type": (
                "list"
                if pa.types.is_list(data_type)
                else "large_list"
                if pa.types.is_large_list(data_type)
                else "fixed_size_list"
            ),
            "fields": [
                _exact_arrow_field(data_type.value_field, inside_collection=True)
            ],
        }
        if pa.types.is_fixed_size_list(data_type):
            value["length"] = data_type.list_size
        return value
    if pa.types.is_map(data_type) and _data_type_has_blob_v2(data_type):
        raise TypeError(_NESTED_BLOB_COLLECTION_ERROR)
    raise TypeError(f"unsupported Arrow type for Function signature: {data_type}")


def _list_of(item: pa.DataType) -> pa.DataType:
    return pa.list_(pa.field("item", item, nullable=False))


def _annotation_type(annotation: Any) -> tuple[pa.DataType, bool]:
    nullable = False
    origin = get_origin(annotation)
    if origin in (Union, types.UnionType):
        arguments = get_args(annotation)
        non_none = tuple(
            argument for argument in arguments if argument is not type(None)
        )
        if len(non_none) != 1 or len(non_none) == len(arguments):
            raise TypeError(f"unsupported union annotation: {annotation!r}")
        annotation = non_none[0]
        nullable = True

    origin = get_origin(annotation)
    if origin is Annotated:
        base, *metadata = get_args(annotation)
        arrow_types = [value for value in metadata if isinstance(value, pa.DataType)]
        if len(arrow_types) != 1:
            raise TypeError(
                "Annotated Function types require exactly one PyArrow DataType"
            )
        # A struct value is written as a dict, which has no Arrow type of its
        # own; the Annotated metadata supplies it.
        if base is dict or get_origin(base) is dict:
            base_nullable = False
        else:
            _, base_nullable = _annotation_type(base)
        return arrow_types[0], nullable or base_nullable

    if isinstance(annotation, pa.DataType):
        return annotation, nullable
    if annotation is bool:
        return pa.bool_(), nullable
    if annotation is int:
        return pa.int64(), nullable
    if annotation is float:
        return pa.float64(), nullable
    if annotation is str:
        return pa.string(), nullable
    if annotation is bytes:
        return pa.binary(), nullable
    if annotation is date:
        return pa.date32(), nullable
    if annotation is datetime:
        return pa.timestamp("us"), nullable
    if get_origin(annotation) is list:
        arguments = get_args(annotation)
        if len(arguments) != 1:
            raise TypeError(f"unsupported list annotation: {annotation!r}")
        value_type, value_nullable = _annotation_type(arguments[0])
        if value_nullable:
            raise TypeError("nullable Function list elements are not supported")
        return _list_of(value_type), nullable
    raise TypeError(f"unsupported Function annotation: {annotation!r}")


def _callable_parameters(function: Callable[..., Any]) -> tuple[inspect.Parameter, ...]:
    parameters = tuple(inspect.signature(function).parameters.values())
    for parameter in parameters:
        if parameter.kind in (
            inspect.Parameter.POSITIONAL_ONLY,
            inspect.Parameter.VAR_POSITIONAL,
            inspect.Parameter.VAR_KEYWORD,
        ):
            raise TypeError("Function callables require named, non-variadic parameters")
        if parameter.default is not inspect.Parameter.empty:
            raise TypeError("Function callable defaults are not supported")
    return parameters


def _function_output(output: pa.DataType | pa.Field | pa.Schema) -> FunctionOutput:
    if isinstance(output, pa.Schema):
        if output.metadata:
            raise TypeError("Function output schema metadata is not supported")
        fields = tuple(output)
    elif (
        isinstance(output, pa.Field)
        and not _is_blob_v2_field(output)
        and pa.types.is_struct(output.type)
    ):
        _validate_exact_arrow_field(output)
        if output.nullable:
            raise ValueError("Function output must be non-nullable")
        fields = tuple(output.type)
    elif isinstance(output, pa.DataType) and pa.types.is_struct(output):
        fields = tuple(output)
    else:
        field = (
            output
            if isinstance(output, pa.Field)
            else pa.field("result", output, nullable=False)
        )
        if not isinstance(field, pa.Field):
            raise TypeError(
                "output_schema must be a PyArrow DataType, Field, or Schema"
            )
        _validate_exact_arrow_field(field)
        if field.nullable:
            raise ValueError("Function output must be non-nullable")
        return FunctionOutput(
            kind="scalar",
            arrow_type=_canonical_arrow_field(field),
            nullable=False,
        )

    if not fields:
        raise ValueError("named-struct Function output must contain at least one field")
    for field in fields:
        _validate_exact_arrow_field(field)
    names = [field.name for field in fields]
    if len(set(names)) != len(names):
        raise ValueError("Function output field names must be unique")
    return FunctionOutput(
        kind="named_struct",
        fields=tuple(
            FunctionResultField(
                name=field.name,
                arrow_type=_canonical_arrow_field(field),
                nullable=field.nullable,
            )
            for field in fields
        ),
    )


def _infer_signature(
    function: Callable[..., Any],
    input_schema: Optional[pa.Schema],
    output_schema: Optional[pa.DataType | pa.Field | pa.Schema],
    *,
    method: bool = False,
) -> FunctionSignature:
    parameters = _callable_parameters(function)
    if method:
        parameters = parameters[1:]
    if (input_schema is None) != (output_schema is None):
        raise ValueError("input_schema and output_schema must be provided together")

    try:
        annotations = get_type_hints(function, include_extras=True)
    except Exception as error:
        raise TypeError(f"failed to resolve Function annotations: {error}") from error

    def input_parameter(field):
        annotation = annotations.get(field.name)
        if get_origin(annotation) in (Union, types.UnionType):
            choices = [
                value for value in get_args(annotation) if value is not type(None)
            ]
            if len(choices) == 1:
                annotation = choices[0]
        arrow_type = _canonical_arrow_field(field)
        if annotation is BlobReader:
            if arrow_type != _FUNCTION_BLOB_V2_TYPE:
                raise TypeError(
                    f"BlobReader input '{field.name}' requires a Blob V2 field"
                )
            access = "reader"
        elif annotation is bytes:
            if arrow_type != _FUNCTION_BLOB_V2_TYPE and not (
                pa.types.is_binary(field.type) or pa.types.is_large_binary(field.type)
            ):
                raise TypeError(
                    f"bytes input '{field.name}' requires binary or Blob V2"
                )
            access = "bytes"
        else:
            access = "bytes" if arrow_type == _FUNCTION_BLOB_V2_TYPE else None
        return FunctionParameter(
            name=field.name,
            arrow_type=arrow_type,
            nullable=field.nullable,
            blob_access=access,
        )

    if input_schema is not None:
        if not isinstance(input_schema, pa.Schema):
            raise TypeError("input_schema must be a PyArrow Schema")
        if input_schema.metadata:
            raise TypeError("Function input schema metadata is not supported")
        for field in input_schema:
            _validate_exact_arrow_field(field)
        expected = tuple(parameter.name for parameter in parameters)
        actual = tuple(input_schema.names)
        if actual != expected:
            raise ValueError(
                "input_schema fields must exactly match callable parameters in order: "
                f"expected {expected!r}, got {actual!r}"
            )
        inputs = tuple(input_parameter(field) for field in input_schema)
        return FunctionSignature(inputs=inputs, output=_function_output(output_schema))

    missing = [
        parameter.name for parameter in parameters if parameter.name not in annotations
    ]
    if missing or "return" not in annotations:
        names = missing + ([] if "return" in annotations else ["return"])
        raise TypeError(f"missing Function annotations: {names!r}")
    inputs = []
    for parameter in parameters:
        annotation = annotations[parameter.name]
        base = annotation
        nullable = False
        if get_origin(base) in (Union, types.UnionType):
            choices = [value for value in get_args(base) if value is not type(None)]
            if len(choices) == 1:
                base, nullable = choices[0], True
        if base is BlobReader:
            inputs.append(
                input_parameter(_blob_field(parameter.name, nullable=nullable))
            )
            continue
        data_type, nullable = _annotation_type(annotation)
        inputs.append(
            input_parameter(pa.field(parameter.name, data_type, nullable=nullable))
        )
    output_type, output_nullable = _annotation_type(annotations["return"])
    if output_nullable:
        raise ValueError("Function output must be non-nullable")
    return FunctionSignature(
        inputs=tuple(inputs),
        output=_function_output(pa.field("result", output_type, nullable=False)),
    )


def _is_udf_decorator(node: ast.expr) -> bool:
    if isinstance(node, ast.Call):
        node = node.func
    return (isinstance(node, ast.Name) and node.id == "udf") or (
        isinstance(node, ast.Attribute) and node.attr == "udf"
    )


def _literal_source(value: Any) -> str:
    if value is None or type(value) in (bool, int, str, bytes):
        return repr(value)
    if type(value) is float and math.isfinite(value):
        return repr(value)
    if type(value) is tuple:
        children = ", ".join(_literal_source(child) for child in value)
        if len(value) == 1:
            children += ","
        return f"({children})"
    raise TypeError(
        "Function source references an unsupported global value of type "
        f"{type(value).__name__}"
    )


_DYNAMIC_NAMESPACE_ACCESS = frozenset(
    {"globals", "locals", "vars", "eval", "exec", "compile", "__import__"}
)
# Modules that hand out namespaces (`sys.modules`, `builtins`, importers,
# introspection). The artifact's module namespace holds only the names it was
# packaged with, so reaching around it cannot be represented.
_NAMESPACE_MODULES = frozenset(
    {"sys", "builtins", "importlib", "inspect", "gc", "ctypes", "types"}
)


def _namespace_acquisition(
    definition: ast.FunctionDef, references: set[str]
) -> list[str]:
    found = set(references & _DYNAMIC_NAMESPACE_ACCESS)
    for node in ast.walk(definition):
        if isinstance(node, ast.Import):
            found.update(
                alias.name
                for alias in node.names
                if alias.name.split(".")[0] in _NAMESPACE_MODULES
            )
        elif isinstance(node, ast.ImportFrom) and node.module:
            if node.module.split(".")[0] in _NAMESPACE_MODULES:
                found.add(node.module)
    return sorted(found)


def _module_references(module_source: str) -> set[str]:
    """Names `module_source` loads from module scope, by Python's own scope
    analysis on the exact text that ships: loads inside any definition, plus
    what the definition statement itself evaluates at module scope (bases,
    decorators, defaults). Free variables belong to an enclosing scope inside
    the definition, and postponed annotations are not runtime loads."""

    def visit(table: symtable.SymbolTable, found: set[str]) -> None:
        for symbol in table.get_symbols():
            if symbol.is_global() and (
                symbol.is_referenced() or symbol.is_declared_global()
            ):
                found.add(symbol.get_name())
        for child in table.get_children():
            visit(child, found)

    found: set[str] = set()
    module = symtable.symtable(module_source, "<udf>", "exec")
    for symbol in module.get_symbols():
        if (
            symbol.is_referenced()
            and not symbol.is_assigned()
            and not symbol.is_imported()
        ):
            found.add(symbol.get_name())
    for table in module.get_children():
        visit(table, found)
    return found


_SOURCE_HEADER = "from __future__ import annotations"
_SOURCE_ENTRY_FILE = "user_function.py"
_PYTHON_CALLABLE_ARTIFACT = "python_callable"
_PYTHON_BUNDLE_ARTIFACT = "python_bundle"
# Modules the Function image places beside user code.
_RESERVED_SOURCE_MODULES = frozenset(
    {"user_function", "lance_source", "lance_callable", "lance_blob"}
)
# The service accepts bundle paths made of these components only.
_SOURCE_PATH_COMPONENT = re.compile(r"[A-Za-z_][A-Za-z0-9_]*")
_REQUIREMENT_NAME = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]*")


def _class_source(cls: type) -> str:
    """Source of a module-level class.

    `inspect.getsource` finds a class through its module's file, which a
    notebook cell or a doctest does not have. A method's code object records
    where the class body really is, so that is consulted first.
    """
    for member in vars(cls).values():
        if isinstance(member, (staticmethod, classmethod)):
            member = member.__func__
        code = (
            getattr(inspect.unwrap(member), "__code__", None)
            if callable(member)
            else None
        )
        if code is None:
            continue
        lines = linecache.getlines(code.co_filename)
        try:
            module = ast.parse("".join(lines))
        except SyntaxError:
            continue
        for node in module.body:
            if (
                isinstance(node, ast.ClassDef)
                and node.name == cls.__name__
                and node.lineno <= code.co_firstlineno <= node.end_lineno
            ):
                start = min(
                    [node.lineno]
                    + [decorator.lineno for decorator in node.decorator_list]
                )
                return "".join(lines[start - 1 : node.end_lineno])
    return inspect.getsource(cls)


def _definition_node(
    target: Any, *, entry: bool
) -> Union[ast.FunctionDef, ast.ClassDef]:
    try:
        source = textwrap.dedent(
            _class_source(target)
            if inspect.isclass(target)
            else inspect.getsource(target)
        )
    except (OSError, TypeError) as error:
        raise ValueError(
            f"@udf requires inspectable Python source for {target.__qualname__!r}"
        ) from error
    module = ast.parse(source)
    is_class = inspect.isclass(target)
    definitions = [
        node
        for node in module.body
        if isinstance(
            node,
            (ast.ClassDef,) if is_class else (ast.FunctionDef, ast.AsyncFunctionDef),
        )
        and node.name == target.__name__
    ]
    if len(definitions) != 1 or isinstance(definitions[0], ast.AsyncFunctionDef):
        raise ValueError(
            "@udf source must contain exactly one "
            + ("class" if is_class else "synchronous function")
        )
    definition = definitions[0]
    if entry:
        if any(
            not _is_udf_decorator(decorator) for decorator in definition.decorator_list
        ):
            raise ValueError("@udf cannot package additional Python decorators")
        definition.decorator_list = []
    return definition


def _definition_namespace(target: Any) -> Mapping[str, Any]:
    """The module namespace a function or class resolves its globals in."""
    if inspect.isclass(target):
        module = sys.modules.get(target.__module__)
        if module is None:
            raise ValueError(
                f"@udf cannot resolve the module of {target.__qualname__!r}"
            )
        namespace = vars(module)
        environment = namespace.get("__builtins__", builtins)
        environment = (
            vars(environment)
            if isinstance(environment, types.ModuleType)
            else environment
        )
    else:
        namespace = target.__globals__
        environment = target.__builtins__
    # The artifact runs under the standard builtins; only the exact mapping is
    # provably equivalent (a subclass or copy can change lookups and hooks).
    if environment is not vars(builtins):
        raise ValueError("@udf cannot package a non-standard builtins environment")
    return namespace


def _is_self_reference(target: Any, namespace: Mapping[str, Any], name: str) -> bool:
    """`name` inside the definition means the definition itself unless the
    module has since bound it to something else."""
    if name != target.__name__:
        return False
    bound = namespace.get(name, target)
    if bound is target:
        return True
    # The decorator's own result is the one wrapper known to behave like
    # `target`; any other binding may differ from the packaged definition.
    return type(bound) is UdfDefinition and bound._function is target


def _packaged_by_source(value: Any) -> bool:
    """Functions and classes defined in `__main__` -- notebook cells and
    scripts -- have no module a worker can import, so their source travels."""
    if not (inspect.isfunction(value) or inspect.isclass(value)):
        return False
    if getattr(value, "__module__", None) != "__main__":
        return False
    qualname = value.__qualname__
    if "<" in qualname or "." in qualname:
        raise ValueError(
            f"@udf cannot package {qualname!r}: only module-level functions and "
            "classes travel with a Function, not lambdas, closures, or nested "
            "definitions"
        )
    return True


def _requirement_name(requirement: str) -> str:
    requirement = requirement.rpartition("::")[2]
    match = _REQUIREMENT_NAME.match(requirement.strip())
    return _normalized_distribution(match.group(0) if match else requirement)


def _normalized_distribution(name: str) -> str:
    return re.sub(r"[-_.]+", "_", name).lower()


def _installed_roots() -> tuple[str, ...]:
    import os
    import site
    import sysconfig

    roots = set(site.getsitepackages())
    roots.add(site.getusersitepackages())
    paths = sysconfig.get_paths()
    for key in ("purelib", "platlib", "stdlib", "platstdlib"):
        if key in paths:
            roots.add(paths[key])
    return tuple(os.path.realpath(root) + os.sep for root in roots)


def _module_locations(module: types.ModuleType) -> tuple[str, ...]:
    import os

    file = getattr(module, "__file__", None)
    if file:
        return (os.path.realpath(file),)
    return tuple(os.path.realpath(path) for path in getattr(module, "__path__", ()))


class _SourcePackager:
    """Builds the entry module of a Function artifact.

    The module holds the definition plus exactly the module-level names it
    references: modules as imports, importable classes and functions as
    imports, literals inline, and functions and classes defined in
    ``__main__`` by source, recursively, dependencies first.
    """

    def __init__(self, code_modules: frozenset[str], requirements: frozenset[str]):
        self._code_modules = code_modules
        self._requirements = requirements
        self._imports: dict[str, str] = {}
        self._definitions: list[str] = []
        self._packaged: set[int] = set()
        self._names: dict[str, Any] = {}

    def package(self, target: Any) -> str:
        entry = self._definition(target, entry=True)
        parts = [_SOURCE_HEADER]
        if self._imports:
            parts.extend(["", *(self._imports[name] for name in sorted(self._imports))])
        for definition in self._definitions:
            parts.extend(["", definition])
        parts.extend(["", entry, ""])
        return "\n".join(parts)

    def _definition(self, target: Any, *, entry: bool) -> str:
        definition = _definition_node(target, entry=entry)
        if not inspect.isclass(target) and inspect.getclosurevars(target).nonlocals:
            raise ValueError(
                "@udf cannot package functions that capture closure values"
            )
        source = ast.unparse(definition)
        references = _module_references(f"{_SOURCE_HEADER}\n\n{source}\n")
        dynamic = _namespace_acquisition(definition, references)
        if dynamic:
            raise ValueError(
                f"@udf cannot package dynamic namespace access: {dynamic!r}"
            )
        namespace = _definition_namespace(target)
        self._packaged.add(id(target))
        self._claim(target.__name__, ("definition", id(target)))
        unresolved = []
        for name in sorted(references):
            if name == target.__name__:
                if not _is_self_reference(target, namespace, name):
                    raise ValueError(
                        f"@udf cannot package {name!r}: the module binds that name "
                        "to another value, which the artifact's own definition "
                        "would shadow"
                    )
                continue
            if name in namespace:
                self._bind(name, namespace[name])
            elif not hasattr(builtins, name):
                unresolved.append(name)
        if unresolved:
            raise ValueError(
                f"@udf source contains unresolved global names: {unresolved!r}"
            )
        return source

    def _claim(self, name: str, meaning: Any) -> None:
        claimed = self._names.setdefault(name, meaning)
        if claimed != meaning:
            raise ValueError(
                f"@udf cannot package {name!r}: packaged definitions bind that "
                "name to different values"
            )

    def _bind(self, name: str, value: Any) -> None:
        if _packaged_by_source(value):
            if value.__name__ != name:
                raise ValueError(
                    f"@udf packages {value.__qualname__!r} by source, so it must be "
                    f"referenced by its own name, not {name!r}"
                )
            if id(value) not in self._packaged:
                self._definitions.append(self._definition(value, entry=False))
            self._claim(name, ("definition", id(value)))
            return
        line = self._global_source(name, value)
        self._claim(name, line)
        self._imports[name] = line

    def _global_source(self, name: str, value: Any) -> str:
        """One module-level line that rebinds `name` to `value` in the artifact:
        an import for modules and importable classes/functions, a literal
        otherwise."""
        if value is BlobReader:
            return f"from lance_blob import BlobReader as {name}"
        if isinstance(value, types.ModuleType):
            if value.__name__.split(".")[0] in _NAMESPACE_MODULES:
                raise ValueError(
                    f"@udf cannot package dynamic namespace access: {value.__name__!r}"
                )
            try:
                imported = importlib.import_module(value.__name__)
            except ImportError:
                imported = None
            if imported is not value:
                raise TypeError(
                    f"Function source references module {name!r} that does not "
                    f"import as {value.__name__!r}"
                )
            self._require_importable(value.__name__)
            return f"import {value.__name__} as {name}"
        module_name = getattr(value, "__module__", None)
        qualname = getattr(value, "__qualname__", None)
        if (
            isinstance(module_name, str)
            and isinstance(qualname, str)
            and module_name != "__main__"
            and "." not in qualname
            and "<" not in qualname
        ):
            try:
                imported = getattr(importlib.import_module(module_name), qualname)
            except (ImportError, AttributeError):
                imported = None
            if imported is value:
                self._require_importable(module_name)
                return f"from {module_name} import {qualname} as {name}"
        return f"{name} = {_literal_source(value)}"

    def _require_importable(self, module_name: str) -> None:
        """Warn about an import the Function's environment likely cannot
        satisfy: a module that lives in a local source tree, is not shipped
        with ``code=``, and names no declared package.

        A warning rather than an error: a distribution name need not match
        the module it installs, so a declared package can still provide the
        module, and callers may rewrite the packaged source before
        registering it."""
        top = module_name.partition(".")[0]
        if (
            top in self._code_modules
            or top in sys.stdlib_module_names
            or top in sys.builtin_module_names
            or _normalized_distribution(top) in self._requirements
        ):
            return
        module = sys.modules.get(top)
        locations = () if module is None else _module_locations(module)
        roots = _installed_roots()
        if not locations or any(location.startswith(roots) for location in locations):
            return
        warnings.warn(
            f"@udf source imports {module_name!r} from {locations[0]}, a local "
            "module the Function's environment may not be able to import; ship "
            f"it with code=[{top}] or declare the package that provides it in "
            "pip/conda",
            UserWarning,
        )


def _code_files(modules: Sequence[Any]) -> tuple[dict[str, str], frozenset[str]]:
    """The Python source files of the modules and packages in ``code=``,
    keyed by their path on the Function's import path."""
    import os
    from pathlib import Path

    files: dict[str, str] = {}
    names: set[str] = set()
    for module in modules:
        if not isinstance(module, types.ModuleType):
            raise TypeError(
                f"code= takes imported modules or packages, got {type(module).__name__}"
            )
        name = module.__name__
        if "." in name:
            raise ValueError(
                f"code= takes top-level modules or packages; pass "
                f"{name.partition('.')[0]!r} instead of {name!r}"
            )
        if (
            name in _RESERVED_SOURCE_MODULES
            or name == "__main__"
            or name in sys.stdlib_module_names
            or not _SOURCE_PATH_COMPONENT.fullmatch(name)
        ):
            raise ValueError(f"code= cannot ship module {name!r}")
        if name in names:
            raise ValueError(f"code= lists module {name!r} more than once")
        names.add(name)
        paths = list(getattr(module, "__path__", ()))
        if paths:
            if len(paths) != 1 or not getattr(module, "__file__", None):
                raise ValueError(
                    f"code= cannot ship namespace package {name!r}; give it an "
                    "__init__.py"
                )
            root = Path(paths[0])
            sources = []
            for directory, subdirectories, filenames in os.walk(root):
                # Only importable modules travel: data files, caches, and
                # directories that are not identifiers cannot be imported.
                subdirectories[:] = sorted(
                    child
                    for child in subdirectories
                    if _SOURCE_PATH_COMPONENT.fullmatch(child)
                    and child != "__pycache__"
                )
                for filename in filenames:
                    stem, suffix = os.path.splitext(filename)
                    if suffix == ".py" and _SOURCE_PATH_COMPONENT.fullmatch(stem):
                        sources.append(Path(directory) / filename)
            for path in sources:
                relative = path.relative_to(root).as_posix()
                files[f"{name}/{relative}"] = _source_text(path)
        else:
            file = getattr(module, "__file__", None)
            if not file or not file.endswith(".py"):
                raise ValueError(f"code= ships Python source; module {name!r} has none")
            files[f"{name}.py"] = _source_text(Path(file))
    return files, frozenset(names)


def _source_text(path) -> str:
    try:
        return path.read_text(encoding="utf-8")
    except UnicodeDecodeError as error:
        raise ValueError(f"code= ships UTF-8 source; {path} is not UTF-8") from error


def _initialization_type(data_type: pa.DataType) -> bool:
    """Types with one unambiguous JSON and SQL literal spelling."""
    if data_type in (
        pa.bool_(),
        pa.int8(),
        pa.int16(),
        pa.int32(),
        pa.int64(),
        pa.uint8(),
        pa.uint16(),
        pa.uint32(),
        pa.uint64(),
        pa.float32(),
        pa.float64(),
        pa.string(),
        pa.large_string(),
    ):
        return True
    if (
        pa.types.is_list(data_type)
        or pa.types.is_large_list(data_type)
        or pa.types.is_fixed_size_list(data_type)
    ):
        return _initialization_type(data_type.value_type)
    if pa.types.is_struct(data_type):
        return all(_initialization_type(field.type) for field in data_type)
    return False


def _initialization_signature(cls: type) -> tuple[FunctionParameter, ...]:
    """The initialization fields a class Function's ``__init__`` declares."""
    if cls.__init__ is object.__init__:
        return ()
    parameters = tuple(inspect.signature(cls.__init__).parameters.values())[1:]
    for parameter in parameters:
        if parameter.kind not in (
            inspect.Parameter.POSITIONAL_OR_KEYWORD,
            inspect.Parameter.KEYWORD_ONLY,
        ):
            raise TypeError(
                "Function initialization parameters must be named and "
                f"non-variadic: {parameter.name!r}"
            )
    try:
        annotations = get_type_hints(cls.__init__, include_extras=True)
    except Exception as error:
        raise TypeError(
            f"failed to resolve Function initialization annotations: {error}"
        ) from error
    missing = [
        parameter.name for parameter in parameters if parameter.name not in annotations
    ]
    if missing:
        raise TypeError(f"missing Function initialization annotations: {missing!r}")
    fields = []
    for parameter in parameters:
        data_type, nullable = _annotation_type(annotations[parameter.name])
        if not _initialization_type(data_type):
            raise TypeError(
                f"Function initialization parameter {parameter.name!r} has type "
                f"{data_type}; initialization takes booleans, integers, floats, "
                "strings, and lists or structs of them"
            )
        # A parameter with a default may be omitted from a binding; a null
        # value then takes the default.
        if parameter.default is not inspect.Parameter.empty:
            nullable = True
        fields.append(
            FunctionParameter(
                name=parameter.name,
                arrow_type=_canonical_arrow_field(
                    pa.field(parameter.name, data_type, nullable=nullable)
                ),
                nullable=nullable,
            )
        )
    return tuple(fields)


def _class_member(cls: type, name: str) -> Any:
    for klass in cls.__mro__:
        if klass is not object and name in vars(klass):
            return vars(klass)[name]
    return None


def _class_call(cls: type) -> Callable[..., Any]:
    call = _class_member(cls, "__call__")
    if not inspect.isfunction(call) or inspect.iscoroutinefunction(call):
        raise TypeError("a class Function must define a synchronous __call__ method")
    close = _class_member(cls, "close")
    if close is not None and (
        not inspect.isfunction(close) or len(inspect.signature(close).parameters) != 1
    ):
        raise TypeError("a class Function's close must be a method taking no arguments")
    return call


class UdfDefinition:
    """A Python callable prepared for remote Function registration.

    Instances are created with :func:`udf`. Calling an instance executes the
    original Python function, or constructs the original class, which keeps
    local unit testing ordinary. Remote execution adapts the callable to the
    internal Arrow batch ABI described by the registration artifact.
    """

    def __init__(
        self,
        function: Any,
        *,
        name: Optional[str],
        input_schema: Optional[pa.Schema],
        output_schema: Optional[pa.DataType | pa.Field | pa.Schema],
        pip: tuple[str, ...],
        env: Mapping[str, str],
        python_version: Optional[str],
        gpu: bool = False,
        conda: tuple[str, ...] = (),
        conda_channels: tuple[str, ...] = (),
        code: tuple[types.ModuleType, ...] = (),
    ):
        is_class = inspect.isclass(function)
        if not is_class and (
            not inspect.isfunction(function) or inspect.iscoroutinefunction(function)
        ):
            raise TypeError("@udf requires a synchronous Python function or a class")
        function_name = name or function.__name__
        if not _FUNCTION_NAME.fullmatch(function_name):
            raise ValueError(f"invalid Function name: {function_name!r}")
        if pip and conda:
            raise ValueError("a Function environment is pip or conda, not both")
        if conda_channels and not conda:
            raise ValueError("conda_channels requires conda packages")
        packages = tuple(sorted(set(conda if conda else pip)))
        if any(not package or package != package.strip() for package in packages):
            raise ValueError("package requirements must be non-empty and trimmed")
        if conda:
            environment_spec = PythonEnvironmentSpec(
                kind="conda", packages=packages, channels=tuple(conda_channels)
            )
        else:
            environment_spec = PythonEnvironmentSpec(kind="pip", packages=packages)
        environment = dict(env)
        if any(
            not isinstance(key, str) or not isinstance(value, str)
            for key, value in environment.items()
        ):
            raise TypeError("Function env keys and values must be strings")
        if is_class:
            if "<" in function.__qualname__ or "." in function.__qualname__:
                raise ValueError(
                    "@udf classes must be defined at module level; nested classes "
                    "would capture their enclosing scope"
                )
            signature = _infer_signature(
                _class_call(function), input_schema, output_schema, method=True
            )
            initialization = _initialization_signature(function)
            shared = sorted(
                {field.name for field in initialization}
                & {parameter.name for parameter in signature.inputs}
            )
            if shared:
                raise ValueError(
                    "Function initialization parameters and inputs share names "
                    f"{shared!r}; a binding passes both as keyword arguments"
                )
            signature = signature._copy(update={"initialization": initialization})
        else:
            signature = _infer_signature(function, input_schema, output_schema)
        code_files, code_modules = _code_files(code)
        entry_source = _SourcePackager(
            code_modules, frozenset(_requirement_name(package) for package in packages)
        ).package(function)
        if code_files:
            kind = _PYTHON_BUNDLE_ARTIFACT
            source = json.dumps(
                {"files": {**code_files, _SOURCE_ENTRY_FILE: entry_source}},
                ensure_ascii=False,
                sort_keys=True,
                separators=(",", ":"),
            ).encode("utf-8")
        else:
            kind = _PYTHON_CALLABLE_ARTIFACT
            source = entry_source.encode("utf-8")
        digest = f"sha256:{hashlib.sha256(source).hexdigest()}"
        gpu_marker = _normalize_gpu_marker(gpu)
        runtime = PythonRuntimeSpec(
            kind="python_v2" if gpu_marker is not None else "python",
            python_version=python_version
            or f"{sys.version_info.major}.{sys.version_info.minor}",
            environment=environment_spec,
            env=environment,
            gpu=gpu_marker,
        )
        self._function = function
        self._request = FunctionRegistrationRequest(
            name=function_name,
            artifact=FunctionArtifactRequest(
                kind=kind,
                digest=digest,
                entrypoint=function.__name__,
                content=FunctionArtifactContent(
                    encoding="base64",
                    data=base64.b64encode(source).decode("ascii"),
                ),
                adapter=PythonAdapterSpec(
                    kind="scalar_to_arrow_batch",
                    version=1,
                ),
            ),
            signature=signature,
            runtime=runtime,
        )
        # A class's own __dict__ holds its methods; copying it would shadow
        # this wrapper's attributes.
        functools.update_wrapper(
            self, function, updated=() if is_class else functools.WRAPPER_UPDATES
        )

    @property
    def registration_request(self) -> FunctionRegistrationRequest:
        """The immutable request sent by ``create_function_async``.

        Carries no secret bindings. Binding is a registration-time decision,
        so a Function bound to Secrets is registered through :meth:`bind_secrets`,
        which is what ``create_function`` calls.
        """
        return self._request

    def bind_secrets(
        self, secrets: Optional[Sequence[EnvVarSecret]]
    ) -> FunctionRegistrationRequest:
        """The registration request for this definition bound to ``secrets``.

        Binding does not change the Function's source: each
        [EnvVarSecret][lancedb.secrets.EnvVarSecret] names a Secret and the
        environment variable its value should arrive in, and the Function reads
        that variable the way it already did. Whether the named Secrets exist is
        the server's answer, not this one.
        """
        bindings = () if secrets is None else tuple(secrets)
        wrong_type = [
            binding for binding in bindings if not isinstance(binding, EnvVarSecret)
        ]
        if wrong_type:
            kinds = sorted({type(binding).__name__ for binding in wrong_type})
            raise TypeError(
                f"Function secrets must be EnvVarSecret values, not {kinds!r}; a "
                "credential value is never sent to this API"
            )
        variables = [binding.env_variable for binding in bindings]
        duplicates = sorted({name for name in variables if variables.count(name) > 1})
        if duplicates:
            raise ValueError(
                "a Function binds each environment variable once; duplicated: "
                f"{duplicates!r}"
            )
        # `env` is ordinary configuration carried in the definition, so a name in
        # both would have a value visible in the Function's record and a value
        # that is not. Refuse rather than pick.
        environment = self._request.runtime.env or {}
        overlap = sorted(set(environment) & set(variables))
        if overlap:
            raise ValueError(
                f"Function env and secret bindings must be disjoint: {overlap!r}"
            )
        if not bindings:
            return self._request
        # Sorted, because the list is carried in the FunctionVersion hash and a
        # caller's argument order is not part of what a Function is.
        resolved = tuple(
            sorted(
                (
                    SecretBinding(
                        kind="env",
                        variable=binding.env_variable,
                        secret_ref=SecretReference(
                            name=binding.secret_name,
                            namespace_path=tuple(binding.secret_namespace_path),
                        ),
                    )
                    for binding in bindings
                ),
                key=lambda binding: (binding.kind, binding.variable or ""),
            )
        )
        return self._request._copy(update={"secret_bindings": resolved})

    def __call__(self, *args, **kwargs):
        return self._function(*args, **kwargs)


@overload
def udf(function: Callable[..., Any]) -> UdfDefinition: ...


@overload
def udf(
    function: None = None,
    *,
    name: Optional[str] = None,
    input_schema: Optional[pa.Schema] = None,
    output_schema: Optional[pa.DataType | pa.Field | pa.Schema] = None,
    pip: tuple[str, ...] | list[str] = (),
    env: Optional[Mapping[str, str]] = None,
    python_version: Optional[str] = None,
    gpu: bool = False,
    conda: tuple[str, ...] | list[str] = (),
    conda_channels: tuple[str, ...] | list[str] = (),
    code: Sequence[types.ModuleType] = (),
) -> Callable[[Callable[..., Any]], UdfDefinition]: ...


def udf(
    function: Optional[Callable[..., Any]] = None,
    *,
    name: Optional[str] = None,
    input_schema: Optional[pa.Schema] = None,
    output_schema: Optional[pa.DataType | pa.Field | pa.Schema] = None,
    pip: tuple[str, ...] | list[str] = (),
    env: Optional[Mapping[str, str]] = None,
    python_version: Optional[str] = None,
    gpu: bool = False,
    conda: tuple[str, ...] | list[str] = (),
    conda_channels: tuple[str, ...] | list[str] = (),
    code: Sequence[types.ModuleType] = (),
):
    """Prepare a Python function or class for remote Function registration.

    Input and output signatures are inferred from supported annotations. For
    Arrow types annotations cannot express precisely, pass ``input_schema``
    and ``output_schema`` together. Scalar outputs must be non-nullable. Every
    named-struct field may be nullable; Enterprise preserves the struct's
    validity when the result is expanded into sibling columns.

    **Functions and classes.** A decorated function is called once per row,
    or once per batch when its parameters are annotated ``pyarrow.Array`` or
    it takes one ``pyarrow.RecordBatch``. A decorated class is the same
    callable with state: each remote instance calls ``__init__`` once, calls
    ``__call__`` for every row or batch of every input it processes, and calls
    ``close()``, if defined, once when it retires. Load models, open clients,
    and build rate limiters in ``__init__``; every batch of the instance then
    reuses them. An instance serves many batches of one binding, so state is
    shared across them, but not across instances or workers.

    **Initialization.** The annotated parameters of a class's ``__init__``
    are the Function's initialization fields. Their values belong to each
    binding, not to the Function: ``function(text=col("body"), model="small")``
    binds column ``body`` and initializes every instance with ``model="small"``,
    and another column can bind the same Function version with other values.
    Initialization fields take booleans, integers, floats, strings, and lists
    or structs of them (``Annotated[dict, pa.struct(...)]``). A parameter with
    a default may be omitted from a binding; a null value then takes the
    default. Use Secrets, not initialization, for credentials.

    **What travels with the Function.** The artifact is a snapshot of the
    decorated source plus exactly the module-level names it references:

    - modules, and classes and functions importable from an installed
      package, become imports; the package must be declared in ``pip`` or
      ``conda``;
    - literals (``None``, booleans, numbers, strings, bytes, tuples of them)
      are inlined;
    - functions and classes defined in ``__main__``, such as notebook cells,
      are packaged by source, together with what they reference;
    - modules and packages listed in ``code`` are shipped as Python source and
      imported normally, so helpers shared by several Functions live in one
      place.

    A reference to a module in a local source tree that is neither shipped
    with ``code`` nor named by a declared package draws a warning at
    registration, because the worker will most likely fail to import it.
    Closures, lambdas, nested definitions, mutable module-level
    objects, code that reaches the module namespace another way
    (``globals()``/``eval``, ``sys.modules``, ``builtins``), and a
    non-standard ``__builtins__`` are rejected where they can be seen and
    otherwise unsupported. Changing any packaged source, including a module in
    ``code``, produces a new Function version; the environment built from
    ``pip`` or ``conda`` is reused.

    Parameters
    ----------
    function : function or class, optional
        The synchronous callable, or the class whose instances are called.
    name : str, optional
        The remote Function name. Defaults to the callable name.
    input_schema : pyarrow.Schema, optional
        Explicit input fields in the exact order of the callable parameters.
        Must be provided together with ``output_schema``.
    output_schema : pyarrow.DataType, pyarrow.Field, or pyarrow.Schema, optional
        Explicit scalar or named-struct output. Scalar outputs must be
        non-nullable. Must be provided together with ``input_schema``.
    pip : sequence of str, optional
        Pip requirements for the remote environment.
    conda : sequence of str, optional
        Conda packages for the remote environment, instead of ``pip``.
    conda_channels : sequence of str, optional
        Conda channels in priority order; requires ``conda``.
    env : mapping of str to str, optional
        Environment variables included in the Function definition. Not for
        credentials -- these are ordinary configuration, stored with the
        Function and visible wherever it is. Prefer initialization for
        values that differ between bindings.
    python_version : str, optional
        Remote Python major/minor version. Defaults to the client version.
    gpu : bool, default False
        Whether every remote execution requires a GPU. The execution platform
        selects one compatible GPU for each worker. The requirement is part of
        the immutable Function version.
    code : sequence of module, optional
        Top-level modules or packages whose importable ``.py`` files ship
        with the Function and are importable by their own names.

    Returns
    -------
    UdfDefinition
        A callable definition accepted by
        :meth:`lancedb.db.DBConnection.create_function`,
        :meth:`lancedb.db.AsyncConnection.create_function_async` and
        :meth:`lancedb.db.DBConnection.create_function_async`.

    Examples
    --------
    >>> from lancedb import udf
    >>> @udf(pip=["numpy==2.2.0"])
    ... def score(value: float) -> float:
    ...     return value * 2
    >>> score(1.5)
    3.0
    >>> @udf(pip=["cupy-cuda12x"], gpu=True)
    ... def gpu_score(value: int) -> int:
    ...     return value * 2
    >>> gpu_score.registration_request.runtime.gpu
    True
    >>> @udf
    ... class scale:
    ...     def __init__(self, factor: float, offset: float = 0.0):
    ...         self.factor, self.offset = factor, offset
    ...     def __call__(self, value: float) -> float:
    ...         return value * self.factor + self.offset
    >>> scale(factor=2.0)(1.5)
    3.0
    >>> [field.name for field in scale.registration_request.signature.initialization]
    ['factor', 'offset']
    """

    def decorate(target: Any) -> UdfDefinition:
        return UdfDefinition(
            target,
            name=name,
            input_schema=input_schema,
            output_schema=output_schema,
            pip=tuple(pip),
            env={} if env is None else env,
            python_version=python_version,
            gpu=gpu,
            conda=tuple(conda),
            conda_channels=tuple(conda_channels),
            code=tuple(code),
        )

    if function is None:
        return decorate
    return decorate(function)


__all__ = [
    "AssignmentMapping",
    "ApplicationInput",
    "FunctionApplication",
    "FunctionArtifact",
    "FunctionArtifactContent",
    "FunctionArtifactRequest",
    "FunctionBinding",
    "FunctionOutput",
    "FunctionParameter",
    "FunctionRegistrationRequest",
    "FunctionResultField",
    "FunctionSignature",
    "FunctionImage",
    "FunctionVersion",
    "FunctionVersionRef",
    "InputBinding",
    "OutputMapping",
    "PythonEnvironmentSpec",
    "PythonAdapterSpec",
    "PythonRuntimeSpec",
    "RefreshColumnResult",
    "UdfDefinition",
    "udf",
]