# SPDX-License-Identifier: PROPRIETARY
# SPDX-FileCopyrightText: Copyright The Geneva Authors

import logging
import sys
from pathlib import Path
from typing import cast

import pyarrow as pa
import pytest

from geneva import connect, udf
from geneva.db import Connection
from geneva.runners.ray.pipeline import run_ray_copy_table
from geneva.table import Table

_LOG = logging.getLogger(__name__)
logging.basicConfig(level=logging.DEBUG, stream=sys.stderr, force=True)
sys.stderr.reconfigure(line_buffering=True)


TABLE = "table"

pytestmark = pytest.mark.ray


@pytest.fixture
def db(tmp_path: Path) -> Connection:
    return connect(tmp_path)


@pytest.fixture
def video_table(db) -> Table:
    tbl = pa.Table.from_pydict(
        {
            "video_uri": ["a", "b", "c", "d", "e", "f"],
            "rating": ["g", "nr", "pg", "pg-13", "r", "t"],
        }
    )
    table = db.create_table(TABLE, tbl)
    return table


def test_db_create_materialized_view(db, video_table) -> None:
    @udf(data_type=pa.binary())
    def load_video(video_uri: pa.Array) -> pa.Array:
        videos = [str(i).encode("utf-8") for i in video_uri]
        return cast("pa.Array", pa.array(videos))

    _LOG.info(video_table.to_arrow())

    q = (
        video_table.search(None)
        .shuffle(seed=42)
        .select(
            {
                "video_uri": "video_uri",
                "video": load_video,
            }
        )
    )

    dl_view = db.create_materialized_view("dl_view", q)
    udf_field = dl_view.schema.field("video")
    assert udf_field.metadata.get(b"virtual_column") == b"true"
    assert udf_field.metadata.get(b"virtual_column.udf_inputs") is not None
    assert udf_field.metadata.get(b"virtual_column.udf_name") == b"load_video"
    assert udf_field.metadata.get(b"virtual_column.udf_backend") is not None
    assert udf_field.metadata.get(b"virtual_column.udf") is not None

    cnt_not_null = dl_view.count_rows(filter="video is not null")
    assert cnt_not_null == 0

    dl_view.refresh()

    cnt_null = dl_view.count_rows(filter="video is null")
    _LOG.info(dl_view.to_arrow())
    assert cnt_null == 0


def test_create_materialized_view(db, video_table) -> None:
    @udf(data_type=pa.binary())
    def load_video(video_uri: pa.Array) -> pa.Array:
        videos = [str(i).encode("utf-8") for i in video_uri]
        return cast("pa.Array", pa.array(videos))

    view_table = (
        video_table.search(None)
        .shuffle(seed=42)
        .select(
            {
                "video_uri": "video_uri",
                "video": load_video,
            }
        )
        .create_materialized_view(db, "dl_view")
    )
    udf_field = view_table.schema.field("video")
    assert udf_field.metadata.get(b"virtual_column") == b"true"
    assert udf_field.metadata.get(b"virtual_column.udf_inputs") is not None
    assert udf_field.metadata.get(b"virtual_column.udf_name") == b"load_video"
    assert udf_field.metadata.get(b"virtual_column.udf_backend") is not None
    assert udf_field.metadata.get(b"virtual_column.udf") is not None

    dl_view = db.open_table("dl_view")
    cnt_not_null = dl_view.count_rows(filter="video is not null")
    assert cnt_not_null == 0

    dl_view.refresh()

    cnt_null = dl_view.count_rows(filter="video is null")
    _LOG.info(dl_view.to_arrow())
    assert cnt_null == 0


