# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import io
import os
import struct
import sys
import tempfile
import types
from pathlib import Path

import pyarrow as pa
import pytest
from datafusion import SessionContext

from pypaimon_rust.datafusion import PaimonCatalog, PythonScalarUDF, SQLContext, udf

WAREHOUSE = os.environ.get("PAIMON_TEST_WAREHOUSE", "/tmp/paimon-warehouse")
PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
BLOB_DESCRIPTOR_MAGIC = 0x424C4F4244455343


def serialize_blob_descriptor(uri: str, offset: int, length: int) -> bytes:
    uri_bytes = uri.encode("utf-8")
    return (
        struct.pack("<BQI", 2, BLOB_DESCRIPTOR_MAGIC, len(uri_bytes))
        + uri_bytes
        + struct.pack("<qq", offset, length)
    )


def write_sample_video(
    path: Path,
    colors: tuple[tuple[int, int, int], ...] = ((240, 40, 80),),
) -> None:
    av = pytest.importorskip("av")
    image_module = pytest.importorskip("PIL.Image")

    with av.open(str(path), mode="w") as container:
        stream = container.add_stream("mpeg4", rate=1)
        stream.width = 32
        stream.height = 32
        stream.pix_fmt = "yuv420p"

        for color in colors:
            image = image_module.new("RGB", (32, 32), color=color)
            frame = av.VideoFrame.from_image(image)
            for packet in stream.encode(frame):
                container.mux(packet)
        for packet in stream.encode():
            container.mux(packet)


def sample_image_bytes() -> bytes:
    image_module = pytest.importorskip("PIL.Image")

    output = io.BytesIO()
    image = image_module.new("RGB", (32, 32), color=(40, 120, 220))
    image.save(output, format="PNG")
    return output.getvalue()


def extract_rows(batches):
    table = pa.Table.from_batches(batches)
    return sorted(zip(table["id"].to_pylist(), table["name"].to_pylist()))


def test_video_snapshot_builtin_registered_on_context_init():
    ctx = SQLContext()

    batches = ctx.sql("SELECT video_snapshot(CAST(NULL AS BYTEA)) AS cover_png")
    table = pa.Table.from_batches(batches)

    assert table["cover_png"].to_pylist() == [None]


def test_sql_context_survives_video_snapshot_registration_failure(monkeypatch):
    monkeypatch.setitem(
        sys.modules,
        "pypaimon_rust.functions",
        types.SimpleNamespace(),
    )

    with pytest.warns(
        RuntimeWarning,
        match="video_snapshot built-in could not be registered",
    ):
        ctx = SQLContext()

    batches = ctx.sql("SELECT 1 AS value")
    table = pa.Table.from_batches(batches)
    assert table["value"].to_pylist() == [1]


def test_video_snapshot_builtin_auto_registered_for_sql():
    with tempfile.TemporaryDirectory() as warehouse:
        video_path = Path(warehouse) / "sample.mp4"
        write_sample_video(video_path)
        video_bytes = video_path.read_bytes()

        ctx = SQLContext()
        ctx.register_catalog("paimon", {"warehouse": warehouse})
        ctx.register_batch(
            "paimon.default.videos",
            pa.record_batch(
                [[1], pa.array([video_bytes], type=pa.binary())],
                names=["id", "video"],
            ),
        )

        batches = ctx.sql(
            """
            SELECT id, video_snapshot(video) AS cover_png
            FROM paimon.default.videos
            """
        )
        table = pa.Table.from_batches(batches)

        assert table["id"].to_pylist() == [1]
        assert table["cover_png"].to_pylist()[0].startswith(PNG_SIGNATURE)

        ctx.sql("DROP TEMPORARY TABLE paimon.default.videos")


