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

from collections.abc import Iterator
from typing import NamedTuple

import pyarrow as pa
import pytest

import geneva
from geneva import connect
from geneva.db import Connection
from geneva.packager import (
    ScalarUDTFSpec,
    marshal_scalar_udtf,
    unmarshal_scalar_udtf,
)
from geneva.table import Table
from geneva.transformer import ScalarUDTF

pytestmark = pytest.mark.ray


# ---------------------------------------------------------------------------
# Shared types and helpers
# ---------------------------------------------------------------------------


class Clip(NamedTuple):
    clip_start: float
    clip_end: float


class Chunk(NamedTuple):
    chunk_text: str
    chunk_index: int


CLIP_SCHEMA = pa.schema(
    [
        pa.field("clip_start", pa.float64()),
        pa.field("clip_end", pa.float64()),
    ]
)


def _make_video_table(tmp_path, name: str = "videos") -> tuple[Connection, Table]:
    db = connect(tmp_path)
    data = pa.table(
        {
            "video_path": pa.array(["/v/a.mp4", "/v/b.mp4", "/v/c.mp4"]),
            "duration": pa.array([30.0, 20.0, 10.0]),
        }
    )
    tbl = db.create_table(name, data)
    return db, tbl


def _make_doc_table(tmp_path, name: str = "docs") -> tuple[Connection, Table]:
    db = connect(tmp_path)
    data = pa.table(
        {
            "doc_id": pa.array([1, 2]),
            "text": pa.array(
                [
                    "Hello world. How are you. Fine thanks.",
                    "One sentence only.",
                ]
            ),
        }
    )
    tbl = db.create_table(name, data)
    return db, tbl


# ---------------------------------------------------------------------------
# Unit: ScalarUDTF decorator and class
# ---------------------------------------------------------------------------


class TestScalarUDTFDecorator:
    def test_basic_decorator_infers_schema(self) -> None:
        @geneva.scalar_udtf
        def extract_clips(
            duration: float,
        ) -> Iterator[Clip]:
            for start in range(0, int(duration), 10):
                end = min(start + 10, duration)
                yield Clip(clip_start=start, clip_end=end)

        assert isinstance(extract_clips, ScalarUDTF)
        assert extract_clips.output_schema == CLIP_SCHEMA
        assert extract_clips.input_columns == ["duration"]
        assert extract_clips.name == "extract_clips"
        assert extract_clips.batch is False

    def test_decorator_with_explicit_schema(self) -> None:
        @geneva.scalar_udtf(output_schema=CLIP_SCHEMA)
        def extract_clips(
            duration: float,
        ) -> Iterator[Clip]:
            for start in range(0, int(duration), 10):
                yield Clip(clip_start=start, clip_end=min(start + 10, duration))

        assert extract_clips.output_schema == CLIP_SCHEMA

    def test_decorator_with_input_columns(self) -> None:
        @geneva.scalar_udtf(input_columns=["dur"])
        def extract_clips(
            dur: float,
        ) -> Iterator[Clip]:
            yield Clip(clip_start=0, clip_end=dur)

        assert extract_clips.input_columns == ["dur"]

    def test_decorator_batch_mode(self) -> None:
        @geneva.scalar_udtf(batch=True, output_schema=CLIP_SCHEMA)
        def batch_extract(
            duration: pa.Array,
            __source_row_id: pa.Array,
        ) -> pa.RecordBatch:
            return pa.RecordBatch.from_pydict(
                {
                    "__source_row_id": __source_row_id,
                    "clip_start": pa.array([0.0]),
                    "clip_end": duration,
                }
            )

        assert isinstance(batch_extract, ScalarUDTF)
        assert batch_extract.batch is True


# ---------------------------------------------------------------------------
# Unit: execute_on_record_batch (scalar mode)
# ---------------------------------------------------------------------------


