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

"""
Integration tests for Phalanx remote server connection.

These tests assume a local Phalanx server is running on localhost.
The server can be started with:
    cd src/rust/phalanx && cargo run -- --config configs/local.toml

Tests verify:
1. Client can connect to Phalanx via db:// URI with host_override
2. Client can access system tables (jobs, clusters, manifests) via remote connection
3. Backfill jobs work correctly through Phalanx
"""

import logging
import uuid
from collections.abc import Generator
from pathlib import Path
from typing import Any

import pyarrow as pa
import pytest

import geneva
from geneva import connect
from geneva.cluster import GenevaCluster
from geneva.db import Connection

_LOG = logging.getLogger(__name__)


@pytest.fixture
def phalanx_conn(session_db_uri: str, host_override: str, api_key: str) -> Connection:
    db = connect(
        uri=session_db_uri,
        api_key=api_key,
        host_override=host_override,
    )
    return db


@pytest.fixture
def local_cluster(phalanx_conn: Connection, slug) -> Generator[str, Any, None]:
    cluster_name = f"local-{slug}"
    gc = GenevaCluster.create_local(cluster_name).build()
    phalanx_conn.define_cluster(cluster_name, gc)
    yield cluster_name
    phalanx_conn.delete_cluster(cluster_name)


@pytest.mark.phalanx
class TestPhalanxConnection:
    """Tests for connecting to Phalanx server via db:// URI."""

    def test_connect_to_phalanx(
        self,
        phalanx_conn: Connection,
        session_db_uri: str,
        host_override: str,
    ) -> None:
        """Test basic connection to Phalanx server."""

        assert phalanx_conn.uri == session_db_uri
        assert phalanx_conn._host_override == host_override

    def test_phalanx_table_names(self, phalanx_conn: Connection) -> None:
        """Test listing table names via Phalanx."""
        # Should be able to list tables (may be empty initially)
        tables = phalanx_conn.table_names()
        assert isinstance(tables, list)
        _LOG.info(f"Tables via Phalanx: {tables}")

    def test_phalanx_create_and_read_table(
        self, tmp_path: Path, phalanx_conn: Connection
    ) -> None:
        """Test creating and reading a table via Phalanx."""
        # Create a test table
        table_name = f"phalanx_test_{uuid.uuid4().hex[:8]}"
        data = pa.Table.from_pydict({"id": [1, 2, 3], "value": ["a", "b", "c"]})
        db = phalanx_conn

        try:
            db.create_table(table_name, data)

            # Verify table exists
            assert table_name in db.table_names()

            # Read back data
            tbl = db.open_table(table_name)
            assert tbl.count_rows() == 3
        finally:
            # Cleanup
            db.drop_table(table_name)

    def test_phalanx_with_upload_bucket(
        self, tmp_path: Path, session_db_uri: str, host_override: str, api_key: str
    ) -> None:
        """Test Phalanx connection with separate upload bucket."""
        upload_bucket = str(tmp_path / "uploads")

        db = connect(
            uri=session_db_uri,
            api_key=api_key,
            host_override=host_override,
            upload_dir=upload_bucket,
        )

        assert db._upload_dir == upload_bucket


@pytest.mark.phalanx
class TestPhalanxSystemTables:
    """Tests for system tables (jobs, clusters, manifests) via Phalanx."""

    def test_list_manifests_via_phalanx(self, phalanx_conn: Connection) -> None:
        """Test listing jobs via Phalanx connection."""
        # Should be able to list jobs (may be empty)
        jobs = phalanx_conn.list_manifests()
        assert isinstance(jobs, list)

    def test_list_clusters_via_phalanx(self, phalanx_conn: Connection) -> None:
        """Test listing clusters via Phalanx connection."""
        # Should be able to list clusters (may be empty)
        clusters = phalanx_conn.list_clusters()
        assert isinstance(clusters, list)