def test_video_snapshot_descriptor_without_table_file_io_returns_null():
    with tempfile.TemporaryDirectory() as warehouse:
        video_path = Path(warehouse) / "sample.mp4"
        write_sample_video(video_path)
        descriptor = serialize_blob_descriptor(
            str(video_path), 0, video_path.stat().st_size
        )

        ctx = SQLContext()
        ctx.register_catalog("paimon", {"warehouse": warehouse})
        ctx.register_batch(
            "paimon.default.videos",
            pa.record_batch(
                [[1], pa.array([descriptor], type=pa.binary())],
                names=["id", "video"],
            ),
        )

        batches = ctx.sql(
            """
            SELECT id, video_snapshot(video) AS cover_png
            FROM paimon.default.videos
            """
        )
        table = pa.Table.from_batches(batches)

        assert table["id"].to_pylist() == [1]
        assert table["cover_png"].to_pylist() == [None]

        ctx.sql("DROP TEMPORARY TABLE paimon.default.videos")


def test_video_snapshot_returns_null_for_image_bytes():
    with tempfile.TemporaryDirectory() as warehouse:
        ctx = SQLContext()
        ctx.register_catalog("paimon", {"warehouse": warehouse})
        ctx.register_batch(
            "paimon.default.media",
            pa.record_batch(
                [[1], pa.array([sample_image_bytes()], type=pa.binary())],
                names=["id", "content"],
            ),
        )

        batches = ctx.sql(
            """
            SELECT id, video_snapshot(content) AS cover_png
            FROM paimon.default.media
            """
        )
        table = pa.Table.from_batches(batches)

        assert table["id"].to_pylist() == [1]
        assert table["cover_png"].to_pylist() == [None]

        ctx.sql("DROP TEMPORARY TABLE paimon.default.media")


def test_video_snapshot_reads_descriptor_with_table_file_io():
    with tempfile.TemporaryDirectory() as warehouse:
        ctx = SQLContext()
        ctx.register_catalog("paimon", {"warehouse": warehouse})
        ctx.sql("CREATE TABLE paimon.default.videos (id INT, video BINARY)")

        video_path = Path(warehouse) / "default.db" / "videos" / "sample.mp4"
        video_path.parent.mkdir(parents=True, exist_ok=True)
        write_sample_video(video_path)
        descriptor = serialize_blob_descriptor(
            str(video_path), 0, video_path.stat().st_size
        )

        ctx.register_batch(
            "source_videos",
            pa.record_batch(
                [[1], pa.array([descriptor], type=pa.binary())],
                names=["id", "video"],
            ),
        )
        ctx.sql(
            """
            INSERT INTO paimon.default.videos
            SELECT id, video FROM paimon.default.source_videos
            """
        )

        batches = ctx.sql(
            """
            SELECT id, video_snapshot(video) AS cover_png
            FROM paimon.default.videos
            """
        )
        table = pa.Table.from_batches(batches)

        assert table["id"].to_pylist() == [1]
        assert table["cover_png"].to_pylist()[0].startswith(PNG_SIGNATURE)

        ctx.sql("DROP TEMPORARY TABLE paimon.default.source_videos")
        ctx.sql("DROP TABLE paimon.default.videos")


def test_video_snapshot_accepts_timestamp_ms():
    image_module = pytest.importorskip("PIL.Image")

    with tempfile.TemporaryDirectory() as warehouse:
        video_path = Path(warehouse) / "sample.mp4"
        write_sample_video(video_path, colors=((240, 40, 80), (40, 220, 80)))
        video_bytes = video_path.read_bytes()

        ctx = SQLContext()
        ctx.register_catalog("paimon", {"warehouse": warehouse})
        ctx.register_batch(
            "paimon.default.videos",
            pa.record_batch(
                [[1], pa.array([video_bytes], type=pa.binary())],
                names=["id", "video"],
            ),
        )

        batches = ctx.sql(
            """
            SELECT
                video_snapshot(video) AS first_png,
                video_snapshot(video, CAST(1000 AS INT)) AS second_png,
                video_snapshot(video, 5000) AS beyond_duration_png
            FROM paimon.default.videos
            """
        )
        row = pa.Table.from_batches(batches).to_pylist()[0]

        assert row["first_png"].startswith(PNG_SIGNATURE)
        assert row["second_png"].startswith(PNG_SIGNATURE)

        first_image = image_module.open(io.BytesIO(row["first_png"])).convert("RGB")
        second_image = image_module.open(io.BytesIO(row["second_png"])).convert("RGB")
        assert first_image.getpixel((16, 16)) != second_image.getpixel((16, 16))
        assert row["beyond_duration_png"].startswith(PNG_SIGNATURE)
        beyond_duration_image = image_module.open(
            io.BytesIO(row["beyond_duration_png"])
        ).convert("RGB")
        assert beyond_duration_image.getpixel((16, 16)) == second_image.getpixel((16, 16))

        ctx.sql("DROP TEMPORARY TABLE paimon.default.videos")


