# SPDX-License-Identifier: PROPRIETARY
# SPDX-FileCopyrightText: Copyright The Geneva Authors
import contextlib
import copy
import logging
import threading
from collections.abc import Callable, Iterable
from datetime import datetime, timedelta, timezone
from functools import cached_property
from pathlib import Path
from typing import TYPE_CHECKING, Any, TypeVar
from urllib.parse import urlparse, urlunparse

import attrs
import lancedb
import pyarrow as pa
from lance_namespace import LanceNamespace
from lance_namespace import connect as namespace_connect
from lancedb import DBConnection, LanceNamespaceDBConnection
from lancedb.common import DATA, Credential
from lancedb.pydantic import LanceModel
from lancedb.table import Table as LanceDBTable
from lancedb.util import get_uri_scheme
from overrides import override
from yarl import URL

from geneva import DEFAULT_UPLOAD_DIR
from geneva.checkpoint import CheckpointStore
from geneva.cluster import GenevaClusterType
from geneva.config import ConfigBase
from geneva.packager import DockerUDFPackager, UDFPackager
from geneva.packager.autodetect import upload_local_env
from geneva.packager.uploader import Uploader

if TYPE_CHECKING:
    import lance

    from geneva.cluster.mgr import ClusterConfigManager, GenevaCluster
    from geneva.jobs.jobs import JobStateManager
    from geneva.manifest.mgr import GenevaManifest, ManifestConfigManager
    from geneva.query import GenevaQueryBuilder
    from geneva.table import Table
    from geneva.transformer import UDTF, ScalarUDTF

_LOG = logging.getLogger(__name__)

T = TypeVar("T")

SYSTEM_NAMESPACE = "__system"
SYSTEM_TABLE_MIGRATIONS_TABLE_NAME = "geneva_migrations"
SYSTEM_TABLE_MIGRATING = "migrating"
SYSTEM_TABLE_COMPLETED = "completed"


@attrs.define
class SystemTableMigrationRecord:
    table_name: str = attrs.field()
    status: str = attrs.field()
    updated_at: datetime = attrs.field(
        factory=lambda: datetime.now(timezone.utc),
        metadata={"pa_type": pa.timestamp("us", tz="UTC")},
    )


def uses_namespace_for_system_tables(
    uri: str | None,
    namespace_impl: str | None,
) -> bool:
    """Whether a connection/reference should locate system tables via namespaces."""
    return namespace_impl is not None or bool(uri and str(uri).startswith("db://"))


def infer_system_namespace(
    *,
    uri: str | None,
    namespace_impl: str | None,
    system_namespace: list[str] | None,
) -> list[str]:
    if system_namespace is not None:
        return system_namespace
    if uses_namespace_for_system_tables(uri, namespace_impl):
        return [SYSTEM_NAMESPACE]
    return []


def infer_system_db_uri(
    *,
    uri: str | None,
    namespace_impl: str | None,
) -> str | None:
    if uri is None:
        return None
    if uses_namespace_for_system_tables(uri, namespace_impl):
        return uri
    return str(URL(str(uri)) / SYSTEM_NAMESPACE)


# Lance internal columns that are always available but not in the user schema
_LANCE_INTERNAL_COLUMNS = {"_rowid", "_rowaddr"}


def has_stable_row_ids(fragments: "Iterable[lance.LanceFragment]") -> bool:
    """Check if Lance fragments have stable row IDs enabled.

    Stable row IDs are indicated by presence of row_id_meta on fragment metadata.
    This is a Lance feature (added in v0.21.0) that ensures row identifiers remain
    constant even when table operations like compaction reorganize the physical data.

    Parameters
    ----------
    fragments : Iterable[lance.LanceFragment]
        Lance fragments to check (from dataset.get_fragments())

    Returns
    -------
    bool
        True if any fragment has stable row IDs enabled, False otherwise
    """
    return any(frag.metadata.row_id_meta is not None for frag in fragments)