@pytest.mark.phalanx
class TestPhalanxBackfill:
    """Integration tests for running backfill jobs through Phalanx."""

    @pytest.fixture
    def test_table(self, phalanx_conn: Connection) -> str:
        """Create a test table for backfill tests."""
        table_name = f"backfill_test_{uuid.uuid4().hex[:8]}"
        data = pa.Table.from_pydict(
            {
                "id": list(range(100)),
                "text": [f"item_{i}" for i in range(100)],
                "vec": list(range(100)),
            }
        )

        phalanx_conn.create_table(table_name, data)
        yield table_name

        # Cleanup
        phalanx_conn.drop_table(table_name)

    def test_backfill_via_phalanx_local_ray(
        self, test_table: str, phalanx_conn: Connection, local_cluster: str
    ) -> None:
        """Test backfill job execution via Phalanx with local Ray cluster.

        This test:
        1. Connects to Phalanx server
        2. Creates a backfill job on a test table
        3. Runs the job with local Ray
        4. Verifies results
        """

        self._backfill_test(local_cluster, phalanx_conn, test_table, None)

    @pytest.mark.timeout(900)
    def test_backfill_via_phalanx_kuberay(
        self,
        test_table: str,
        phalanx_conn: Connection,
        session_cluster: str,
        session_manifest: str,
    ) -> None:
        self._backfill_test(session_cluster, phalanx_conn, test_table, session_manifest)

    def _backfill_test(
        self, cluster_name: str, phalanx_conn, test_table, session_manifest
    ) -> None:
        @geneva.udf(data_type=pa.string(), version=uuid.uuid4().hex)
        def uppercase_text(text: str) -> str:
            return text.upper() if text else text

        tbl = phalanx_conn.open_table(test_table)
        # note create_materialized_view and add_columns with UDF not supported
        # by remote table
        with phalanx_conn.context(cluster=cluster_name, manifest=session_manifest):
            tbl.add_columns(
                {"upper": uppercase_text},  # type: ignore[arg-type]
                batch_size=10,
                concurrency=4,
                intra_applier_concurrency=4,
            )
            tbl.checkout_latest()
            # Disable admission control - can't query cluster resources through Phalanx
            tbl.backfill("upper", _admission_check=False)
        # Data assertions: verify the backfill created the 'upper' column correctly
        tbl = phalanx_conn.open_table(test_table)
        tbl.checkout_latest()
        # Check if column exists and has data
        result = tbl.search().select(["id", "text", "upper"]).to_arrow()
        # Verify the 'upper' column exists
        assert "upper" in result.column_names, (
            f"Expected 'upper' column, got {result.column_names}"
        )
        # Verify all rows were processed
        assert result.num_rows == 100, f"Expected 100 rows, got {result.num_rows}"
        # Verify the transformation is correct (text should be uppercased)
        text_col = result.column("text").to_pylist()
        upper_col = result.column("upper").to_pylist()
        for i, (text_val, upper_val) in enumerate(
            zip(text_col, upper_col, strict=True)
        ):
            expected = text_val.upper() if text_val else text_val
            assert upper_val == expected, (
                f"Row {i}: expected '{expected}', got '{upper_val}'"
            )
        # Spot check a few values (need to find them by id since order may vary)
        id_col = result.column("id").to_pylist()
        data = {id_col[i]: upper_col[i] for i in range(len(id_col))}
        assert data[0] == "ITEM_0", f"Expected 'ITEM_0' for id=0, got '{data[0]}'"
        assert data[50] == "ITEM_50", f"Expected 'ITEM_50' for id=50, got '{data[50]}'"
        assert data[99] == "ITEM_99", f"Expected 'ITEM_99' for id=99, got '{data[99]}'"
        _LOG.info(f"Backfill via Phalanx completed for table {test_table}")

    def test_matview_via_phalanx_local_ray(
        self, test_table: str, phalanx_conn: Connection, local_cluster: str
    ) -> None:
        tbl = phalanx_conn.open_table(test_table)

        # note create_materialized_view and add_columns with UDF not supported
        # by remote table
        with phalanx_conn.context(cluster=local_cluster):
            mv = tbl.search().create_materialized_view(
                phalanx_conn, f"test-matview-{tbl.name}"
            )
            mv.refresh()

        assert mv.count_rows() == 100

        _LOG.info(f"Refresh via Phalanx completed for table {test_table}")