def test_query_simple_table_via_catalog_provider():
    catalog = PaimonCatalog({"warehouse": WAREHOUSE})
    ctx = SessionContext()
    ctx.register_catalog_provider("paimon", catalog)

    df = ctx.sql("SELECT id, name FROM paimon.default.simple_log_table")

    assert extract_rows(df.collect()) == [
        (1, "alice"),
        (2, "bob"),
        (3, "carol"),
    ]


def test_sql_context_ddl_dml():
    with tempfile.TemporaryDirectory() as warehouse:
        ctx = SQLContext()
        ctx.register_catalog("paimon", {"warehouse": warehouse})

        ctx.sql("CREATE SCHEMA paimon.test_db")
        ctx.sql(
            "CREATE TABLE paimon.test_db.users "
            "(id INT, name STRING, PRIMARY KEY (id))"
        )

        ctx.sql("INSERT INTO paimon.test_db.users VALUES (1, 'alice'), (2, 'bob')")

        batches = ctx.sql("SELECT id, name FROM paimon.test_db.users")
        table = pa.Table.from_batches(batches)
        rows = sorted(zip(table["id"].to_pylist(), table["name"].to_pylist()))
        assert rows == [(1, "alice"), (2, "bob")]

        ctx.sql("DROP TABLE paimon.test_db.users")
        ctx.sql("DROP SCHEMA paimon.test_db")


def test_register_batch_fully_qualified():
    with tempfile.TemporaryDirectory() as warehouse:
        ctx = SQLContext()
        ctx.register_catalog("paimon", {"warehouse": warehouse})

        batch = pa.record_batch([[1, 2], ["alice", "bob"]], names=["id", "name"])
        ctx.register_batch("paimon.default.my_temp", batch)

        batches = ctx.sql("SELECT id, name FROM paimon.default.my_temp")
        table = pa.Table.from_batches(batches)
        rows = sorted(zip(table["id"].to_pylist(), table["name"].to_pylist()))
        assert rows == [(1, "alice"), (2, "bob")]

        ctx.sql("DROP TEMPORARY TABLE paimon.default.my_temp")


def test_register_batch_bare_name():
    with tempfile.TemporaryDirectory() as warehouse:
        ctx = SQLContext()
        ctx.register_catalog("paimon", {"warehouse": warehouse})

        batch = pa.record_batch([[1, 2], ["alice", "bob"]], names=["id", "name"])
        # Bare name uses current catalog and current database
        ctx.register_batch("my_temp", batch)

        batches = ctx.sql("SELECT id, name FROM paimon.default.my_temp")
        table = pa.Table.from_batches(batches)
        rows = sorted(zip(table["id"].to_pylist(), table["name"].to_pylist()))
        assert rows == [(1, "alice"), (2, "bob")]

        ctx.sql("DROP TEMPORARY TABLE paimon.default.my_temp")


