# SPDX-License-Identifier: PROPRIETARY
# SPDX-FileCopyrightText: Copyright The Geneva Authors
from pathlib import Path
from unittest.mock import MagicMock, patch

import lancedb
import pyarrow as pa
import pytest

from geneva import connect
from geneva.db import (
    SYSTEM_TABLE_COMPLETED,
    SYSTEM_TABLE_MIGRATING,
    SYSTEM_TABLE_MIGRATIONS_TABLE_NAME,
    Connection,
    _ensure_system_namespace_exists,
)
from geneva.debug.error_store import GENEVA_ERRORS_TABLE_NAME, ErrorStore
from geneva.jobs.jobs import GENEVA_JOBS_TABLE_NAME, JobStateManager
from geneva.manifest.mgr import ManifestConfigManager
from geneva.table import TableReference


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

    # Use lancedb to verify the results are the same
    ldb = lancedb.connect(tmp_path)

    # Create a Table with integer columns
    tbl = pa.Table.from_pydict({"a": [1, 2, 3], "b": [4, 5, 6]})
    db.create_table("table1", tbl)
    ldb_tbls = db.table_names()
    assert "table1" in ldb_tbls
    db.open_table("table1")

    db_tbls = db.table_names()
    assert db_tbls == ldb_tbls

    # Use lancedb to read the data back
    ldb_tbl = ldb.open_table("table1")
    assert ldb_tbl.to_arrow() == tbl
    db.drop_table("table1")


def test_connect_non_existent(tmp_path: Path) -> None:
    db = connect(tmp_path)
    with pytest.raises(ValueError, match=r".*was not found"):
        db.open_table("non_existent")


def test_error_store_from_system_table_ref_stays_in_single_system_db(
    tmp_path: Path,
) -> None:
    db = connect(tmp_path)
    db.create_table("table1", pa.Table.from_pydict({"a": [1]}))

    system_ref = (
        db.open_table("table1")
        .get_reference()
        .as_system_table(GENEVA_ERRORS_TABLE_NAME)
    )
    reopened = system_ref.open_db()
    store = ErrorStore(reopened)

    assert store.table.uri.endswith("/__system/geneva_errors.lance")
    assert "/__system/__system/" not in store.table.uri


def test_error_store_nests_under_system_db_named_system(tmp_path: Path) -> None:
    db_path = tmp_path / "__system"
    db = connect(db_path)
    db.create_table("table1", pa.Table.from_pydict({"a": [1]}))

    system_ref = (
        db.open_table("table1")
        .get_reference()
        .as_system_table(GENEVA_ERRORS_TABLE_NAME)
    )
    reopened = system_ref.open_db()
    store = ErrorStore(reopened)

    assert store.table.uri.endswith("/__system/__system/geneva_errors.lance")


def test_legacy_system_table_migration_writes_completed_marker(tmp_path: Path) -> None:
    db = connect(tmp_path)
    legacy_jobs = JobStateManager(db, namespace=[])
    legacy_jobs.launch("tbl", "col")

    migrated_jobs = JobStateManager(db)

    assert migrated_jobs.get_table().count_rows() == 1
    assert db.open_table(GENEVA_JOBS_TABLE_NAME).count_rows() == 1

    migrations = db.system_table_connection.open_table(
        SYSTEM_TABLE_MIGRATIONS_TABLE_NAME
    )
    rows = (
        migrations.search()
        .where(f"table_name = '{GENEVA_JOBS_TABLE_NAME}'")
        .limit(1)
        .to_arrow()
        .to_pylist()
    )
    assert rows[0]["status"] == SYSTEM_TABLE_COMPLETED