class Connection(DBConnection):
    """Geneva Connection."""

    def __init__(
        self,
        uri: str,
        *,
        region: str = "us-east-1",
        api_key: Credential | None = None,
        host_override: str | None = None,
        storage_options: dict[str, str] | None = None,
        checkpoint_store: CheckpointStore | None = None,
        packager: UDFPackager | None = None,
        namespace_impl: str | None = None,
        namespace_properties: dict[str, str] | None = None,
        system_namespace: list[str],
        is_system_db_root: bool = False,
        upload_dir: str | None = None,
        **kwargs,
    ) -> None:
        super().__init__()

        self._uri = uri
        self._region = region
        self._api_key = api_key
        self._host_override = host_override
        self._storage_options = storage_options
        self._ldb: DBConnection | None = None
        self._checkpoint_store = checkpoint_store
        self._packager = packager or DockerUDFPackager()
        self.namespace_impl = namespace_impl
        self.namespace_properties = namespace_properties
        self._is_system_db_root = is_system_db_root
        # Use namespace-based system tables for namespace and remote connections.
        # For local/object-store direct connections, system tables live under a
        # sibling "__system" database URI.
        uses_namespace_system_tables = uses_namespace_for_system_tables(
            uri, namespace_impl
        )
        self.system_namespace = system_namespace if uses_namespace_system_tables else []
        self._upload_dir = upload_dir

        self._jobs_manager: JobStateManager | None = None
        self._cluster_manager: ClusterConfigManager | None = None
        self._manifest_manager: ManifestConfigManager | None = None
        self._flight_client: Any | None = None
        self._namespace_connection: LanceNamespaceDBConnection | None = None
        self._system_table_conn: Connection | None = None
        self._system_table_conn_lock = threading.Lock()
        self._system_table_migrations: set[tuple[str, tuple[str, ...], str]] = set()
        self._system_namespace_ensured = False
        self._system_namespace_fallback: list[str] | None = None
        self._kwargs = kwargs

    def __repr__(self) -> str:
        return f"<Geneva uri={self.uri}>"

    def __getstate__(self) -> dict:
        return {
            "uri": self._uri,
            "api_key": self._api_key,
            "host_override": self._host_override,
            "storage_options": self._storage_options,
            "region": self._region,
            "namespace_impl": self.namespace_impl,
            "namespace_properties": self.namespace_properties,
            "system_namespace": self.effective_system_namespace,
            "is_system_db_root": self._is_system_db_root,
            "upload_dir": self._upload_dir,
        }

    def __setstate__(self, state) -> None:
        self.__init__(state.pop("uri"), **state)

    def __enter__(self) -> "Connection":
        return self

    def __exit__(self, exc_type, exc_value, traceback) -> None:
        self.close()
        return None  # Don't suppress exceptions

    def close(self) -> None:
        """Close the connection."""
        if self._flight_client is not None:
            self._flight_client.close()
        if self._ldb is not None and hasattr(self._ldb, "_conn"):
            # go to the async client and eagerly close the connection
            self._ldb._conn.close()  # type: ignore[attr-defined]
        if self._system_table_conn is not None:
            self._system_table_conn.close()

    def _uses_namespace_for_system_tables(self) -> bool:
        """Whether system tables should use namespace APIs on this connection."""
        return uses_namespace_for_system_tables(self._uri, self.namespace_impl)

    def _is_system_table_db_root(self) -> bool:
        """Whether this direct connection already points at the system DB root."""
        return self._is_system_db_root

    @property
    def user_db_uri(self) -> str:
        """Original user database URI for this connection."""
        if not self._is_system_table_db_root():
            return str(self.uri)

        parsed = urlparse(str(self.uri))
        path = parsed.path.rstrip("/")
        if path.endswith(f"/{SYSTEM_NAMESPACE}"):
            path = path[: -(len(SYSTEM_NAMESPACE) + 1)] or "/"
        return urlunparse(parsed._replace(path=path))

    @property
    def effective_system_namespace(self) -> list[str]:
        """Effective namespace after backend-compatibility fallback."""
        if self._system_namespace_fallback is not None:
            return self._system_namespace_fallback
        return self.system_namespace

    @property
    def system_table_db_uri(self) -> str:
        """Database URI where system tables should live."""
        if self._uses_namespace_for_system_tables() or self._is_system_table_db_root():
            return str(self.uri)
        return str(URL(str(self.uri)) / SYSTEM_NAMESPACE)

    @property
    def system_table_connection(self) -> "Connection":
        """Connection used to access system tables."""
        if self._uses_namespace_for_system_tables() or self._is_system_table_db_root():
            return self

        if self._system_table_conn is None:
            with self._system_table_conn_lock:
                if self._system_table_conn is None:
                    self._system_table_conn = Connection(
                        self.system_table_db_uri,
                        region=self._region,
                        api_key=self._api_key,
                        host_override=self._host_override,
                        storage_options=self._storage_options,
                        checkpoint_store=self._checkpoint_store,
                        packager=self._packager,
                        namespace_impl=None,
                        namespace_properties=None,
                        system_namespace=[],
                        is_system_db_root=True,
                        upload_dir=self._upload_dir,
                        **self._kwargs,
                    )
        return self._system_table_conn

    def resolve_system_table_location(
        self, namespace: list[str] | None = None
    ) -> tuple["Connection", list[str]]:
        """Resolve the target connection/namespace for system table operations."""
        if namespace is not None:
            return self, namespace
        if self._uses_namespace_for_system_tables():
            return self, self.effective_system_namespace
        return self.system_table_connection, []

    def _run_system_table_operation(
        self,
        table_name: str,
        operation: Callable[["Connection", list[str]], T],
        *,
        namespace: list[str] | None = None,
        migrate_legacy: bool = True,
    ) -> T:
        """Run a system-table operation with one remote fallback path."""

        def run_on_target(
            target_db: "Connection",
            target_namespace: list[str],
        ) -> T:
            if target_namespace and target_db._uses_namespace_for_system_tables():
                _ensure_system_namespace_exists(target_db)
            if migrate_legacy:
                self.migrate_legacy_system_table_if_needed(
                    table_name=table_name,
                    target_conn=target_db,
                    target_namespace=target_namespace,
                )
            return operation(target_db, target_namespace)

        target_db, target_namespace = self.resolve_system_table_location(namespace)
        try:
            return run_on_target(target_db, target_namespace)
        except Exception as e:
            if (
                namespace is not None
                or not self.is_remote()
                or not target_namespace
                or not self._is_remote_namespace_location_error(e)
            ):
                raise

            _LOG.warning(
                "Remote backend rejected namespace-backed system table '%s'; "
                "falling back to root namespace for compatibility",
                table_name,
            )
            self._system_namespace_fallback = []
            fallback_db, fallback_namespace = self.resolve_system_table_location(None)
            return run_on_target(fallback_db, fallback_namespace)

    @staticmethod
    def _is_remote_namespace_location_error(exc: Exception) -> bool:
        return (
            "location must be provided when namespace is not empty" in str(exc).lower()
        )

    @staticmethod
    def _is_not_found_error(exc: Exception) -> bool:
        msg = str(exc).lower()
        return "not found" in msg or "does not exist" in msg

    @staticmethod
    def _table_exists(
        conn: "Connection", table_name: str, namespace: list[str] | None
    ) -> bool:
        ns = namespace if namespace is not None else []
        try:
            conn.open_table(table_name, namespace=ns)
            return True
        except Exception as e:
            if Connection._is_not_found_error(e):
                return False
            _LOG.debug(
                "Error checking existence of table '%s' in namespace=%s on uri=%s",
                table_name,
                ns,
                conn.uri,
                exc_info=True,
            )
            raise

    @staticmethod
    def _replace_table_from_source(
        *,
        src_conn: "Connection",
        src_namespace: list[str] | None,
        dst_conn: "Connection",
        dst_namespace: list[str] | None,
        table_name: str,
    ) -> None:
        src_ns = src_namespace if src_namespace is not None else []
        dst_ns = dst_namespace if dst_namespace is not None else []

        src_tbl = src_conn.open_table(table_name, namespace=src_ns)
        src_data = src_tbl.search().to_arrow()
        src_row_count = src_tbl.count_rows()

        if Connection._table_exists(dst_conn, table_name, dst_ns):
            dst_conn.drop_table(table_name, namespace=dst_ns)

        dst_conn.create_table(
            table_name,
            schema=src_tbl.schema,
            mode="overwrite",
            namespace=dst_ns,
        )

        if src_row_count > 0:
            dst_tbl = dst_conn.open_table(table_name, namespace=dst_ns)
            dst_tbl.add(src_data)

        dst_tbl = dst_conn.open_table(table_name, namespace=dst_ns)
        if dst_tbl.count_rows() != src_row_count:
            raise RuntimeError(
                f"Migration validation failed for '{table_name}': "
                f"expected {src_row_count} rows but copied {dst_tbl.count_rows()}"
            )

    def _legacy_system_table_candidates(
        self,
        target_conn: "Connection",
        target_namespace: list[str] | None,
    ) -> list[tuple["Connection", list[str]]]:
        dst_ns = target_namespace if target_namespace is not None else []
        candidates: list[tuple[Connection, list[str]]] = []

        if self._uses_namespace_for_system_tables():
            if self.namespace_impl is not None and dst_ns != ["default"]:
                candidates.append((self, ["default"]))
            if self.is_remote() and dst_ns != []:
                candidates.append((self, []))
        else:
            candidates.append((self, []))

        return [
            (legacy_conn, legacy_ns)
            for legacy_conn, legacy_ns in candidates
            if not (legacy_conn is target_conn and legacy_ns == dst_ns)
        ]

    def _find_legacy_system_table_source(
        self,
        *,
        table_name: str,
        target_conn: "Connection",
        target_namespace: list[str] | None,
    ) -> tuple["Connection", list[str]] | None:
        for legacy_conn, legacy_ns in self._legacy_system_table_candidates(
            target_conn,
            target_namespace,
        ):
            if self._table_exists(legacy_conn, table_name, legacy_ns):
                return legacy_conn, legacy_ns
        return None

    def _open_system_table_migrations_table(
        self, create_if_missing: bool
    ) -> tuple["Connection", LanceDBTable, list[str]] | None:
        from geneva.utils.schema import alter_or_create_table

        def open_or_create(
            target_db: "Connection",
            target_namespace: list[str],
        ) -> tuple["Connection", LanceDBTable, list[str]] | None:
            if not create_if_missing and not self._table_exists(
                target_db,
                SYSTEM_TABLE_MIGRATIONS_TABLE_NAME,
                target_namespace,
            ):
                return None

            table = alter_or_create_table(
                target_db,
                SYSTEM_TABLE_MIGRATIONS_TABLE_NAME,
                SystemTableMigrationRecord(
                    table_name="dummy",
                    status=SYSTEM_TABLE_COMPLETED,
                ),
                namespace=target_namespace,
            )
            return target_db, table, target_namespace

        return self._run_system_table_operation(
            SYSTEM_TABLE_MIGRATIONS_TABLE_NAME,
            open_or_create,
            migrate_legacy=False,
        )

    def _get_system_table_migration_status(self, table_name: str) -> str | None:
        table_info = self._open_system_table_migrations_table(create_if_missing=False)
        if table_info is None:
            return None
        _, table, _ = table_info
        rows = (
            table._ltbl.search()  # pyright: ignore[reportAttributeAccessIssue]
            .where(f"table_name = '{table_name}'")
            .limit(1)
            .to_arrow()
            .to_pylist()
        )
        if not rows:
            return None
        return rows[0]["status"]

    def _set_system_table_migration_status(self, table_name: str, status: str) -> None:
        table_info = self._open_system_table_migrations_table(create_if_missing=True)
        assert table_info is not None, (
            "migration table must exist when create_if_missing=True"
        )
        _, table, _ = table_info
        ltbl = table._ltbl  # pyright: ignore[reportAttributeAccessIssue]
        ltbl.delete(f"table_name = '{table_name}'")
        ltbl.add(
            [
                attrs.asdict(
                    SystemTableMigrationRecord(
                        table_name=table_name,
                        status=status,
                    )
                )
            ]
        )

    def _clear_system_table_migration_status(self, table_name: str) -> None:
        table_info = self._open_system_table_migrations_table(create_if_missing=False)
        if table_info is None:
            return
        _, table, _ = table_info
        ltbl = table._ltbl  # pyright: ignore[reportAttributeAccessIssue]
        ltbl.delete(f"table_name = '{table_name}'")

    def migrate_legacy_system_table_if_needed(
        self,
        *,
        table_name: str,
        target_conn: "Connection",
        target_namespace: list[str] | None,
    ) -> None:
        """Best-effort migration from legacy system-table locations.

        Legacy locations:
        - Local/object-store direct connections: <db>/<table>.lance (root)
        - Remote connections: root namespace []
        - Namespace connections: default namespace ["default"]
        """
        dst_ns = target_namespace if target_namespace is not None else []
        key = (table_name, tuple(dst_ns), str(target_conn.uri))
        if key in self._system_table_migrations:
            return
        if table_name == SYSTEM_TABLE_MIGRATIONS_TABLE_NAME:
            self._system_table_migrations.add(key)
            return

        status = self._get_system_table_migration_status(table_name)
        legacy_source = self._find_legacy_system_table_source(
            table_name=table_name,
            target_conn=target_conn,
            target_namespace=target_namespace,
        )

        if status == SYSTEM_TABLE_COMPLETED:
            self._system_table_migrations.add(key)
            return

        if status is None and legacy_source is None:
            self._system_table_migrations.add(key)
            return

        if legacy_source is None:
            if status == SYSTEM_TABLE_MIGRATING:
                _LOG.warning(
                    "Clearing stale migration marker for system table '%s' on uri=%s "
                    "namespace=%s because no legacy source table was found",
                    table_name,
                    target_conn.uri,
                    dst_ns,
                )
                self._clear_system_table_migration_status(table_name)
                self._system_table_migrations.add(key)
                return
            raise RuntimeError(
                f"System table '{table_name}' is marked as migrating but no legacy "
                "source table was found"
            )

        legacy_conn, legacy_ns = legacy_source
        self._set_system_table_migration_status(table_name, SYSTEM_TABLE_MIGRATING)
        _LOG.info(
            "Migrating legacy system table '%s' from namespace=%s uri=%s "
            "to namespace=%s uri=%s",
            table_name,
            legacy_ns,
            legacy_conn.uri,
            dst_ns,
            target_conn.uri,
        )
        self._replace_table_from_source(
            src_conn=legacy_conn,
            src_namespace=legacy_ns,
            dst_conn=target_conn,
            dst_namespace=dst_ns,
            table_name=table_name,
        )
        self._set_system_table_migration_status(table_name, SYSTEM_TABLE_COMPLETED)
        _LOG.warning(
            "Legacy system table '%s' remains at namespace=%s uri=%s after "
            "copying to uri=%s for rollback safety",
            table_name,
            legacy_ns,
            legacy_conn.uri,
            target_conn.uri,
        )
        self._system_table_migrations.add(key)

    def prepare_system_table_target(
        self,
        table_name: str,
        namespace: list[str] | None = None,
    ) -> tuple["Connection", list[str]]:
        """Resolve and prepare the destination for a system-table operation."""
        return self._run_system_table_operation(
            table_name,
            lambda target_db, target_namespace: (target_db, target_namespace),
            namespace=namespace,
        )

    def alter_or_create_system_table(
        self,
        table_name: str,
        model: Any,
        namespace: list[str] | None = None,
    ) -> tuple["Connection", LanceDBTable]:
        """Open or create a Geneva internal table in its resolved location."""
        from geneva.utils.schema import alter_or_create_table

        return self._run_system_table_operation(
            table_name,
            lambda target_db, target_namespace: (
                target_db,
                alter_or_create_table(
                    target_db,
                    table_name,
                    model,
                    namespace=target_namespace,
                ),
            ),
            namespace=namespace,
        )

    @cached_property
    def namespace_client(self) -> LanceNamespace | None:
        """Returns namespace client if using namespace connection."""
        if self.namespace_impl is not None and self.namespace_properties is not None:
            return namespace_connect(self.namespace_impl, self.namespace_properties)
        return None

    @cached_property
    def _connect(self) -> DBConnection:
        conn = self._get_conn(uri=self.uri)

        if self.namespace_impl is not None:
            if self._namespace_connection is None:
                self._namespace_connection = conn  # type: ignore[assignment]
        elif self._ldb is None:
            self._ldb = conn
        return conn

    def _get_conn(self, uri: str | None) -> DBConnection:
        """Returns the underlying lancedb connection. If the connection
        was initialized using a namespace, this will return the
        LanceNamespaceDBConnection.
        """
        # use the namespace connection if set
        if self.namespace_impl is not None:
            if not self._namespace_connection:
                assert self.namespace_properties is not None, (
                    "namespace_properties must be set when namespace_impl is set"
                )
                namespace_client = namespace_connect(
                    self.namespace_impl, self.namespace_properties
                )

                ns_conn = LanceNamespaceDBConnection(
                    namespace_client,
                    storage_options=self._storage_options,
                    **self._kwargs,
                )
            _LOG.info(f"using namespace connection: {self._namespace_connection=}")
            return ns_conn

        conn = lancedb.connect(
            uri or self.uri,
            region=self._region,
            api_key=self._api_key,
            host_override=self._host_override,
            storage_options=self._storage_options,
            **self._kwargs,
        )
        return conn

    @cached_property
    def _history(self) -> "JobStateManager":  # noqa: F821
        """Returns a JobStateManager that persists job executions and statuses"""
        from geneva.jobs import JobStateManager

        if self._jobs_manager is None:
            self._jobs_manager = JobStateManager(self)
        return self._jobs_manager

    @cached_property
    def flight_client(self) -> Any:
        if self._flight_client is not None:
            return self._flight_client
        import flightsql

        flightsql_client = flightsql.FlightSQLClient
        url = urlparse(self._host_override)
        hostname = url.hostname
        client = flightsql_client(
            host=hostname,
            port=10025,
            token="DATABASE_TOKEN",  # Dummy auth, not plugged in yet
            metadata={"database": self.uri},  # Name of the project-id
            features={"metadata-reflection": "true"},
            insecure=True,  # or False, up to you
        )
        self._flight_client = client
        return client

    @override
    def table_names(
        self, page_token: str | None = None, limit: int | None = None, *args, **kwargs
    ) -> Iterable[str]:
        """List all available tables and views."""
        return self._connect.table_names(
            *args, page_token=page_token, limit=limit or 10, **kwargs
        )

    @override
    def open_table(
        self,
        name: str,
        storage_options: dict[str, str] | None = None,
        index_cache_size: int | None = None,
        version: int | None = None,
        namespace: list[str] | None = None,
        *args,
        **kwargs,
    ) -> "Table":
        """Open a Lance Table.

        Parameters
        ----------
        name: str
            Name of the table.
        storage_options: dict[str, str], optional
            Additional options for the storage backend.
            Options already set on the connection will be inherited by the table,
            but can be overridden here. See available options at
            [https://lancedb.github.io/lancedb/guides/storage/](https://lancedb.github.io/lancedb/guides/storage/)
        namespace: list[str], optional
            Namespace path for the table (e.g., ["workspace"])

        """
        from .table import Table

        storage_options = storage_options or self._storage_options
        namespace = namespace if namespace is not None else []

        return Table(
            self,
            name,
            namespace=namespace,
            index_cache_size=index_cache_size,
            storage_options=storage_options,
            version=version,
        )

    def _validate_existing_table_has_stable_row_ids(self, name: str) -> None:
        """Validate that an existing table has stable row IDs enabled.

        Called when exist_ok=True and stable row IDs are requested, to ensure
        the existing table matches the requested configuration.

        Raises:
            ValueError: If table exists with data but doesn't have stable row IDs
        """
        # Use lancedb table API (not lance.dataset directly) to support namespaces
        if name not in self._connect.table_names():
            return

        try:
            table = self._connect.open_table(name)
            fragments = list(table.to_lance().get_fragments())  # type: ignore[union-attr]
            if fragments:
                if not has_stable_row_ids(fragments):
                    raise ValueError(
                        f"Cannot open table '{name}' with exist_ok=True: "
                        f"table exists but does not have stable row IDs enabled.\n\n"
                        f"You requested stable row IDs, but the existing table "
                        f"was created without them.\n\n"
                        f"Options:\n"
                        f"  1. Drop and recreate the table with stable row IDs\n"
                        f"  2. Remove storage_options if stable row IDs are "
                        f"not required\n"
                        f"  3. Use mode='overwrite' to replace the table"
                    )
            else:
                _LOG.warning(
                    f"Table '{name}' exists but is empty - cannot verify "
                    f"if stable row IDs are enabled."
                )
        except ValueError:
            raise
        except Exception as e:
            _LOG.debug(
                f"Could not check stable row ID status for existing table '{name}': {e}"
            )

    @override
    def create_table(  # type: ignore
        self,
        name: str,
        data: DATA | None = None,
        schema: pa.Schema | LanceModel | None = None,
        mode: str = "create",
        exist_ok: bool = False,
        on_bad_vectors: str = "error",
        fill_value: float = 0.0,
        storage_options: dict[str, str] | None = None,
        *args,
        **kwargs,
    ) -> "Table":  # type: ignore
        """Create a Table in the lake

        Parameters
        ----------
        name: str
            The name of the table
        data: The data to initialize the table, *optional*
            User must provide at least one of `data` or `schema`.
            Acceptable types are:

            - list-of-dict
            - pandas.DataFrame
            - pyarrow.Table or pyarrow.RecordBatch
        schema: The schema of the table, *optional*
            Acceptable types are:

            - pyarrow.Schema
            - lancedb.pydantic.LanceModel
        mode: str; default "create"
            The mode to use when creating the table.
            Can be either "create" or "overwrite".
            By default, if the table already exists, an exception is raised.
            If you want to overwrite the table, use mode="overwrite".
        exist_ok: bool, default False
            If a table by the same name already exists, then raise an exception
            if exist_ok=False. If exist_ok=True, then open the existing table;
            it will not add the provided data but will validate against any
            schema that's specified.
        on_bad_vectors: str, default "error"
            What to do if any of the vectors are not the same size or contain NaNs.
            One of "error", "drop", "fill".
        """
        from .table import Table

        # Handle new_table_enable_stable_row_ids validation and normalization
        # Accept both string "true" and boolean True for user convenience
        if storage_options:
            enable_stable = storage_options.get("new_table_enable_stable_row_ids")
            if enable_stable in ("true", True, 1):
                # Normalize to string "true" for lancedb
                storage_options = dict(storage_options)
                storage_options["new_table_enable_stable_row_ids"] = "true"

                # Validate exist_ok: ensure existing table has stable row IDs
                if exist_ok:
                    self._validate_existing_table_has_stable_row_ids(name)

        # Create table through lancedb
        conn = self._connect

        if self._host_override:
            # in OSS, exist_ok is a separate param, but in phalanx it is set as "mode"
            # workaround until https://github.com/lancedb/lancedb/issues/2900
            if exist_ok and mode == "create":
                mode = "exist_ok"

            if storage_options:
                _LOG.warning(
                    "storage_options parameter is not supported when creating "
                    "tables on remote connections, ignoring"
                )
        else:
            # these params are not supported in remote connections
            kwargs.update(
                exist_ok=exist_ok,
                on_bad_vectors=on_bad_vectors,
                storage_options=storage_options,
            )

        conn.create_table(
            name,
            data,
            schema,
            mode,
            *args,
            fill_value=fill_value,
            **kwargs,
        )
        # Extract namespace from kwargs to pass to Table
        namespace = kwargs.get("namespace")
        return Table(self, name, namespace=namespace, storage_options=storage_options)

    def create_view(
        self,
        name: str,
        query: str,
        materialized: bool = False,
    ) -> "Table":
        """Create a View from a Query.

        Parameters
        ----------
        name: str
            Name of the view.
        query: str
            SQL query to create the view.
        materialized: bool, optional
            If True, the view is materialized.
        """
        if materialized:
            # idea, rename the provided name, and use it as the basis for the
            # materialized view.
            # - how do we add the udfs to the final materialized view table?
            NotImplementedError(
                "creating materialized view via sql query is not supported yet."
            )

        # TODO add test coverage here
        self.sql(f"CREATE VIEW {name} AS ({query})")  # type: ignore[attr-defined]
        return self.open_table(name)

    def create_materialized_view(
        self,
        name: str,
        query: "GenevaQueryBuilder",
        with_no_data: bool = True,
    ) -> "Table":
        """
        Create a materialized view

        Parameters
        ----------
        name: str
            Name of the materialized view.
        query: GenevaQueryBuilder
            Query to create the view.
        with_no_data: bool, optional
            If True, the view is materialized, if false it is ready for refresh.
        """
        from geneva.query import GenevaQueryBuilder

        if not isinstance(query, GenevaQueryBuilder):
            raise ValueError(
                "Materialized views only support plain queries (where, select)"
            )

        tbl = query.create_materialized_view(self, name)
        if not with_no_data and hasattr(tbl, "refresh_view"):
            tbl.refresh_view(name)  # type: ignore[attr-defined]

        return tbl

    def create_udtf_view(
        self,
        name: str,
        source: "GenevaQueryBuilder",
        udtf: "UDTF",
    ) -> "Table":
        """Create a UDTF-backed materialized view.

        .. warning::
            This API is in **beta** and may change in future releases.

        The view is created empty; call ``view.refresh()`` to populate it.

        Parameters
        ----------
        name : str
            Name for the new view table.
        source : GenevaQueryBuilder
            Query defining the source data.
        udtf : UDTF
            The UDTF to execute on refresh.
        """

        from geneva.packager import marshal_udtf
        from geneva.query import (
            MATVIEW_META_BASE_DBURI,
            MATVIEW_META_BASE_TABLE,
            MATVIEW_META_BASE_VERSION,
            MATVIEW_META_NAMESPACE_PATH,
            MATVIEW_META_SOURCE_QUERY,
            MATVIEW_META_UDTF,
            MATVIEW_META_VERSION,
            GenevaQueryBuilder,
        )

        if not isinstance(source, GenevaQueryBuilder):
            raise ValueError("source must be a GenevaQueryBuilder")

        # Validate input_columns against source schema if specified
        # Note: _rowid/_rowaddr are Lance internal columns always available
        if udtf.input_columns:
            src_schema = source._table.schema
            missing = [
                c
                for c in udtf.input_columns
                if c not in src_schema.names and c not in _LANCE_INTERNAL_COLUMNS
            ]
            if missing:
                raise ValueError(
                    f"UDTF input_columns {missing} not found in source schema. "
                    f"Available: {src_schema.names}"
                )

        # Serialize UDTF and source query
        udtf_spec = marshal_udtf(udtf)
        source_query_json = source.to_query_object().model_dump_json()

        # Get source table info
        local_tbl = source._table._native_tbl
        db_uri = _get_db_uri(source._table.uri)

        # Build metadata
        metadata = {
            MATVIEW_META_UDTF: udtf_spec.to_json(),
            MATVIEW_META_SOURCE_QUERY: source_query_json,
            MATVIEW_META_BASE_TABLE: local_tbl.name,
            MATVIEW_META_BASE_DBURI: db_uri,
            MATVIEW_META_BASE_VERSION: str(local_tbl.version),
            MATVIEW_META_VERSION: "udtf",
        }

        # Store namespace path as $-delimited string (if source is in a child namespace)
        if source._table._namespace:
            metadata[MATVIEW_META_NAMESPACE_PATH] = "$".join(source._table._namespace)

        # Create empty table with UDTF output schema + metadata.
        # Delegates to self.create_table() which is namespace-aware.
        schema_with_meta = udtf.output_schema.with_metadata(metadata)
        empty_table = schema_with_meta.empty_table()

        return self.create_table(name, data=empty_table)

    def create_scalar_udtf_view(
        self,
        name: str,
        source: "GenevaQueryBuilder",
        scalar_udtf: "ScalarUDTF",
    ) -> "Table":
        """Create a scalar UDTF-backed materialized view (1:N row expansion).

        .. warning::
            This API is in **beta** and may change in future releases.

        The view is created with placeholder rows (one per source row) that are
        populated on ``view.refresh()``. Each source row expands to zero or more
        output rows via the scalar UDTF.

        Parameters
        ----------
        name : str
            Name for the new view table.
        source : GenevaQueryBuilder
            Query defining the source data.
        scalar_udtf : ScalarUDTF
            The scalar UDTF to execute on refresh.
        """

        from geneva.packager import marshal_scalar_udtf
        from geneva.query import (
            MATVIEW_META_BASE_DBURI,
            MATVIEW_META_BASE_TABLE,
            MATVIEW_META_BASE_VERSION,
            MATVIEW_META_NAMESPACE_PATH,
            MATVIEW_META_SCALAR_UDTF,
            MATVIEW_META_SOURCE_QUERY,
            MATVIEW_META_VERSION,
            MATVIEW_VERSION_SCALAR_UDTF,
            GenevaQueryBuilder,
        )

        if not isinstance(source, GenevaQueryBuilder):
            raise ValueError("source must be a GenevaQueryBuilder")

        # Validate input columns against query projection (not base table)
        if scalar_udtf.input_columns:
            projected = source._schema_for_query(include_metacols=False)
            projected_names = projected.names
            missing = [c for c in scalar_udtf.input_columns if c not in projected_names]
            if missing:
                raise ValueError(
                    f"ScalarUDTF input_columns {missing} not found "
                    f"in query projection. "
                    f"Available: {projected_names}"
                )

        # Serialize scalar UDTF and source query
        spec = marshal_scalar_udtf(scalar_udtf)
        source_query_json = source.to_query_object().model_dump_json()

        # Get source table info
        local_tbl = source._table._native_tbl
        db_uri = _get_db_uri(source._table.uri)

        # Build the output schema: __source_row_id + __child_index
        # + inherited source columns + UDTF output columns
        inherited_schema = source._schema_for_query(include_metacols=False)
        view_fields = [
            pa.field("__source_row_id", pa.int64()),
            pa.field("__child_index", pa.int32()),
        ]
        # Add inherited source columns
        view_fields.extend(
            field
            for field in inherited_schema
            if field.name not in ("__source_row_id", "__child_index", "__is_set")
        )
        # Add UDTF output columns
        view_fields.extend(scalar_udtf.output_schema)

        view_schema = pa.schema(view_fields)

        # Build metadata
        metadata = {
            MATVIEW_META_SCALAR_UDTF: spec.to_json(),
            MATVIEW_META_SOURCE_QUERY: source_query_json,
            MATVIEW_META_BASE_TABLE: local_tbl.name,
            MATVIEW_META_BASE_DBURI: db_uri,
            MATVIEW_META_BASE_VERSION: str(local_tbl.version),
            MATVIEW_META_VERSION: MATVIEW_VERSION_SCALAR_UDTF,
        }

        # Store namespace path as $-delimited string (if source is in a child namespace)
        if source._table._namespace:
            metadata[MATVIEW_META_NAMESPACE_PATH] = "$".join(source._table._namespace)

        # Create empty table with the full schema + metadata
        schema_with_meta = view_schema.with_metadata(metadata)
        empty_table = schema_with_meta.empty_table()

        return self.create_table(name, data=empty_table)

    def drop_view(self, name: str) -> pa.Table:
        """Drop a view."""
        return self.sql(f"DROP VIEW {name}")  # type: ignore[attr-defined]

    @override
    def drop_table(self, name: str, *args, **kwargs) -> None:
        """Drop a table."""
        self._connect.drop_table(name, *args, **kwargs)

    def define_cluster(self, name: str, cluster: "GenevaCluster") -> None:  # noqa: F821
        """
        Define a persistent Geneva cluster. This will upsert the cluster definition by
        name. The cluster can then be provisioned using `context(cluster=name)`.

        Parameters
        ----------
        name: str
            Name of the cluster. This will be used as the key when upserting and
            provisioning the cluster. The cluster name must comply with RFC 12123.
        cluster: GenevaCluster
            The cluster definition to store.
        """
        from geneva.cluster.mgr import ClusterConfigManager

        if self._cluster_manager is None:
            self._cluster_manager = ClusterConfigManager(self)

        cluster.name = name
        self._cluster_manager.upsert(cluster)

    def list_clusters(self) -> list["GenevaCluster"]:  # noqa: F821
        """
        List the cluster definitions. These can be defined using `define_cluster()`.

        Returns
        -------
        list[GenevaCluster]
            List of Geneva cluster definitions
        """
        from geneva.cluster.mgr import ClusterConfigManager

        if self._cluster_manager is None:
            self._cluster_manager = ClusterConfigManager(self)
        return self._cluster_manager.list()

    def delete_cluster(self, name: str) -> None:  # noqa: F821
        """
        Delete a Geneva cluster definition.

        Parameters
        ----------
        name: str
            Name of the cluster to delete.
        """
        from geneva.cluster.mgr import ClusterConfigManager

        if self._cluster_manager is None:
            self._cluster_manager = ClusterConfigManager(self)

        self._cluster_manager.delete(name)

    def define_manifest(
        self,
        name: str,
        manifest: "GenevaManifest",  # noqa: F821
        uploader: Uploader | None = None,
    ) -> None:
        """
        Define a persistent Geneva Manifest that represents the files and dependencies
        used in the execution environment. This will upsert the manifest definition by
        name and upload the required artifacts. The manifest can then be used with
        `context(manifest=name)`.

        Parameters
        ----------
        name: str
            Name of the manifest. This will be used as the key when upserting and
            loading the manifest.
        manifest: GenevaManifest
            The manifest definition to use.
        uploader: Uploader, optional
            An optional, custom Uploader to use. If not provided, the uploader will be
            auto-detected based on the
            environment configuration.
        """
        from geneva.manifest.mgr import ManifestConfigManager

        if self._manifest_manager is None:
            self._manifest_manager = ManifestConfigManager(self)

        # Ensure the manifest table exists before creating Uploader
        # This guarantees the table location can be queried
        _ = self._manifest_manager.get_table()

        # If no uploader is provided, create one with manifest table context
        if uploader is None:
            from geneva.manifest.mgr import MANIFEST_TABLE_NAME

            # If upload_dir is configured, use it directly for manifest uploads
            if self._upload_dir:
                uploader = Uploader(upload_dir=self._upload_dir)
            elif self.namespace_impl and self.namespace_properties:
                # Create uploader for namespace table
                # Build table_id for manifest table
                table_id = self.system_namespace + [MANIFEST_TABLE_NAME]
                uploader = Uploader(
                    namespace_impl=self.namespace_impl,
                    namespace_properties=self.namespace_properties,
                    table_id=table_id,
                )
            else:
                # For local/object-store direct connections, use the
                # "__system" DB URI as the manifest-table root.
                table_id = [MANIFEST_TABLE_NAME]
                uploader = Uploader(
                    db_uri=self.system_table_db_uri,
                    table_id=table_id,
                )

        with upload_local_env(
            # todo: implement excludes
            uploader=uploader,
            zip_output_dir=manifest.local_zip_output_dir,
            delete_local_zips=manifest.delete_local_zips,
            skip_site_packages=manifest.skip_site_packages,
        ) as zips:
            m = copy.deepcopy(manifest)
            m.name = name
            m.zips = zips
            m.checksum = manifest.compute_checksum()
            self._manifest_manager.upsert(m)

    def list_manifests(self) -> list["GenevaManifest"]:  # noqa: F821
        """
        List the manifest definitions. These can be defined using `define_manifest()`.

        Returns
        -------
        list[GenevaManifest]
            List of Geneva manifest definitions
        """
        from geneva.manifest.mgr import ManifestConfigManager

        if self._manifest_manager is None:
            self._manifest_manager = ManifestConfigManager(self)
        return self._manifest_manager.list()

    def delete_manifest(self, name: str) -> None:  # noqa: F821
        """
        Delete a Geneva manifest definition.

        Parameters
        ----------
        name: str
            Name of the manifest to delete.
        """
        from geneva.manifest.mgr import ManifestConfigManager

        if self._manifest_manager is None:
            self._manifest_manager = ManifestConfigManager(self)

        self._manifest_manager.delete(name)

    @staticmethod
    def local_ray_context() -> contextlib.AbstractContextManager[None]:
        """Context manager for a local Ray instance.
        This will provision a local Ray instance and return a context manager.
        This is useful for development or small jobs.
        """
        from geneva.runners.ray._mgr import ray_cluster

        return ray_cluster(local=True, log_to_driver=True, logging_level=logging.INFO)

    def context(
        self,
        cluster: str,
        manifest: str | None = None,
        on_exit=None,
        wait_timeout: float | None = None,
        log_to_driver: bool = True,
        logging_level=logging.INFO,
    ) -> contextlib.AbstractContextManager[None]:
        """Context manager for a Geneva Execution Environment.
            This will provision a cluster based on the cluster
            definition and the manifest provided.
            By default, the context manager will delete the cluster on exit.
            This can be configured with the on_exit parameter.
        Parameters
        ----------
        cluster: str
            Name of the persisted cluster definition to use. Required.
            This will raise an exception if the cluster definition was not
            defined via `define_cluster()`.
        manifest: str, optional
            Optional name of the persisted manifest to use. This will
            raise an exception if the manifest definition was not
            defined via `define_manifest()`. If manifest is not provided,
            the local environment will be uploaded.
        on_exit: ExitMode, optional, default ExitMode.DELETE
            Exit mode for the cluster. By default, the cluster waits for all
            running jobs to complete before deleting.
            To retain the cluster when any job fails or the context body
            raises an exception, use `ExitMode.RETAIN_ON_FAILURE`.
            To always retain the cluster, use `ExitMode.RETAIN`.
        wait_timeout: float, optional, default None
            Internal/experimental. Maximum seconds to wait for tracked jobs
            during context exit. Only applies with DELETE or
            RETAIN_ON_FAILURE. None means wait indefinitely. For
            RETAIN_ON_FAILURE, a timeout is treated as a failure and the
            cluster is retained.
        log_to_driver: bool, optional, default True
            Whether to send Ray worker logs to the driver. Defaults to True for
            better visibility in tests and debugging.
        logging_level: int, optional, default logging.INFO
            The logging level for Ray workers. Use logging.DEBUG for detailed logs.
        """
        from geneva.cluster.mgr import ClusterConfigManager
        from geneva.manifest.mgr import ManifestConfigManager
        from geneva.runners.ray._mgr import ray_cluster
        from geneva.runners.ray.raycluster import ExitMode

        if self._cluster_manager is None:
            self._cluster_manager = ClusterConfigManager(self)
        if self._manifest_manager is None:
            self._manifest_manager = ManifestConfigManager(self)

        if cluster is None:
            raise ValueError(
                "cluster name is required. Use conn.context(cluster=...) or "
                "conn.local_ray_context() for a local Ray instance."
            )

        cluster_def = self._cluster_manager.load(cluster)
        if cluster_def is None:
            raise Exception(
                f"cluster definition '{cluster}' not found. "
                f"Create a new cluster with define_cluster()"
            )

        if cluster_def.cluster_type == GenevaClusterType.LOCAL_RAY:
            if manifest is not None:
                raise ValueError(
                    "Manifests are not supported with LOCAL_RAY cluster type"
                )
            return ray_cluster(
                local=True, log_to_driver=log_to_driver, logging_level=logging_level
            )

        ray_env = {
            "RAY_BACKEND_LOG_LEVEL": "info",
            "RAY_LOG_TO_DRIVER": "1",
            "RAY_ENABLE_RECORD_ACTOR_TASK_LOGGING": "0",
            "RAY_RUNTIME_ENV_LOG_TO_DRIVER_ENABLED": "true",
        }

        # load the manifest if provided
        manifest_def = None
        if manifest is not None:
            manifest_def = self._manifest_manager.load(manifest)
            if manifest_def is None:
                raise Exception(
                    f"manifest definition '{manifest}' not found. "
                    f"Create a new manifest with define_manifest()"
                )

        if cluster_def.cluster_type == GenevaClusterType.EXTERNAL_RAY:
            return ray_cluster(
                addr=cluster_def.ray_address,
                use_portforwarding=False,
                manifest=manifest_def,
                extra_env=ray_env,
                log_to_driver=log_to_driver,
                logging_level=logging_level,
                ray_init_kwargs=cluster_def.ray_init_kwargs,
            )

        use_portforwarding = cluster_def.as_dict()["kuberay"].get(
            "use_portforwarding", True
        )
        rc = cluster_def.to_ray_cluster()
        rc.on_exit = on_exit or ExitMode.DELETE
        if wait_timeout is not None:
            rc.wait_timeout = wait_timeout
        if manifest_def:
            # image explicitly provided in manifest takes precedence over cluster
            if img := manifest_def.head_image:
                _LOG.debug(f"overriding cluster head image from manifest: {img}")
                rc.head_group.image = img
            if img := manifest_def.worker_image:
                _LOG.debug(f"overriding cluster worker image from manifest: {img}")
                for wg in rc.worker_groups:
                    wg.image = img

        return ray_cluster(
            use_portforwarding=use_portforwarding,
            ray_cluster=rc,
            manifest=manifest_def,
            extra_env=ray_env,
            log_to_driver=log_to_driver,
            logging_level=logging_level,
        )

    def is_remote(self) -> bool:
        return self.uri.startswith("db://")