class TestScalarExecution:
    def test_1_to_n_expansion(self) -> None:
        @geneva.scalar_udtf
        def extract_clips(
            duration: float,
        ) -> Iterator[Clip]:
            for start in range(0, int(duration), 10):
                end = min(start + 10, duration)
                yield Clip(clip_start=start, clip_end=end)

        batch = pa.RecordBatch.from_pydict(
            {
                "__source_row_id": pa.array([100, 200], type=pa.int64()),
                "duration": pa.array([25.0, 10.0]),
            }
        )

        result = extract_clips.execute_on_record_batch(batch)

        # Row 100 (duration=25): 3 clips [0-10, 10-20, 20-25]
        # Row 200 (duration=10): 1 clip [0-10]
        assert result.num_rows == 4

        src_ids = result["__source_row_id"].to_pylist()
        assert src_ids == [100, 100, 100, 200]

        child_indices = result["__child_index"].to_pylist()
        assert child_indices == [0, 1, 2, 0]

        starts = result["clip_start"].to_pylist()
        assert starts == [0.0, 10.0, 20.0, 0.0]

        ends = result["clip_end"].to_pylist()
        assert ends == [10.0, 20.0, 25.0, 10.0]

    def test_zero_yield_rows(self) -> None:
        """A source row can produce zero output rows."""

        @geneva.scalar_udtf
        def maybe_expand(
            value: int,
        ) -> Iterator[Clip]:
            if value > 10:
                yield Clip(clip_start=0, clip_end=float(value))

        batch = pa.RecordBatch.from_pydict(
            {
                "__source_row_id": pa.array([1, 2, 3], type=pa.int64()),
                "value": pa.array([5, 20, 3]),
            }
        )

        result = maybe_expand.execute_on_record_batch(batch)
        # Only row 2 (value=20) yields output
        assert result.num_rows == 1
        assert result["__source_row_id"].to_pylist() == [2]

    def test_dict_yield(self) -> None:
        @geneva.scalar_udtf(output_schema=CLIP_SCHEMA)
        def extract_clips(
            duration: float,
        ) -> Iterator[dict]:
            yield {"clip_start": 0.0, "clip_end": duration}

        batch = pa.RecordBatch.from_pydict(
            {
                "__source_row_id": pa.array([1], type=pa.int64()),
                "duration": pa.array([10.0]),
            }
        )

        result = extract_clips.execute_on_record_batch(batch)
        assert result.num_rows == 1
        assert result["clip_end"].to_pylist() == [10.0]


class TestSerialization:
    def test_marshal_unmarshal_round_trip(self) -> None:
        @geneva.scalar_udtf
        def extract_clips(
            duration: float,
        ) -> Iterator[Clip]:
            for start in range(0, int(duration), 10):
                yield Clip(
                    clip_start=start,
                    clip_end=min(start + 10, duration),
                )

        spec = marshal_scalar_udtf(extract_clips)
        assert isinstance(spec, ScalarUDTFSpec)
        assert spec.name == "extract_clips"

        # Round-trip via JSON
        json_str = spec.to_json()
        spec2 = ScalarUDTFSpec.from_json(json_str)
        assert spec2.name == spec.name
        assert spec2.version == spec.version

        # Unmarshal
        restored = unmarshal_scalar_udtf(spec2)
        assert isinstance(restored, ScalarUDTF)
        assert restored.name == "extract_clips"
        assert restored.output_schema == CLIP_SCHEMA