def test_register_udf_from_python():
    with tempfile.TemporaryDirectory() as warehouse:
        ctx = SQLContext()
        ctx.register_catalog("paimon", {"warehouse": warehouse})

        batch = pa.record_batch([[1, None, 3]], names=["id"])
        ctx.register_batch("my_temp", batch)

        def plus_ten(values):
            return pa.array(
                [None if value is None else value + 10 for value in values.to_pylist()],
                type=pa.int64(),
            )

        ctx.register_udf(udf(plus_ten, [pa.int64()], pa.int64(), "volatile", "plus_ten"))

        batches = ctx.sql(
            "SELECT plus_ten(id) AS id FROM paimon.default.my_temp ORDER BY id"
        )
        table = pa.Table.from_batches(batches)
        assert table["id"].to_pylist() == [11, 13, None]

        ctx.sql("DROP TEMPORARY TABLE paimon.default.my_temp")


def test_register_udf_default_name_is_sql_identifier_for_closure():
    with tempfile.TemporaryDirectory() as warehouse:
        ctx = SQLContext()
        ctx.register_catalog("paimon", {"warehouse": warehouse})

        batch = pa.record_batch([[1, 2]], names=["id"])
        ctx.register_batch("my_temp", batch)

        def build_udf():
            def plus_one(values):
                return pa.array(
                    [value + 1 for value in values.to_pylist()], type=pa.int64()
                )

            return plus_one

        scalar_udf = udf(build_udf(), [pa.int64()], pa.int64(), "volatile")
        assert scalar_udf.name == "plus_one"
        ctx.register_udf(scalar_udf)

        batches = ctx.sql(
            "SELECT plus_one(id) AS id FROM paimon.default.my_temp ORDER BY id"
        )
        table = pa.Table.from_batches(batches)
        assert table["id"].to_pylist() == [2, 3]

        ctx.sql("DROP TEMPORARY TABLE paimon.default.my_temp")


def test_register_udf_multiple_arguments():
    with tempfile.TemporaryDirectory() as warehouse:
        ctx = SQLContext()
        ctx.register_catalog("paimon", {"warehouse": warehouse})

        batch = pa.record_batch(
            [
                pa.array([1, 2, None], type=pa.int64()),
                pa.array([10, 20, 30], type=pa.int64()),
            ],
            names=["id", "delta"],
        )
        ctx.register_batch("my_temp", batch)

        def add_values(left, right):
            values = []
            for left_value, right_value in zip(left.to_pylist(), right.to_pylist()):
                if left_value is None or right_value is None:
                    values.append(None)
                else:
                    values.append(left_value + right_value)
            return pa.array(values, type=pa.int64())

        ctx.register_udf(
            udf(
                add_values,
                [pa.int64(), pa.int64()],
                pa.int64(),
                "volatile",
                "add_values",
            )
        )

        batches = ctx.sql(
            """
            SELECT add_values(id, delta) AS value
            FROM paimon.default.my_temp
            ORDER BY id
            """
        )
        table = pa.Table.from_batches(batches)
        assert table["value"].to_pylist() == [11, 22, None]

        ctx.sql("DROP TEMPORARY TABLE paimon.default.my_temp")


def test_register_udf_multi_partition_union_plan():
    with tempfile.TemporaryDirectory() as warehouse:
        ctx = SQLContext()
        ctx.register_catalog("paimon", {"warehouse": warehouse})

        batch = pa.record_batch([[1, 2, 3]], names=["id"])
        ctx.register_batch("my_temp", batch)

        def plus_ten(values):
            return pa.array([value + 10 for value in values.to_pylist()], type=pa.int64())

        ctx.register_udf(udf(plus_ten, [pa.int64()], pa.int64(), "volatile", "plus_ten"))

        batches = ctx.sql(
            """
            SELECT plus_ten(id) AS id FROM paimon.default.my_temp
            UNION ALL
            SELECT plus_ten(id) AS id FROM paimon.default.my_temp
            ORDER BY id
            """
        )
        table = pa.Table.from_batches(batches)
        assert table["id"].to_pylist() == [11, 11, 12, 12, 13, 13]

        ctx.sql("DROP TEMPORARY TABLE paimon.default.my_temp")


def test_udf_rejects_non_callable():
    try:
        udf(1, [pa.int64()], pa.int64(), "volatile")
        pytest.fail("expected non-callable UDF creation to fail")
    except TypeError as e:
        assert "`func` argument must be callable" in str(e)