def sql(self, query: str) -> pa.Table:
    """Execute a raw SQL query.

    It uses the Flight SQL engine to execute the query.

    Parameters
    ----------
    query: str
        SQL query to execute

    Returns
    -------
    pyarrow.Table
        Result of the query in a `pyarrow.Table`

    """
    info = self.flight_client.execute(query)
    return self.flight_client.do_get(info.endpoints[0].ticket).read_all()


@attrs.define
class _GenavaConnectionConfig(ConfigBase):
    region: str = attrs.field(default="us-east-1")
    api_key: str | None = attrs.field(default=None)
    host_override: str | None = attrs.field(default=None)
    checkpoint: str | None = attrs.field(default=None)
    system_namespace: list[str] = attrs.field(factory=lambda: [SYSTEM_NAMESPACE])

    @classmethod
    @override
    def name(cls) -> str:
        return "connection"


def _ensure_system_namespace_exists(conn: Connection) -> None:
    """Ensure system_namespace exists, creating it if necessary.

    This is called once at connection time for namespace connections to ensure
    the system namespace is ready before any system tables are created.
    For nested namespaces, creates parent namespaces first.

    Args:
        conn: The Geneva connection

    Raises:
        RuntimeError: If the namespace doesn't exist and cannot be created
    """
    from lance_namespace import CreateNamespaceRequest, DescribeNamespaceRequest

    if not conn.system_namespace or conn._system_namespace_ensured:
        return

    # Remote db:// connections rely on the server to handle namespace creation
    # during table operations. Explicit namespace creation is not reliable across
    # phalanx/server versions and can fail before any useful work begins.
    if conn.namespace_impl is None and conn.is_remote():
        _LOG.info(
            "Deferring system namespace creation for remote connection: %s",
            conn.system_namespace,
        )
        conn._system_namespace_ensured = True
        return

    if conn.namespace_client is None:
        return

    try:
        # Check if namespace exists
        conn.namespace_client.describe_namespace(
            DescribeNamespaceRequest(id=conn.system_namespace)
        )
        _LOG.info(f"System namespace {conn.system_namespace} exists")
        conn._system_namespace_ensured = True
    except Exception as e:
        # Namespace doesn't exist, try to create it
        if "not found" in str(e).lower():
            try:
                # For nested namespaces, create parents first
                for i in range(1, len(conn.system_namespace) + 1):
                    parent_ns = conn.system_namespace[:i]
                    try:
                        conn.namespace_client.describe_namespace(
                            DescribeNamespaceRequest(id=parent_ns)
                        )
                        _LOG.debug(f"Namespace {parent_ns} already exists")
                    except Exception:
                        _LOG.info(f"Creating namespace {parent_ns}...")
                        conn.namespace_client.create_namespace(
                            CreateNamespaceRequest(id=parent_ns)
                        )
                        _LOG.info(f"Created namespace {parent_ns}")

                _LOG.info(f"System namespace {conn.system_namespace} is ready")
                conn._system_namespace_ensured = True
            except Exception as create_error:
                raise RuntimeError(
                    f"Cannot create system namespace {conn.system_namespace}: "
                    f"{create_error}"
                ) from create_error
        else:
            # Some other error, re-raise
            raise


