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

from collections.abc import Generator
from pathlib import Path

import lancedb
import numpy as np
import pandas as pd
import pyarrow as pa
import pytest
from pyarrow.fs import FileSystem, FileType

from geneva import connect, udf
from geneva.packager import DockerUDFPackager


def test_add_column(tmp_path: Path) -> None:
    db = connect(tmp_path)

    # create a basic table
    tbl = pa.Table.from_pydict({"id": [1, 2, 3, 4, 5, 6]})
    table = db.create_table("table1", tbl)

    # add a basic column
    table.add_columns(
        {"id2": "cast(null as string)"},
    )

    schema = table.schema

    assert len(schema) == 2
    field = schema.field("id2")
    assert field is not None
    assert field.type == pa.string()


@pytest.fixture(params=["geneva", "lance"])
def db(tmp_path: Path, request) -> lancedb.DBConnection:
    """Create a temporary database for testing."""
    if request.param == "geneva":
        return connect(tmp_path)
    else:
        return lancedb.connect(tmp_path)


def test_create_table_and_index(db: lancedb.DBConnection) -> None:
    assert len(list(db.table_names())) == 0

    df = pd.DataFrame(
        {
            "id": [1, 2, 3, 4, 5, 6],
            "name": ["Alice", "Bob", "Charlie", "David", "Eve", "Frank"],
        }
    )
    db.create_table("tbl", df)

    assert len(list(db.table_names())) == 1
    assert db.table_names()[0] == "tbl"

    tbl = db.open_table("tbl")
    assert tbl.count_rows() == 6

    df = pd.DataFrame(
        {
            "id": [11, 12, 13, 14, 15, 16],
            "name": ["alice", "bob", "charlie", "david", "eve", "frank"],
        }
    )
    tbl.add(df)
    assert tbl.count_rows() == 12

    tbl.create_scalar_index("id")
    tbl.create_fts_index("name", use_tantivy=False)

    tbl.delete("id % 2 == 0")
    assert tbl.version == 5
    assert tbl.to_arrow().combine_chunks() == pa.Table.from_pydict(
        {
            "id": [1, 3, 5, 11, 13, 15],
            "name": ["Alice", "Charlie", "Eve", "alice", "charlie", "eve"],
        }
    )
    assert tbl.count_rows() == 6


def test_create_vector_idx(db: lancedb.DBConnection) -> None:
    def producer() -> Generator[pa.Table, None, None]:
        dim = 128

        rng = np.random.default_rng()
        for i in range(100):
            ids = pa.array([i * 20 + j for j in range(20)])
            values = pa.array(rng.random(20 * dim).astype(np.float32))
            fsl = pa.FixedSizeListArray.from_arrays(values, dim)  # type: ignore
            yield pa.Table.from_arrays([ids, fsl], ["id", "vector"])

    tbl = db.create_table("table", producer())
    assert tbl.count_rows() == 2000
    tbl.create_index(num_sub_vectors=8)

    indices = tbl.list_indices()
    assert indices[0].index_type == "IvfPq"


def test_add_virtual_column(tmp_path: Path) -> None:
    packager = DockerUDFPackager(
        # use prebuilt image tag so we don't have to build the image
        prebuilt_docker_img="test-image:latest"
    )
    db = connect(tmp_path, packager=packager)

    # create a basic table
    tbl = pa.Table.from_pydict({"id": [1, 2, 3, 4, 5, 6]})
    table = db.create_table("table1", tbl)

    @udf(data_type=pa.int64())
    def double_id(id: int):  # noqa A002
        return id * 2

    table.add_columns(
        {"id2": double_id},  # implicit udf arg name mapping
    )

    schema = table.schema

    assert len(schema) == 2
    field = schema.field("id2")
    assert field is not None
    assert field.type == pa.int64()

    assert len(field.metadata) == 5
    assert field.metadata[b"virtual_column"] == b"true"
    assert field.metadata[b"virtual_column.udf_backend"] == b"DockerUDFSpecV1"
    assert field.metadata[b"virtual_column.udf_name"] == b"double_id"
    assert field.metadata[b"virtual_column.udf_inputs"] == b'["id"]'

    # check that the UDF was actually uploaded
    fs, root_path = FileSystem.from_uri(f"{str(tmp_path)}/table1.lance")
    file_info = fs.get_file_info(
        f"{root_path}/{field.metadata[b'virtual_column.udf'].decode('utf-8')}"
    )
    assert file_info.type is not FileType.NotFound

    try:
        import ray  # noqa: F401

        # before materializing, the virtual column should have nulls
        expected = pa.Table.from_pydict(
            {
                "id": [1, 2, 3, 4, 5, 6],
                "id2": pa.array([None] * 6, pa.int64()),
            }
        )
        assert table.to_arrow().equals(expected)

        # backfill and check values of virtual column
        table.backfill("id2")
        assert table.to_pandas().equals(
            pd.DataFrame(
                {
                    "id": [1, 2, 3, 4, 5, 6],
                    "id2": [2, 4, 6, 8, 10, 12],
                }
            )
        )
    except ImportError:
        "Ray not installed, skipping value checks part of test"