def test_migrate_legacy_system_table_retries_after_failure(tmp_path: Path) -> None:
    db = connect(tmp_path)
    target_db = db.system_table_connection
    copy_calls = 0
    status_updates: list[str] = []

    def fake_copy(**_: object) -> None:
        nonlocal copy_calls
        copy_calls += 1
        if copy_calls == 1:
            raise RuntimeError("copy failed")

    with (
        patch.object(
            Connection, "_get_system_table_migration_status", return_value=None
        ),
        patch.object(
            Connection,
            "_find_legacy_system_table_source",
            return_value=(db, []),
        ),
        patch.object(
            Connection,
            "_set_system_table_migration_status",
            side_effect=lambda table_name, status: status_updates.append(status),
        ),
        patch.object(Connection, "_replace_table_from_source", side_effect=fake_copy),
    ):
        with pytest.raises(RuntimeError, match="copy failed"):
            db.migrate_legacy_system_table_if_needed(
                table_name=GENEVA_JOBS_TABLE_NAME,
                target_conn=target_db,
                target_namespace=[],
            )

        db.migrate_legacy_system_table_if_needed(
            table_name=GENEVA_JOBS_TABLE_NAME,
            target_conn=target_db,
            target_namespace=[],
        )

    assert copy_calls == 2


def test_migrate_legacy_system_table_clears_stale_marker_without_legacy_source(
    tmp_path: Path,
) -> None:
    db = connect(tmp_path)
    target_db = db.system_table_connection

    with (
        patch.object(
            Connection,
            "_get_system_table_migration_status",
            return_value=SYSTEM_TABLE_MIGRATING,
        ),
        patch.object(
            Connection,
            "_find_legacy_system_table_source",
            return_value=None,
        ),
        patch.object(Connection, "_clear_system_table_migration_status") as mock_clear,
        patch.object(Connection, "_replace_table_from_source") as mock_copy,
    ):
        db.migrate_legacy_system_table_if_needed(
            table_name=GENEVA_JOBS_TABLE_NAME,
            target_conn=target_db,
            target_namespace=[],
        )

    mock_clear.assert_called_once_with(GENEVA_JOBS_TABLE_NAME)
    mock_copy.assert_not_called()


class TestConnectRemote:
    """Tests for connect() with remote db:// URI (Phalanx server)."""

    @pytest.fixture(autouse=True)
    def _mock_set_upload_dir(self) -> None:
        with patch("geneva.db._set_upload_dir"):
            yield

    def test_connect_remote_uri_uses_host_override(self) -> None:
        """Test that db:// URI with host_override triggers remote connection."""
        with patch("geneva.db.lancedb.connect") as mock_lancedb_connect:
            mock_conn = MagicMock()
            mock_conn.table_names.return_value = []
            mock_lancedb_connect.return_value = mock_conn

            db = connect(
                uri="db://my_database",
                api_key="test-api-key",
                host_override="https://phalanx.internal:8080",
            )

            # Trigger the lazy connection
            _ = db._connect

            # Verify lancedb.connect was called with correct parameters
            mock_lancedb_connect.assert_called_once()
            call_kwargs = mock_lancedb_connect.call_args.kwargs
            assert call_kwargs.get("api_key") is not None
            assert call_kwargs.get("host_override") == "https://phalanx.internal:8080"

            # Verify the URI was passed correctly
            call_args = mock_lancedb_connect.call_args.args
            assert call_args[0] == "db://my_database"

    def test_connect_remote_stores_connection_params(self) -> None:
        """Test that remote connection parameters are stored on Connection object."""
        db = connect(
            uri="db://test_db",
            api_key="my-api-key",
            host_override="https://phalanx.example.com",
            region="us-west-2",
        )

        assert db.uri == "db://test_db"
        assert db._host_override == "https://phalanx.example.com"
        assert db._region == "us-west-2"
        # api_key is wrapped in Credential
        assert db._api_key is not None

    def test_connect_remote_system_namespace_preserved(self) -> None:
        """Remote connections should preserve configured system_namespace."""
        db = connect(
            uri="db://test_db",
            api_key="test-key",
            host_override="https://phalanx.example.com",
            system_namespace=["prod", "system"],
        )

        assert db.system_namespace == ["prod", "system"]

    def test_connect_remote_system_namespace_defaults_to_system(self) -> None:
        """Remote connections should default internal tables to __system."""
        db = connect(
            uri="db://test_db",
            api_key="test-key",
            host_override="https://phalanx.example.com",
        )

        assert db.system_namespace == ["__system"]

    def test_connect_remote_creates_checkpoint_store(self) -> None:
        """Test that checkpoint store is created for remote connections."""
        db = connect(
            uri="db://test_db",
            api_key="test-key",
            host_override="https://phalanx.example.com",
            upload_dir="test",
        )

        # Checkpoint store should be created
        assert db._checkpoint_store is not None

    def test_connect_local_vs_remote_differentiation(self, tmp_path: Path) -> None:
        """Test that local and remote URIs are handled differently."""
        # Local connection
        local_db = connect(tmp_path)
        assert not local_db.uri.startswith("db://")

        # Remote connection
        remote_db = connect(
            uri="db://remote_db",
            api_key="key",
            host_override="https://phalanx.example.com",
        )
        assert remote_db.uri == "db://remote_db"

    def test_connect_remote_flight_client_creation(self) -> None:
        """Test that flight_client is properly configured for remote connections."""
        db = connect(
            uri="db://test_db",
            api_key="test-key",
            host_override="https://phalanx.example.com:8080",
        )

        # Access flight_client to trigger creation
        with patch("flightsql.FlightSQLClient") as mock_flight:
            mock_flight.return_value = MagicMock()
            _ = db.flight_client

            # Verify FlightSQLClient was created with correct host
            mock_flight.assert_called_once()
            call_kwargs = mock_flight.call_args.kwargs
            assert call_kwargs.get("host") == "phalanx.example.com"