def connect(
    uri: str | Path | None = None,
    *,
    region: str | None = None,
    api_key: Credential | str | None = None,
    host_override: str | None = None,
    storage_options: dict[str, str] | None = None,
    checkpoint: str | CheckpointStore | None = None,
    system_namespace: list[str] | None = None,
    namespace_impl: str | None = None,
    namespace_properties: dict[str, str] | None = None,
    upload_dir: str | None = None,
    initialize_upload_dir: bool = True,
    **kwargs,
) -> Connection:
    """Create a Geneva Connection to an existing database.

    Examples
    --------
        import geneva
        # Connect to a database in object storage
        conn = geneva.connect("s3://my-storage-bucket/my-database")
        # Connect using directory namespace
        conn = geneva.connect(
            namespace_impl="dir", namespace_properties={"root": "/path"}
        )
        # Connect using REST namespace
        conn = geneva.connect(
            namespace_impl="rest", namespace_properties={"uri": f"http://127.0.0.1:1234"}
        )
        # Connect with Phalanx remote server and separate upload bucket
        conn = geneva.connect(
            uri="db://my_database",
            api_key="my-api-key",
            host_override="https://phalanx.example.com",
            upload_dir="s3://my-upload-bucket/manifests",
        )
        tbl = conn.open_table("youtube_dataset")

    Parameters
    ----------
    uri: geneva URI, or Path, optional
        LanceDB Database URI, or a S3/GCS path.
        If not provided and namespace_impl is set, defaults to "namespace://".
    region: str | None
        LanceDB cloud region. Set to `None` on LanceDB Enterprise
    api_key: str | None
        Optional API key for enterprise endpoint
    host_override: str | None
        Optional host URI for enterprise endpoint
    system_namespace: list[str] | None
        Namespace for system tables (manifests, clusters, jobs, errors).
        Defaults to config value if not provided.
    namespace_impl: str | None
        The namespace implementation to use (e.g., "dir", "rest").
        If provided, connects using namespace instead of local database.
    namespace_properties: dict[str, str] | None
        Configuration properties for the namespace implementation.
    upload_dir: str | None
        Optional separate bucket/path for manifest uploads. This allows the client to
        upload manifests to a separate bucket instead of the data bucket
        where manifests are uploaded by default.
    Returns
    -------
    Connection - A LanceDB connection
    """

    api_key, checkpoint_store, host_override, region, uri, system_namespace = (
        _pre_connect(
            api_key,
            checkpoint,
            host_override,
            region,
            uri,
            namespace_impl,
            namespace_properties,
            system_namespace,
        )
    )

    if host_override:
        _LOG.info(f"Using host_override: {host_override}")

    conn = Connection(
        str(uri),
        region=region,
        api_key=api_key,
        host_override=host_override,
        storage_options=storage_options,
        checkpoint_store=checkpoint_store,
        namespace_impl=namespace_impl,
        namespace_properties=namespace_properties,
        system_namespace=system_namespace,
        upload_dir=upload_dir,
        **kwargs,
    )

    if initialize_upload_dir:
        _set_upload_dir(conn, upload_dir)

    # Validate and create system namespace for namespace_impl connections.
    # For remote db:// connections, defer creation until system tables are used.
    if namespace_impl is not None:
        _ensure_system_namespace_exists(conn)

    return conn