def test_create_materialized_view_of_view_ints(db) -> None:
    tbl = pa.Table.from_pydict({"video_uri": [0, 1, 2, 3, 4, 5]})
    video_table = db.create_table(TABLE, tbl)

    @udf
    def load_video(video_uri: int) -> int:  # avoiding binary for now
        return video_uri * 10

    _LOG.info(f"original video_table: {video_table.to_arrow()}")

    dl_view = (
        video_table.search(None)
        .shuffle(seed=42)
        .select(
            {
                "video_uri": "video_uri",
                "video": load_video,
            }
        )
        .create_materialized_view(db, "dl_view")
    )

    udf_field = dl_view.schema.field("video")
    assert udf_field.metadata.get(b"virtual_column") == b"true"
    assert udf_field.metadata.get(b"virtual_column.udf_inputs") is not None
    assert udf_field.metadata.get(b"virtual_column.udf_name") == b"load_video"
    assert udf_field.metadata.get(b"virtual_column.udf_backend") is not None
    assert udf_field.metadata.get(b"virtual_column.udf") is not None

    @udf
    def caption_video(video: int) -> int:
        return video * 10

    q = (
        dl_view.search(None)
        .shuffle(seed=42)
        .select(
            {
                "video_uri": "video_uri",
                "video": "video",
                "caption": caption_video,
            }
        )
    )

    _LOG.info(f"query: table: {q._table} udf_cols: {q._column_udfs}")

    caption_view = q.create_materialized_view(db, "caption_view")

    # caption should be a UDF
    udf_field = caption_view.schema.field("caption")
    assert udf_field.metadata.get(b"virtual_column") == b"true"
    assert udf_field.metadata.get(b"virtual_column.udf_inputs") is not None
    assert udf_field.metadata.get(b"virtual_column.udf_name") == b"caption_video"
    assert udf_field.metadata.get(b"virtual_column.udf_backend") is not None
    assert udf_field.metadata.get(b"virtual_column.udf") is not None

    # video in matview should a copy of the values from the source UDF col, and not
    # the UDF
    udf_field = caption_view.schema.field("video")
    assert udf_field.metadata is None

    # Nothing has been refresh so no data in matview
    _LOG.info(f"dl_view before refresh: {dl_view.to_arrow()}")
    cnt_null = caption_view.count_rows(filter="video is not null")
    assert cnt_null == 0

    dl_view.refresh()

    # refreshed source table but not refreshed to matview yet
    _LOG.info(f"dl_view after refresh: {dl_view.to_arrow()}")
    _LOG.info(f"caption mv before refresh: {caption_view.to_arrow()}")
    cnt_null = caption_view.count_rows(filter="video is not null")
    assert cnt_null == 0

    caption_view.refresh()

    # Now all values should be in matview
    _LOG.info(f"caption mv after refresh: {caption_view.to_arrow()}")
    cnt_null = caption_view.count_rows(filter="video is null")
    assert cnt_null == 0


def test_create_materialized_view_of_view(db, video_table) -> None:
    @udf
    def load_video(video_uri: str) -> bytes:
        return str(video_uri).encode("utf-8")

    _LOG.info(f"original video_table: {video_table.to_arrow()}")

    dl_view = (
        video_table.search(None)
        .shuffle(seed=42)
        .select(
            {
                "video_uri": "video_uri",
                "video": load_video,
            }
        )
        .create_materialized_view(db, "dl_view")
    )

    udf_field = dl_view.schema.field("video")
    assert udf_field.metadata.get(b"virtual_column") == b"true"
    assert udf_field.metadata.get(b"virtual_column.udf_inputs") is not None
    assert udf_field.metadata.get(b"virtual_column.udf_name") == b"load_video"
    assert udf_field.metadata.get(b"virtual_column.udf_backend") is not None
    assert udf_field.metadata.get(b"virtual_column.udf") is not None

    @udf
    def caption_video(video: bytes) -> str:
        return f"this is video {video.decode('utf-8')}"

    q = (
        dl_view.search(None)
        .shuffle(seed=42)
        .select(
            {
                "video_uri": "video_uri",
                "video": "video",
                "caption": caption_video,
            }
        )
    )

    _LOG.info(f"query: table: {q._table} udf_cols: {q._column_udfs}")

    caption_view = q.create_materialized_view(db, "caption_view")

    # caption should be a UDF
    udf_field = caption_view.schema.field("caption")
    assert udf_field.metadata.get(b"virtual_column") == b"true"
    assert udf_field.metadata.get(b"virtual_column.udf_inputs") is not None
    assert udf_field.metadata.get(b"virtual_column.udf_name") == b"caption_video"
    assert udf_field.metadata.get(b"virtual_column.udf_backend") is not None
    assert udf_field.metadata.get(b"virtual_column.udf") is not None

    # video in matview should a copy of the values from the source UDF col, and not
    # the UDF
    udf_field = caption_view.schema.field("video")
    assert udf_field.metadata is None

    # Nothing has been refresh so no data in matview
    _LOG.info(f"dl_view before refresh: {dl_view.to_arrow()}")
    cnt_null = caption_view.count_rows(filter="video is not null")
    assert cnt_null == 0

    dl_view.refresh()

    # refreshed source table but not refreshed to matview yet
    _LOG.info(f"dl_view after refresh: {dl_view.to_arrow()}")
    _LOG.info(f"caption mv before refresh: {caption_view.to_arrow()}")
    cnt_null = caption_view.count_rows(filter="video is not null")
    assert cnt_null == 0

    caption_view.refresh()

    # Now all values should be in matview
    _LOG.info(f"caption mv after refresh: {caption_view.to_arrow()}")
    cnt_null = caption_view.count_rows(filter="video is null")
    assert cnt_null == 0


def test_ray_materialized_view(db, video_table) -> None:
    @udf(data_type=pa.binary())
    def load_video(video_uri: pa.Array) -> pa.Array:
        videos = [str(i).encode("utf-8") for i in video_uri]
        return cast("pa.Array", pa.array(videos))

    view_table = (
        video_table.search(None)
        .shuffle(seed=42)
        .select(
            {
                "video_uri": "video_uri",
                "video": load_video,
            }
        )
        .create_materialized_view(db, "table_view")
    )

    run_ray_copy_table(view_table.get_reference(), db._packager, db._checkpoint_store)

    view_table.checkout_latest()
    assert view_table.to_arrow() == pa.Table.from_pydict(
        {
            "__source_row_id": [3, 2, 5, 4, 1, 0],
            "__is_set": [False] * 6,
            "video_uri": ["d", "c", "f", "e", "b", "a"],
            "video": [b"d", b"c", b"f", b"e", b"b", b"a"],
        }
    )