def test_udf_rejects_unsupported_type():
    def identity(values):
        return values

    try:
        udf(identity, [object()], pa.int64(), "volatile", "identity")
        pytest.fail("expected unsupported type registration to fail")
    except TypeError as e:
        assert "Expected a pyarrow.DataType" in str(e)


def test_python_scalar_udf_constructor_matches_datafusion_shape():
    def identity(values):
        return values

    scalar_udf = PythonScalarUDF(
        "identity", identity, [pa.field("value", pa.int64())], pa.int64(), "stable"
    )

    assert scalar_udf.name == "identity"
    assert repr(scalar_udf) == "PythonScalarUDF(identity)"


def test_python_udf_exception_surfaces():
    with tempfile.TemporaryDirectory() as warehouse:
        ctx = SQLContext()
        ctx.register_catalog("paimon", {"warehouse": warehouse})
        ctx.register_batch("my_temp", pa.record_batch([[1]], names=["id"]))

        def boom(values):
            raise RuntimeError("boom")

        ctx.register_udf(udf(boom, [pa.int64()], pa.int64(), "volatile", "boom"))

        try:
            ctx.sql("SELECT boom(id) AS id FROM paimon.default.my_temp")
            pytest.fail("expected Python UDF exception to fail the query")
        except Exception as e:
            message = str(e)
            assert "Python UDF 'boom' failed" in message
            assert "boom" in message


def test_python_udf_rejects_wrong_length():
    with tempfile.TemporaryDirectory() as warehouse:
        ctx = SQLContext()
        ctx.register_catalog("paimon", {"warehouse": warehouse})
        ctx.register_batch("my_temp", pa.record_batch([[1, 2]], names=["id"]))

        def wrong_length(values):
            return pa.array([1], type=pa.int64())

        ctx.register_udf(
            udf(wrong_length, [pa.int64()], pa.int64(), "volatile", "wrong_length")
        )

        try:
            ctx.sql("SELECT wrong_length(id) AS id FROM paimon.default.my_temp")
            pytest.fail("expected wrong-length UDF result to fail the query")
        except Exception as e:
            message = str(e)
            assert "Python UDF 'wrong_length' returned 1 rows, expected 2" in message


def test_python_udf_rejects_wrong_type():
    with tempfile.TemporaryDirectory() as warehouse:
        ctx = SQLContext()
        ctx.register_catalog("paimon", {"warehouse": warehouse})
        ctx.register_batch("my_temp", pa.record_batch([[1]], names=["id"]))

        def wrong_type(values):
            return pa.array(["not an int"], type=pa.string())

        ctx.register_udf(
            udf(wrong_type, [pa.int64()], pa.int64(), "volatile", "wrong_type")
        )

        try:
            ctx.sql("SELECT wrong_type(id) AS id FROM paimon.default.my_temp")
            pytest.fail("expected wrong-type UDF result to fail the query")
        except Exception as e:
            message = str(e)
            assert "Python UDF 'wrong_type' returned Utf8, expected Int64" in message


def test_temp_table_shadows_paimon_table():
    with tempfile.TemporaryDirectory() as warehouse:
        ctx = SQLContext()
        ctx.register_catalog("paimon", {"warehouse": warehouse})

        ctx.sql("CREATE SCHEMA paimon.test_db")
        ctx.sql("CREATE TABLE paimon.test_db.users (id INT, name STRING)")
        ctx.sql("INSERT INTO paimon.test_db.users VALUES (1, 'real')")

        batch = pa.record_batch([[2], ["temp"]], names=["id", "name"])
        ctx.register_batch("paimon.test_db.users", batch)

        # Temp table should shadow the real Paimon table
        batches = ctx.sql("SELECT id, name FROM paimon.test_db.users")
        table = pa.Table.from_batches(batches)
        rows = sorted(zip(table["id"].to_pylist(), table["name"].to_pylist()))
        assert rows == [(2, "temp")]

        ctx.sql("DROP TEMPORARY TABLE paimon.test_db.users")

        # After dropping, the real table is visible again
        batches = ctx.sql("SELECT id, name FROM paimon.test_db.users")
        table = pa.Table.from_batches(batches)
        rows = sorted(zip(table["id"].to_pylist(), table["name"].to_pylist()))
        assert rows == [(1, "real")]

        ctx.sql("DROP TABLE paimon.test_db.users")
        ctx.sql("DROP SCHEMA paimon.test_db")