def resolve_table_physical_uri(
    table_id: list[str],
    *,
    db_uri: str | None,
    api_key: Credential | str | None = None,
    host_override: str | None = None,
    namespace_impl: str | None = None,
    namespace_properties: dict[str, str] | None = None,
) -> str:
    """Resolve a table reference to the backing physical storage URI."""
    table_name = table_id[-1]
    namespace = table_id[:-1]

    if namespace_impl is not None and namespace_properties is not None:
        from lance_namespace import DescribeTableRequest

        namespace_client = namespace_connect(namespace_impl, namespace_properties)
        response = namespace_client.describe_table(DescribeTableRequest(id=table_id))
        if response.location is None:
            raise ValueError(f"Table {table_id} does not have a physical location")
        return response.location

    if db_uri is not None and db_uri.startswith("db://"):
        remote_db = connect(
            db_uri,
            api_key=api_key,
            host_override=host_override,
            checkpoint="memory",
            initialize_upload_dir=False,
            system_namespace=infer_system_namespace(
                uri=db_uri,
                namespace_impl=namespace_impl,
                system_namespace=None,
            ),
            read_consistency_interval=timedelta(0),
        )
        try:
            table = remote_db.open_table(table_name, namespace=namespace)
            return table.uri  # pyright: ignore[reportAttributeAccessIssue]
        finally:
            remote_db.close()

    assert db_uri is not None, "db_uri must be set"
    return str(URL(str(db_uri)) / f"{table_name}.lance")