def test_materialized_view_refresh_with_updated_source(db) -> None:
    """Test that incremental refresh correctly picks up new rows.

    This test verifies that new rows added to the source table are correctly
    picked up by incremental refresh.

    This test verifies that the append-only strategy for new fragments works correctly:
    1. First refresh processes initial rows
    2. New data is added to source (creating new fragments)
    3. Second refresh incrementally adds only the new rows using append semantics
    """
    # Create initial table with some data
    tbl = pa.Table.from_pydict({"video_uri": [0, 1, 2]})
    video_table = db.create_table("source_table", tbl)
    initial_version = video_table.version

    @udf
    def double_value(video_uri: int) -> int:
        return video_uri * 2

    # Create materialized view
    mv = (
        video_table.search(None)
        .select({"video_uri": "video_uri", "doubled": double_value})
        .create_materialized_view(db, "mv_test")
    )

    # First refresh - should process initial 3 rows
    mv.refresh()
    assert mv.count_rows() == 3
    result = mv.to_arrow().to_pydict()
    assert sorted(result["video_uri"]) == [0, 1, 2]
    assert sorted(result["doubled"]) == [0, 2, 4]

    # Update source table with more rows
    video_table.add(pa.Table.from_pydict({"video_uri": [3, 4, 5]}))
    updated_version = video_table.version
    assert updated_version > initial_version
    assert video_table.count_rows() == 6

    # Log fragment information
    src_lance_ds = video_table.to_lance()
    fragments = list(src_lance_ds.get_fragments())
    _LOG.info(f"Source table after add: {len(fragments)} fragments")
    for frag in fragments:
        _LOG.info(f"  Fragment {frag.fragment_id}: {frag.count_rows()} rows")

    # Second refresh - should incrementally add the 3 new rows
    _LOG.info(
        f"Source table before second refresh: "
        f"count={video_table.count_rows()}, version={video_table.version}"
    )
    _LOG.info(
        f"MV before second refresh: count={mv.count_rows()}, version={mv.version}"
    )
    mv.refresh()
    _LOG.info(f"MV after second refresh: count={mv.count_rows()}, version={mv.version}")

    # Verify all 6 rows are present after incremental refresh
    _LOG.info(f"MV data after incremental refresh: {mv.to_arrow().to_pydict()}")
    assert mv.count_rows() == 6
    result = mv.to_arrow().to_pydict()
    assert sorted(result["video_uri"]) == [0, 1, 2, 3, 4, 5]
    assert sorted(result["doubled"]) == [0, 2, 4, 6, 8, 10]

    _LOG.info(
        "Test passed: materialized view successfully refreshed with "
        "incremental new data"
    )


def test_materialized_view_refresh_with_specific_version(db) -> None:
    """Test that refresh can incrementally target specific source versions.

    This test verifies that incremental refresh works correctly when refreshing
    to specific source table versions.
    """
    # Create initial table with some data
    tbl = pa.Table.from_pydict({"value": [1, 2, 3]})
    source_table = db.create_table("versioned_source", tbl)
    version_1 = source_table.version

    @udf
    def triple_value(value: int) -> int:
        return value * 3

    # Create materialized view
    mv = (
        source_table.search(None)
        .select({"value": "value", "tripled": triple_value})
        .create_materialized_view(db, "mv_versioned")
    )

    # Add more rows to create version 2
    source_table.add(pa.Table.from_pydict({"value": [4, 5]}))
    version_2 = source_table.version
    assert version_2 > version_1

    # Add even more rows to create version 3
    source_table.add(pa.Table.from_pydict({"value": [6, 7]}))
    version_3 = source_table.version
    assert version_3 > version_2

    # Refresh to specific version (version 2, which has 5 rows)
    mv.refresh(src_version=version_2)
    assert mv.count_rows() == 5
    result = mv.to_arrow().to_pydict()
    _LOG.info(f"MV data after refresh to version_2: {result}")
    assert sorted(result["value"]) == [1, 2, 3, 4, 5]
    assert sorted(result["tripled"]) == [3, 6, 9, 12, 15]

    # Refresh to latest (should get all 7 rows)
    mv.refresh()  # defaults to latest
    assert mv.count_rows() == 7
    result = mv.to_arrow().to_pydict()
    assert sorted(result["value"]) == [1, 2, 3, 4, 5, 6, 7]
    assert sorted(result["tripled"]) == [3, 6, 9, 12, 15, 18, 21]

    _LOG.info("Test passed: materialized view refresh with specific version works")