def test_connect_remote_does_not_probe_system_tables_for_upload_dir() -> None:
    """Remote connect should not touch system tables just to infer upload_dir."""
    with (
        patch("geneva.db.Uploader.get", side_effect=ValueError("unset")),
        patch("geneva.db.lancedb.connect") as mock_lancedb_connect,
    ):
        db = connect(
            uri="db://test_db",
            api_key="test-key",
            host_override="https://phalanx.example.com",
        )

    assert db.uri == "db://test_db"
    mock_lancedb_connect.assert_not_called()


def test_connect_namespace_does_not_probe_system_tables_for_upload_dir(
    tmp_path: Path,
) -> None:
    """Namespace connect should not create/open system tables just for upload_dir."""
    with (
        patch("geneva.db.Uploader.get", side_effect=ValueError("unset")),
        patch(
            "geneva.manifest.mgr.ManifestConfigManager",
            side_effect=AssertionError("should not probe manifest table"),
        ),
    ):
        db = connect(
            namespace_impl="dir",
            namespace_properties={"root": str(tmp_path)},
        )

    assert db.namespace_impl == "dir"


def test_ensure_system_namespace_exists_remote_defers_namespace_creation() -> None:
    with (
        patch("geneva.db._set_upload_dir"),
        patch("geneva.db.lancedb.connect") as mock_lancedb_connect,
    ):
        mock_remote = MagicMock()
        mock_lancedb_connect.return_value = mock_remote
        db = connect(
            uri="db://test_db",
            api_key="test-key",
            host_override="https://phalanx.example.com",
        )

        _ensure_system_namespace_exists(db)

    mock_lancedb_connect.assert_not_called()
    assert db._system_namespace_ensured is True


REMOTE_NAMESPACE_LOCATION_ERROR = (
    "Invalid input, Location must be provided when namespace is not empty"
)


def test_remote_system_tables_fallback_to_root_namespace_on_location_error() -> None:
    with (
        patch("geneva.db._set_upload_dir"),
        patch("geneva.utils.schema.alter_or_create_table") as mock_alter_or_create,
    ):
        db = connect(
            uri="db://test_db",
            api_key="test-key",
            host_override="https://phalanx.example.com",
        )
        db.migrate_legacy_system_table_if_needed = MagicMock()
        mock_table = MagicMock()
        mock_alter_or_create.side_effect = [
            RuntimeError(REMOTE_NAMESPACE_LOCATION_ERROR),
            mock_table,
        ]

        mgr = ManifestConfigManager(db)

    assert mgr.table is mock_table
    assert db.system_namespace == ["__system"]
    assert db.effective_system_namespace == []
    assert mock_alter_or_create.call_args_list[0].kwargs["namespace"] == ["__system"]
    assert mock_alter_or_create.call_args_list[1].kwargs["namespace"] == []