def test_alter_virtual_column(tmp_path: Path) -> None:
    packager = DockerUDFPackager(
        # use prebuilt image tag so we don't have to build the image
        prebuilt_docker_img="test-image:latest"
    )
    db = connect(tmp_path, packager=packager)

    # create a basic table
    tbl = pa.Table.from_pydict({"id": [1, 2, 3, 4, 5, 6]})
    table = db.create_table("table1", tbl)

    @udf(data_type=pa.int64())
    def double_id(id: int):  # noqa A002
        return id * 2

    table.add_columns(
        {"id2": double_id},  # implicit udf arg name mapping
    )

    schema = table.schema
    field = schema.field("id2")
    assert field.metadata[b"virtual_column.udf_name"] == b"double_id"

    try:
        import kubernetes  # noqa: F401
        import ray  # noqa: F401

        # before materializing, the virtual column should have nulls
        expected = pa.Table.from_pydict(
            {
                "id": [1, 2, 3, 4, 5, 6],
                "id2": pa.array([None] * 6, pa.int64()),
            }
        )
        assert table.to_arrow().equals(expected)

        # backfill and check values of virtual column
        table.backfill("id2")
        assert table.to_pandas().equals(
            pd.DataFrame(
                {
                    "id": [1, 2, 3, 4, 5, 6],
                    "id2": [2, 4, 6, 8, 10, 12],
                }
            )
        )
    except ImportError:
        "Ray not installed, skipping value checks part of test"

    # now check that we can replace the UDF with a new version:
    @udf(data_type=pa.int64())
    def triple_id(id: int):  # noqa A002
        return id * 3

    table.alter_columns(
        *[
            {
                "path": "id2",
                "virtual_column": triple_id,
            }
        ]
    )

    schema = table.schema
    field = schema.field("id2")
    assert field.metadata[b"virtual_column"] == b"true"
    assert field.metadata[b"virtual_column.udf_name"] == b"triple_id"
    assert field.metadata[b"virtual_column.udf_backend"] == b"DockerUDFSpecV1"

    # check that the UDF was actually uploaded
    fs, root_path = FileSystem.from_uri(f"{str(tmp_path)}/table1.lance")
    file_info = fs.get_file_info(
        f"{root_path}/{field.metadata[b'virtual_column.udf'].decode('utf-8')}"
    )
    assert file_info.type is not FileType.NotFound

    try:
        import kubernetes  # noqa: F401
        import ray  # noqa: F401

        # After the alter but before materializing, the virtual column should
        # not have the old UDF's values.  It should have nulls.
        expected = pa.Table.from_pydict(
            {
                "id": [1, 2, 3, 4, 5, 6],
                "id2": [2, 4, 6, 8, 10, 12],
                #  TODO This should be "id2": pa.array([None] * 6, pa.int64()),
            }
        )
        assert table.to_arrow().equals(expected)

        # backfill and check values of virtual column
        table.backfill("id2")
        assert table.to_pandas().equals(
            pd.DataFrame(
                {
                    "id": [1, 2, 3, 4, 5, 6],
                    "id2": [3, 6, 9, 12, 15, 18],
                }
            )
        )
    except ImportError:
        "Ray not installed, skipping value checks part of test"


@pytest.mark.skip(reason="overriding input parameter name is not supported")
def test_alter_virtual_column_diff_col_name(tmp_path: Path) -> None:
    packager = DockerUDFPackager(
        # use prebuilt image tag so we don't have to build the image
        prebuilt_docker_img="test-image:latest"
    )
    db = connect(tmp_path, packager=packager)

    # create a basic table
    tbl = pa.Table.from_pydict({"seq": [1, 2, 3, 4, 5, 6]})
    table = db.create_table("table1", tbl)

    @udf(data_type=pa.int64())
    def double_id(id: int):  # noqa A002
        return id * 2

    table.add_columns(
        {"id2": (double_id, ["seq"])},  # explicit udf arg name mapping
    )

    schema = table.schema
    field = schema.field("id2")
    assert field.metadata[b"virtual_column.udf_name"] == b"double_id"

    # before materializing, the virtual column should have nulls
    expected = pa.Table.from_pydict(
        {
            "seq": [1, 2, 3, 4, 5, 6],
            "id2": pa.array([None] * 6, pa.int64()),
        }
    )
    assert table.to_arrow().equals(expected)

    # backfill and check values of virtual column
    table.backfill("id2")
    assert table.to_pandas().equals(
        pd.DataFrame(
            {
                "seq": [1, 2, 3, 4, 5, 6],
                "id2": [2, 4, 6, 8, 10, 12],
            }
        )
    )

    # now check that we can replace the UDF with a new version:
    @udf(data_type=pa.int64())
    def triple_id(id: int):  # noqa A002
        return id * 3

    table.alter_columns(
        *[
            {
                "path": "id2",
                "virtual_column": triple_id,
                "input_columns": ["seq"],
            }
        ]
    )

    schema = table.schema
    field = schema.field("id2")
    assert field.metadata[b"virtual_column"] == b"true"
    assert field.metadata[b"virtual_column.udf_name"] == b"triple_id"
    assert field.metadata[b"virtual_column.udf_backend"] == b"DockerUDFSpecV1"

    # check that the UDF was actually uploaded
    fs, root_path = FileSystem.from_uri(f"{str(tmp_path)}/table1.lance")
    file_info = fs.get_file_info(
        f"{root_path}/{field.metadata[b'virtual_column.udf'].decode('utf-8')}"
    )
    assert file_info.type is not FileType.NotFound

    # After the alter but before materializing, the virtual column should
    # not have the old UDF's values.  It should have nulls.
    expected = pa.Table.from_pydict(
        {
            "seq": [1, 2, 3, 4, 5, 6],
            "id2": [2, 4, 6, 8, 10, 12],
            #  TODO This should be "id2": pa.array([None] * 6, pa.int64()),
        }
    )
    assert table.to_arrow().equals(expected)

    # backfill and check values of virtual column
    table.backfill("id2")
    assert table.to_pandas().equals(
        pd.DataFrame(
            {
                "seq": [1, 2, 3, 4, 5, 6],
                "id2": [3, 6, 9, 12, 15, 18],
            }
        )
    )