def _set_upload_dir(conn: Connection, upload_dir: str | None) -> None:
    # Set up default uploader if not already configured
    # This is needed for cluster operations (like upload_local_env) that don't have
    # table context
    from geneva.config import override_config_kv

    if upload_dir:
        override_config_kv({"uploader.upload_dir": upload_dir})
    else:
        # Remote and namespace-backed connections already have enough table
        # context to build an uploader lazily when needed. Probing the manifest
        # table here can create concurrent side effects during worker-side
        # TableReference reopens.
        if conn.is_remote() or conn.namespace_impl is not None:
            return
        try:
            Uploader.get()
        except (TypeError, ValueError):
            # Uploader not configured - try to derive a default upload_dir
            try:
                from geneva.manifest.mgr import ManifestConfigManager

                mgr = ManifestConfigManager(conn)
                # get the underlying object storage URI
                tbl_uri: str = mgr.table.uri  # pyright: ignore[reportAttributeAccessIssue]
                inner_uri = _get_db_uri(tbl_uri)

                default_upload_dir = f"{inner_uri}/{DEFAULT_UPLOAD_DIR}"
                override_config_kv({"uploader.upload_dir": default_upload_dir})
            except (AttributeError, TypeError, ValueError, OSError, RuntimeError):
                # Can't derive upload_dir - leave unconfigured
                # Uploads will fail later if attempted without explicit upload_dir
                pass