def test_remote_system_tables_fallback_when_probe_requires_location() -> None:
    with (
        patch("geneva.db._set_upload_dir"),
        patch("geneva.utils.schema.alter_or_create_table") as mock_alter_or_create,
    ):
        db = connect(
            uri="db://test_db",
            api_key="test-key",
            host_override="https://phalanx.example.com",
        )
        db.migrate_legacy_system_table_if_needed = MagicMock(
            side_effect=[
                RuntimeError(REMOTE_NAMESPACE_LOCATION_ERROR),
                None,
            ]
        )
        mock_table = MagicMock()
        mock_alter_or_create.return_value = mock_table

        mgr = ManifestConfigManager(db)

    assert mgr.table is mock_table
    assert db.system_namespace == ["__system"]
    assert db.effective_system_namespace == []
    assert db.migrate_legacy_system_table_if_needed.call_args_list[0].kwargs[
        "target_namespace"
    ] == ["__system"]
    assert (
        db.migrate_legacy_system_table_if_needed.call_args_list[1].kwargs[
            "target_namespace"
        ]
        == []
    )


def test_table_reference_open_db_preserves_remote_system_namespace() -> None:
    """Reopened remote refs should keep their explicit system namespace."""
    with patch("geneva.db._set_upload_dir"):
        ref = TableReference(
            table_id=["tbl"],
            version=None,
            db_uri="db://test_db",
            system_namespace=["prod", "system"],
        )

        db = ref.open_db()

    assert db.system_namespace == ["prod", "system"]


def test_table_reference_open_db_preserves_remote_credentials() -> None:
    with patch("geneva.table.connect") as mock_connect:
        ref = TableReference(
            table_id=["tbl"],
            version=None,
            db_uri="db://test_db",
            api_key="test-key",
            host_override="https://phalanx.example.com",
        )

        ref.open_db()

    assert mock_connect.call_args.kwargs["api_key"] == "test-key"
    assert mock_connect.call_args.kwargs["host_override"] == (
        "https://phalanx.example.com"
    )
    assert mock_connect.call_args.kwargs["initialize_upload_dir"] is False


def test_table_reference_open_db_does_not_probe_upload_dir(tmp_path: Path) -> None:
    with (
        patch("geneva.db.Uploader.get", side_effect=ValueError("unset")),
        patch(
            "geneva.manifest.mgr.ManifestConfigManager",
            side_effect=AssertionError("should not probe manifest table"),
        ),
    ):
        ref = TableReference(
            table_id=["tbl"],
            version=None,
            db_uri=str(tmp_path),
        )

        ref.open_db()


@pytest.mark.asyncio
async def test_table_reference_open_system_db_async_preserves_remote_credentials() -> (
    None
):
    with patch("geneva.table.connect_async") as mock_connect_async:
        ref = TableReference(
            table_id=["tbl"],
            version=None,
            db_uri="db://test_db",
            api_key="test-key",
            host_override="https://phalanx.example.com",
        )

        await ref.open_system_db_async()

    assert mock_connect_async.call_args.kwargs["api_key"] == "test-key"
    assert (
        mock_connect_async.call_args.kwargs["host_override"]
        == "https://phalanx.example.com"
    )


