from __future__ import annotations

from collections.abc import Mapping, Sequence
from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import lancedb

_connections: dict[str, lancedb.DBConnection] = {}


def get_db(workspace: Path) -> lancedb.DBConnection:
    import lancedb

    key = str(workspace)
    if key not in _connections:
        db_path = workspace / "lancedb"
        db_path.mkdir(parents=True, exist_ok=True)
        _connections[key] = lancedb.connect(str(db_path))
    return _connections[key]


def close_db(workspace: Path) -> None:
    key = str(workspace)
    db = _connections.pop(key, None)
    if db is None:
        return
    close = getattr(db, "close", None)
    if callable(close):
        close()


def ensure_table(
    db: lancedb.DBConnection, name: str, sample: Sequence[Mapping[str, object]]
) -> lancedb.table.Table:
    table, _created = open_or_create_table(db, name, sample)
    return table


def open_or_create_table(
    db: lancedb.DBConnection, name: str, sample: Sequence[Mapping[str, object]]
) -> tuple[lancedb.table.Table, bool]:
    try:
        return db.open_table(name), False
    except Exception:
        return db.create_table(name, data=sample), True