def _pre_connect(
    api_key: Credential | str | None,
    checkpoint: str | CheckpointStore | None,
    host_override: str | None,
    region: str | None,
    uri: str | Path | None,
    namespace_impl: str | None,
    namespace_properties: dict[str, str] | None,
    system_namespace: list[str] | None,
) -> tuple[
    Credential | None,
    CheckpointStore,
    str | None,
    str,
    str | Path,
    list[str],
]:
    # load values from config if not provided via arguments
    config = _GenavaConnectionConfig.get()
    region = region or config.region or "us-east-1"
    api_key = api_key or config.api_key
    api_key = Credential(api_key) if isinstance(api_key, str) else api_key
    host_override = host_override or config.host_override
    if system_namespace is None:
        system_namespace = config.system_namespace
    # handle local relative paths
    is_local = isinstance(uri, Path) or (
        uri is not None and get_uri_scheme(uri) == "file"
    )
    if is_local and uri:
        if isinstance(uri, str):
            uri = Path(uri)
        uri = uri.expanduser().absolute()
        uri.mkdir(parents=True, exist_ok=True)

    # Default URI for namespace connections
    if uri is None:
        uri = "namespace://"

    if checkpoint is None:
        checkpoint = str(URL(str(uri)) / "ckp")
    if isinstance(checkpoint, str):
        checkpoint_store = CheckpointStore.from_uri(checkpoint)
    else:
        checkpoint_store = checkpoint

    _LOG.debug(f"using checkpoint store: {type(checkpoint_store)}")

    return api_key, checkpoint_store, host_override, region, uri, system_namespace


def _get_db_uri(tbl_uri: str) -> str:
    """Get the database URI from a table URI.

    For example,
        s3://bucket/path/to/table.lance -> s3://bucket/path/to
        db://path/to/table.lance -> db://path/to
        /tmp/foo/table.lance -> /tmp/foo
        s3+ddb://bucket/path/table.lance?ddbTableName=x -> s3+ddb://bucket/path?ddbTableName=x
    """
    parsed = urlparse(tbl_uri)
    path_parts = parsed.path.rsplit("/", 1)
    db_path = path_parts[0] if len(path_parts) > 1 else ""

    # Handle local paths (no scheme)
    if not parsed.scheme:
        return db_path

    base = f"{parsed.scheme}://{parsed.netloc}{db_path}"
    if parsed.query:
        base = f"{base}?{parsed.query}"
    return base