class TestCreateAndRefresh:
    def test_create_view_empty(self, tmp_path) -> None:
        db, videos = _make_video_table(tmp_path)

        @geneva.scalar_udtf
        def extract_clips(
            duration: float,
        ) -> Iterator[Clip]:
            for start in range(0, int(duration), 10):
                yield Clip(
                    clip_start=start,
                    clip_end=min(start + 10, duration),
                )

        query = videos.search(None).select(["video_path", "duration"])
        view = db.create_scalar_udtf_view("clips", query, extract_clips)

        assert view.count_rows() == 0
        schema = view.schema
        assert "__source_row_id" in schema.names
        assert "__child_index" in schema.names
        assert "video_path" in schema.names
        assert "duration" in schema.names
        assert "clip_start" in schema.names
        assert "clip_end" in schema.names

    def test_refresh_expands_rows(self, tmp_path, ray_with_test_path) -> None:
        db, videos = _make_video_table(tmp_path)

        @geneva.scalar_udtf
        def extract_clips(
            duration: float,
        ) -> Iterator[Clip]:
            for start in range(0, int(duration), 10):
                yield Clip(
                    clip_start=start,
                    clip_end=min(start + 10, duration),
                )

        query = videos.search(None).select(["video_path", "duration"])
        view = db.create_scalar_udtf_view("clips", query, extract_clips)

        view.refresh(_admission_check=False)

        # video a: duration=30 -> 3 clips
        # video b: duration=20 -> 2 clips
        # video c: duration=10 -> 1 clip
        # Total: 6 clips
        assert view.count_rows() == 6

        result = view.to_pandas()
        # Check inherited columns are present
        assert "video_path" in result.columns
        assert "duration" in result.columns
        # Check UDTF output columns
        assert "clip_start" in result.columns
        assert "clip_end" in result.columns
        # Check internal columns
        assert "__source_row_id" in result.columns
        assert "__child_index" in result.columns

    def test_inherited_columns_correct(self, tmp_path, ray_with_test_path) -> None:
        db, videos = _make_video_table(tmp_path)

        @geneva.scalar_udtf
        def extract_clips(
            duration: float,
        ) -> Iterator[Clip]:
            for start in range(0, int(duration), 10):
                yield Clip(
                    clip_start=start,
                    clip_end=min(start + 10, duration),
                )

        query = videos.search(None).select(["video_path", "duration"])
        view = db.create_scalar_udtf_view("clips", query, extract_clips)
        view.refresh(_admission_check=False)

        result = view.to_pandas()
        # All clips from video a should inherit video_path="/v/a.mp4"
        a_clips = result[result["video_path"] == "/v/a.mp4"]
        assert len(a_clips) == 3
        assert all(a_clips["duration"] == 30.0)

        # All clips from video c should have video_path="/v/c.mp4"
        c_clips = result[result["video_path"] == "/v/c.mp4"]
        assert len(c_clips) == 1
        assert c_clips.iloc[0]["clip_start"] == 0.0
        assert c_clips.iloc[0]["clip_end"] == 10.0

    def test_zero_yield_source_row(self, tmp_path, ray_with_test_path) -> None:
        """Source rows that yield zero output are valid."""
        db, docs = _make_doc_table(tmp_path)

        @geneva.scalar_udtf
        def split_long_text(
            text: str,
        ) -> Iterator[Chunk]:
            sentences = text.split(". ")
            if len(sentences) > 1:
                for i, s in enumerate(sentences):
                    yield Chunk(chunk_text=s, chunk_index=i)
            # Single-sentence docs yield nothing

        query = docs.search(None).select(["doc_id", "text"])
        view = db.create_scalar_udtf_view("chunks", query, split_long_text)
        view.refresh(_admission_check=False)

        # Doc 1 has 3 sentences -> 3 chunks
        # Doc 2 has 1 sentence -> 0 chunks
        assert view.count_rows() == 3

    def test_empty_source_table(self, tmp_path, ray_with_test_path) -> None:
        """Refresh with no source rows produces empty view."""
        db = connect(tmp_path)
        empty_data = pa.table(
            {
                "duration": pa.array([], type=pa.float64()),
            }
        )
        src = db.create_table("empty_src", empty_data)

        @geneva.scalar_udtf
        def extract_clips(
            duration: float,
        ) -> Iterator[Clip]:
            yield Clip(clip_start=0, clip_end=duration)

        query = src.search(None)
        view = db.create_scalar_udtf_view("clips", query, extract_clips)
        view.refresh(_admission_check=False)
        assert view.count_rows() == 0

    def test_double_refresh(self, tmp_path, ray_with_test_path) -> None:
        """Refreshing twice overwrites with same results."""
        db, videos = _make_video_table(tmp_path)

        @geneva.scalar_udtf
        def extract_clips(
            duration: float,
        ) -> Iterator[Clip]:
            for start in range(0, int(duration), 10):
                yield Clip(
                    clip_start=start,
                    clip_end=min(start + 10, duration),
                )

        query = videos.search(None).select(["video_path", "duration"])
        view = db.create_scalar_udtf_view("clips", query, extract_clips)

        view.refresh(_admission_check=False)
        assert view.count_rows() == 6

        # Second refresh should produce the same result (no new source rows)
        view.refresh(_admission_check=False)
        assert view.count_rows() == 6

    def test_incremental_refresh(self, tmp_path, ray_with_test_path) -> None:
        """Adding source rows and refreshing appends only the new expansions."""
        db, videos = _make_video_table(tmp_path)

        @geneva.scalar_udtf
        def extract_clips(
            duration: float,
        ) -> Iterator[Clip]:
            for start in range(0, int(duration), 10):
                yield Clip(
                    clip_start=start,
                    clip_end=min(start + 10, duration),
                )

        query = videos.search(None).select(["video_path", "duration"])
        view = db.create_scalar_udtf_view("clips", query, extract_clips)

        view.refresh(_admission_check=False)
        initial_rows = view.count_rows()
        assert initial_rows == 6

        # Add a new source row (duration=5 → 1 clip)
        videos.add(
            pa.table(
                {
                    "video_path": ["new.mp4"],
                    "duration": [5.0],
                }
            )
        )

        # Refresh should pick up only the new row
        view.refresh(_admission_check=False)
        assert view.count_rows() == initial_rows + 1

    def test_input_columns_validation_against_projection(self, tmp_path) -> None:
        """input_columns should be validated against query projection."""
        db, videos = _make_video_table(tmp_path)

        @geneva.scalar_udtf(input_columns=["duration"])
        def extract_clips(
            duration: float,
        ) -> Iterator[Clip]:
            yield Clip(clip_start=0, clip_end=duration)

        # Select only video_path — duration is NOT in the projection
        query = videos.search(None).select(["video_path"])

        with pytest.raises(ValueError, match="not found in query projection"):
            db.create_scalar_udtf_view("clips", query, extract_clips)