def test_table_reference_checkpoint_store_uses_remote_table_location() -> None:
    with (
        patch("geneva.table.resolve_table_physical_uri") as mock_resolve_uri,
        patch("geneva.checkpoint.LanceCheckpointStore") as mock_checkpoint_store,
    ):
        mock_resolve_uri.return_value = "s3://bucket/db/tbl.lance"

        ref = TableReference(
            table_id=["tbl"],
            version=None,
            db_uri="db://test_db",
            api_key="test-key",
            host_override="https://phalanx.example.com",
        )

        ref.open_checkpoint_store()

    assert mock_resolve_uri.call_args.kwargs["api_key"] == "test-key"
    assert (
        mock_resolve_uri.call_args.kwargs["host_override"]
        == "https://phalanx.example.com"
    )
    assert mock_checkpoint_store.call_args.args[0] == "s3://bucket/db/tbl.lance/_ckp"


def test_table_reference_open_db_preserves_explicit_empty_remote_system_namespace() -> (
    None
):
    """Explicit root namespace should not be rewritten to __system."""
    with patch("geneva.db._set_upload_dir"):
        ref = TableReference(
            table_id=["tbl"],
            version=None,
            db_uri="db://test_db",
            system_namespace=[],
        )

        db = ref.open_db()

    assert ref.system_table_namespace == []
    assert db.system_namespace == []


def test_colocated_system_table_ref_reuses_physical_db_root() -> None:
    with patch("geneva.table.connect") as mock_connect:
        ref = TableReference(
            table_id=["tbl"],
            version=None,
            db_uri="s3://bucket/db",
            api_key="test-key",
            host_override="http://localhost:10024",
            system_namespace=[],
            co_located_system_tables=True,
        )

        ref.as_system_table(GENEVA_JOBS_TABLE_NAME).open_db()

    assert mock_connect.call_args.args[0] == "s3://bucket/db"
    assert "api_key" not in mock_connect.call_args.kwargs
    assert "host_override" not in mock_connect.call_args.kwargs
    assert mock_connect.call_args.kwargs["is_system_db_root"] is False


class TestConnectConfig:
    """Tests for connect() configuration loading."""

    def test_connect_loads_config_defaults(self, tmp_path: Path) -> None:
        """Test that connect() loads defaults from config."""
        db = connect(tmp_path)
        # Default region should be set
        assert db._region == "us-east-1"

    def test_connect_explicit_params_override_config(self, tmp_path: Path) -> None:
        """Test that explicit parameters override config defaults."""
        db = connect(tmp_path, region="eu-west-1")
        assert db._region == "eu-west-1"


