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

"""Examples for docs/web/management/function.mdx.

Registering a Function and filling its column need a LanceDB Enterprise
deployment, so `test_function_enterprise` runs only when the connection
variables the Enterprise examples already use are set: LANCEDB_URI,
LANCEDB_API_KEY and LANCEDB_HOST_OVERRIDE, plus LANCEDB_REGION if the
deployment needs one. The walkthrough calls its table `function_demo` and its
Function `double`; the test creates both under names unique to its run, never
replaces an existing object, and drops only what it created.

The other tests run locally and check the client side of the same code: what
`@udf` sends for registration, how a registered version binds a column, and
where a local connection stops. They do not run a Function remotely.
"""

import json
import os
import sys
import uuid
from pathlib import Path

import lancedb
import pyarrow as pa
import pytest
from lancedb import col
from lancedb.functions import FunctionApplication, FunctionVersion

# --8<-- [start:function_define]
from lancedb import udf


@udf
def double(value: float) -> float:
    return value * 2


# --8<-- [end:function_define]

# The definition the walkthrough registers. The Enterprise test registers a copy
# of it under a name that only its run uses.
WALKTHROUGH_DOUBLE = double


def describe(value: float):
    return {"half": value / 2, "note": None}


# A Function version as the service returns it after registration: the SDK's
# shared golden reply, with the name and signature of the definition above.
SERVICE_REPLY = (
    Path(__file__).parents[3]
    / "rust"
    / "lancedb"
    / "tests"
    / "fixtures"
    / "first_class_functions"
    / "v1"
    / "remote_function_job.json"
)


def registered_version(definition) -> FunctionVersion:
    request = definition.registration_request
    reply = json.loads(SERVICE_REPLY.read_text())["result"]
    reply["name"] = request.name
    reply["signature"] = json.loads(request.signature.to_canonical_json())
    return FunctionVersion.from_json(json.dumps(reply))


def require_env(var_name: str) -> str:
    """Skip the test unless the required environment variable is provided."""
    value = os.environ.get(var_name)
    if not value:
        pytest.skip(f"Set {var_name} to run this Enterprise example")
    return value


def test_function_runs_locally_as_python():
    # --8<-- [start:function_local_call]
    print(double(1.5))
    # 3.0
    # --8<-- [end:function_local_call]
    assert double(1.5) == 3.0


def test_function_registration_request():
    request = double.registration_request
    assert request.name == "double"
    assert [
        (parameter.name, parameter.arrow_type, parameter.nullable)
        for parameter in request.signature.inputs
    ] == [("value", "float64", False)]
    output = request.signature.output
    assert (output.kind, output.arrow_type, output.nullable) == (
        "scalar",
        "float64",
        False,
    )
    version = request.runtime.python_version
    assert version == f"{sys.version_info.major}.{sys.version_info.minor}"
    assert request.runtime.environment.packages == ()


def test_function_output_nullability():
    inputs = pa.schema([pa.field("value", pa.float64(), nullable=False)])

    # A named-struct output may declare nullable fields.
    fields = [
        pa.field("half", pa.float64(), nullable=False),
        pa.field("note", pa.string(), nullable=True),
    ]
    struct = udf(input_schema=inputs, output_schema=pa.schema(fields))(describe)
    output = struct.registration_request.signature.output
    assert output.kind == "named_struct"
    assert [(field.name, field.nullable) for field in output.fields] == [
        ("half", False),
        ("note", True),
    ]

    # A scalar output, and a struct declared as one nullable field, may not be.
    for nullable in (
        pa.field("result", pa.float64(), nullable=True),
        pa.field("result", pa.struct(fields), nullable=True),
    ):
        with pytest.raises(ValueError, match="non-nullable"):
            udf(input_schema=inputs, output_schema=nullable)(describe)


def test_function_registration_is_remote(tmp_path):
    db = lancedb.connect(tmp_path)
    with pytest.raises(NotImplementedError):
        db.create_function(double)


def test_function_version_binds_a_column(tmp_path):
    version = registered_version(double)

    application = version(value=col("value"))

    assert isinstance(application, FunctionApplication)
    assert application.function.name == "double"
    assert application.function.version == version.version
    assert [
        (binding.parameter, binding.kind, binding.value["path"])
        for binding in application.inputs
    ] == [("value", "column", "value")]
    with pytest.raises(TypeError, match="direct col"):
        version(value=1.5)
    with pytest.raises(TypeError, match="missing"):
        version(x=col("value"))

    # The binding is accepted by add_columns, which a local table refuses:
    # a Function column needs LanceDB Cloud or Enterprise.
    table = lancedb.connect(tmp_path).create_table("t", [{"value": 1.0}])
    with pytest.raises(NotImplementedError, match="Cloud and Enterprise"):
        table.add_columns({"doubled": application})


def test_function_enterprise(monkeypatch):
    require_env("LANCEDB_URI")
    require_env("LANCEDB_API_KEY")
    require_env("LANCEDB_HOST_OVERRIDE")

    # --8<-- [start:function_connect]
    import os

    import lancedb

    db = lancedb.connect(
        uri=os.environ["LANCEDB_URI"],  # db://your-database
        api_key=os.environ["LANCEDB_API_KEY"],
        region=os.environ.get("LANCEDB_REGION", "us-east-1"),
        host_override=os.environ["LANCEDB_HOST_OVERRIDE"],
    )
    # --8<-- [end:function_connect]

    # In a shared database `function_demo` or `double` may already exist, so
    # this run creates both under names only it uses, and drops only those.
    run = uuid.uuid4().hex[:12]
    table_name = f"function_demo_{run}"
    double = udf(name=f"double_{run}")(WALKTHROUGH_DOUBLE.__wrapped__)
    create_table = db.create_table

    def create_run_table(name, *args, **kwargs):
        assert name == "function_demo"
        return create_table(table_name, *args, **kwargs)

    monkeypatch.setattr(db, "create_table", create_run_table)

    table = None
    version = None
    try:
        # --8<-- [start:function_create_table]
        import pyarrow as pa

        schema = pa.schema([pa.field("value", pa.float64(), nullable=False)])
        table = db.create_table(
            "function_demo",
            [{"value": 1.0}, {"value": 2.0}, {"value": 3.0}],
            schema=schema,
        )
        # --8<-- [end:function_create_table]

        # --8<-- [start:function_register]
        job = db.create_function_async(double)
        version = job.wait()
        # --8<-- [end:function_register]
        assert version.name == f"double_{run}"

        # --8<-- [start:function_apply]
        from lancedb import col

        table.add_columns({"doubled": version(value=col("value"))})
        # --8<-- [end:function_apply]

        # --8<-- [start:function_refresh]
        refresh = table.refresh_column_async("doubled")
        result = refresh.wait()
        print(result.rows_assigned)
        # 3
        # --8<-- [end:function_refresh]
        assert result.rows_assigned == 3
        assert result.rows_failed == 0

        # --8<-- [start:function_inspect]
        rows = table.search().select(["value", "doubled"]).to_list()
        print(sorted((row["value"], row["doubled"]) for row in rows))
        # [(1.0, 2.0), (2.0, 4.0), (3.0, 6.0)]
        # --8<-- [end:function_inspect]
        assert sorted((row["value"], row["doubled"]) for row in rows) == [
            (1.0, 2.0),
            (2.0, 4.0),
            (3.0, 6.0),
        ]
    finally:
        if table is not None:
            db.drop_table(table_name)
        if version is not None:
            db.drop_function(version.name, version=version.version)