def test_drop_temp_table_if_exists():
    with tempfile.TemporaryDirectory() as warehouse:
        ctx = SQLContext()
        ctx.register_catalog("paimon", {"warehouse": warehouse})

        batch = pa.record_batch([[1]], names=["id"])
        ctx.register_batch("paimon.default.my_temp", batch)

        ctx.sql("DROP TEMPORARY TABLE IF EXISTS paimon.default.my_temp")

        # Should be able to drop again without error
        ctx.sql("DROP TEMPORARY TABLE IF EXISTS paimon.default.my_temp")


def test_multi_catalog_temp_table():
    with tempfile.TemporaryDirectory() as wh1, tempfile.TemporaryDirectory() as wh2:
        ctx = SQLContext()
        ctx.register_catalog("cat1", {"warehouse": wh1})
        ctx.register_catalog("cat2", {"warehouse": wh2})

        batch1 = pa.record_batch([[1]], names=["id"])
        batch2 = pa.record_batch([[2]], names=["id"])

        ctx.register_batch("cat1.default.t1", batch1)
        ctx.register_batch("cat2.default.t2", batch2)

        result1 = ctx.sql("SELECT id FROM cat1.default.t1")
        assert pa.Table.from_batches(result1)["id"].to_pylist() == [1]

        result2 = ctx.sql("SELECT id FROM cat2.default.t2")
        assert pa.Table.from_batches(result2)["id"].to_pylist() == [2]

        ctx.sql("DROP TEMPORARY TABLE cat1.default.t1")
        ctx.sql("DROP TEMPORARY TABLE cat2.default.t2")


def test_register_batch_invalid_catalog():
    with tempfile.TemporaryDirectory() as warehouse:
        ctx = SQLContext()
        ctx.register_catalog("paimon", {"warehouse": warehouse})

        batch = pa.record_batch([[1]], names=["id"])
        try:
            ctx.register_batch("unknown_catalog.default.my_temp", batch)
            pytest.fail("Expected an error for unknown catalog")
        except Exception as e:
            assert "unknown_catalog" in str(e).lower() or "not a paimon" in str(e).lower() or "unknown" in str(e).lower()


def test_table_functions_registered_with_catalog():
    """register_catalog auto-registers vector_search / full_text_search as
    UDTFs. Calling one with the wrong argument count surfaces the function's
    own validation error, which proves it is registered — an unregistered
    name would instead fail with 'table function not found'."""
    with tempfile.TemporaryDirectory() as warehouse:
        ctx = SQLContext()
        ctx.register_catalog("paimon", {"warehouse": warehouse})

        for fn in ("vector_search", "full_text_search"):
            try:
                ctx.sql(f"SELECT * FROM {fn}('only_one_arg')")
                pytest.fail(f"expected {fn} to reject a single argument")
            except Exception as e:
                assert "requires 4 arguments" in str(e), str(e)


def test_list_databases_and_tables():
    catalog = PaimonCatalog({"warehouse": WAREHOUSE})

    assert "default" in catalog.list_databases()

    tables = catalog.list_tables("default")
    assert "simple_log_table" in tables

    table = catalog.get_table("default.simple_log_table")
    assert table.identifier() == "default.simple_log_table"
    assert table.location().endswith("/default.db/simple_log_table") or table.location()

    schema = table.schema()
    field_names = [f.name() for f in schema.fields()]
    assert "id" in field_names
    assert "name" in field_names
    # simple_log_table is non-partitioned, so partition keys are empty.
    assert schema.partition_keys() == []