_BATCH_EXTRACT_SCHEMA = pa.schema(
    [
        pa.field("clip_start", pa.float64()),
        pa.field("clip_end", pa.float64()),
    ]
)


# Defined at module level to avoid Python name mangling of __source_row_id
@geneva.scalar_udtf(batch=True, output_schema=_BATCH_EXTRACT_SCHEMA)
def _batch_extract(
    duration: pa.Array,
    __source_row_id: pa.Array,
) -> pa.RecordBatch:
    starts = []
    ends = []
    src_ids = []
    for i in range(len(duration)):
        dur = duration[i].as_py()
        sid = __source_row_id[i].as_py()
        for s in range(0, int(dur), 10):
            starts.append(float(s))
            ends.append(min(float(s + 10), dur))
            src_ids.append(sid)
    return pa.RecordBatch.from_pydict(
        {
            "__source_row_id": pa.array(src_ids, type=pa.int64()),
            "clip_start": pa.array(starts),
            "clip_end": pa.array(ends),
        }
    )


@geneva.scalar_udtf(batch=True, output_schema=CLIP_SCHEMA)
def _batch_fn(
    duration: pa.Array,
    __source_row_id: pa.Array,
) -> pa.RecordBatch:
    return pa.RecordBatch.from_pydict(
        {
            "__source_row_id": __source_row_id,
            "clip_start": pa.array([0.0]),
            "clip_end": pa.array([1.0]),
        }
    )


class TestBatchExecution:
    def test_batch_mode_execution(self) -> None:
        """batch=True variant processes Arrow Arrays."""
        batch = pa.RecordBatch.from_pydict(
            {
                "__source_row_id": pa.array([100, 200], type=pa.int64()),
                "duration": pa.array([25.0, 10.0]),
            }
        )

        result = _batch_extract.execute_on_record_batch(batch)
        assert result.num_rows == 4
        assert result["__source_row_id"].to_pylist() == [
            100,
            100,
            100,
            200,
        ]
        assert "__child_index" in result.schema.names

    def test_batch_mode_infers_input_columns(self) -> None:
        """batch=True infers input_columns, excluding __source_row_id."""
        assert _batch_fn.input_columns == ["duration"]


class TestErrorHandling:
    def test_skip_on_error(self) -> None:
        """on_error with Skip skips failing rows."""
        from geneva.debug.error_store import Skip

        @geneva.scalar_udtf(on_error=[Skip(ValueError)])
        def flaky_expand(
            value: int,
        ) -> Iterator[Clip]:
            if value == 2:
                raise ValueError("bad row")
            yield Clip(clip_start=0, clip_end=float(value))

        batch = pa.RecordBatch.from_pydict(
            {
                "__source_row_id": pa.array([1, 2, 3], type=pa.int64()),
                "value": pa.array([10, 2, 30]),
            }
        )

        result = flaky_expand.execute_on_record_batch(batch)
        # Row 2 (value=2) raises ValueError -> skipped
        assert result.num_rows == 2
        assert result["__source_row_id"].to_pylist() == [1, 3]

    def test_fail_on_error_default(self) -> None:
        """Without on_error, exceptions propagate."""

        @geneva.scalar_udtf
        def always_fail(
            value: int,
        ) -> Iterator[Clip]:
            raise ValueError("boom")
            yield  # noqa: RET503

        batch = pa.RecordBatch.from_pydict(
            {
                "__source_row_id": pa.array([1], type=pa.int64()),
                "value": pa.array([10]),
            }
        )

        with pytest.raises(ValueError, match="boom"):
            always_fail.execute_on_record_batch(batch)