class TestUploadBucket:
    """Tests for upload_bucket configuration (Phalanx security model)."""

    def test_connect_with_upload_bucket(self, tmp_path: Path) -> None:
        """Test that upload_bucket is stored on Connection object."""
        db = connect(
            tmp_path,
            upload_dir="s3://my-upload-bucket/manifests",
        )
        assert db._upload_dir == "s3://my-upload-bucket/manifests"

    def test_connect_without_upload_bucket(self, tmp_path: Path) -> None:
        """Test that upload_bucket defaults to None."""
        db = connect(tmp_path)
        assert db._upload_dir is None

    def test_upload_bucket_with_remote_connection(self) -> None:
        """Test upload_bucket with remote db:// URI (Phalanx scenario)."""
        db = connect(
            uri="db://my_database",
            api_key="test-api-key",
            host_override="https://phalanx.example.com",
            upload_dir="gs://enterprise-upload-bucket/manifests",
        )
        assert db._upload_dir == "gs://enterprise-upload-bucket/manifests"
        assert db.uri == "db://my_database"

    def test_upload_bucket_serialization(self, tmp_path: Path) -> None:
        """Test that upload_bucket is properly serialized/deserialized."""
        db = connect(
            tmp_path,
            upload_dir="s3://my-upload-bucket/manifests",
        )

        # Get serialized state
        state = db.__getstate__()
        assert state["upload_dir"] == "s3://my-upload-bucket/manifests"

        # Create new connection from state
        new_db = connect(tmp_path)
        new_db.__setstate__(state)
        assert new_db._upload_dir == "s3://my-upload-bucket/manifests"

    def test_namespace_connection_serialization(self, tmp_path: Path) -> None:
        """Namespace connection state should survive a pickle-style round trip."""
        db = connect(
            namespace_impl="dir",
            namespace_properties={"root": str(tmp_path)},
            system_namespace=["prod", "system"],
        )

        state = db.__getstate__()
        assert state["namespace_impl"] == "dir"
        assert state["namespace_properties"] == {"root": str(tmp_path)}

        new_db = connect(tmp_path / "other")
        new_db.__setstate__(state)
        assert new_db.namespace_impl == "dir"
        assert new_db.namespace_properties == {"root": str(tmp_path)}
        assert new_db.system_namespace == ["prod", "system"]

    def test_define_manifest_uses_upload_bucket(self, tmp_path: Path) -> None:
        """Test that define_manifest uses upload_bucket when configured."""
        db = connect(
            tmp_path,
            upload_dir="s3://my-upload-bucket/manifests",
        )

        # Mock the Uploader to capture how it's instantiated
        with patch("geneva.db.Uploader") as mock_uploader_class:
            mock_uploader = MagicMock()
            mock_uploader_class.return_value = mock_uploader

            # Mock upload_local_env to avoid actual file operations
            with patch("geneva.db.upload_local_env") as mock_upload:
                mock_upload.return_value.__enter__ = MagicMock(return_value=[])
                mock_upload.return_value.__exit__ = MagicMock(return_value=None)

                # Mock the manifest manager
                with patch.object(db, "_manifest_manager") as mock_mgr:
                    mock_mgr.get_table.return_value = MagicMock()
                    db._manifest_manager = mock_mgr

                    # Create a minimal manifest
                    from geneva.manifest.mgr import GenevaManifest

                    manifest = GenevaManifest(name="test_manifest")

                    db.define_manifest("test_manifest", manifest)

                    # Verify Uploader was called with upload_dir (not table context)
                    mock_uploader_class.assert_called_once_with(
                        upload_dir="s3://my-upload-bucket/manifests"
                    )


@pytest.mark.parametrize(
    ("tbl_uri", "expected"),
    [
        # Simple bucket
        ("s3://bucket/table.lance", "s3://bucket"),
        # Bucket with path
        (
            "s3://bucket/path/table.lance",
            "s3://bucket/path",
        ),
        # local
        (
            "/foo/bar/table.lance",
            "/foo/bar",
        ),
        # GCS bucket
        ("gs://gcs-bucket/foo/table.lance", "gs://gcs-bucket/foo"),
        # Remote URI
        ("db://foo/table.lance", "db://foo"),
        # S3+DDB with query params
        (
            "s3+ddb://bucket/path/table.lance?ddbTableName=external-manifest",
            "s3+ddb://bucket/path?ddbTableName=external-manifest",
        ),
        # S3+DDB nested path with query params
        (
            "s3+ddb://bucket/path/to/table.lance?ddbTableName=my-table",
            "s3+ddb://bucket/path/to?ddbTableName=my-table",
        ),
    ],
)
def test_get_db_uri_from_table(tbl_uri, expected) -> None:
    from geneva.db import _get_db_uri

    result = _get_db_uri(tbl_uri)
    assert result == expected


@pytest.mark.parametrize(
    ("db_uri", "table_id", "expected_suffix"),
    [
        # Simple S3 URI
        (
            "s3://bucket/path",
            ["my_table"],
            "s3://bucket/path/my_table.lance/_geneva_uploads",
        ),
        # S3+DDB URI with query params — query string must come after the full path
        (
            "s3+ddb://bucket/path?ddbTableName=external-manifest",
            ["my_table"],
            "s3+ddb://bucket/path/my_table.lance/_geneva_uploads?ddbTableName=external-manifest",
        ),
    ],
)
def test_uploader_db_uri_with_query_params(db_uri, table_id, expected_suffix) -> None:
    from geneva.packager.uploader import Uploader

    uploader = Uploader(db_uri=db_uri, table_id=table_id)
    assert uploader.upload_dir == expected_suffix
