"""DB helper utilities: upsert, bulk operations, health check."""
from __future__ import annotations

from typing import Any

from sqlalchemy import select, text
from sqlalchemy.dialects.postgresql import insert
from sqlalchemy.ext.asyncio import AsyncSession

from libs.db.models import Base


async def upsert(
    session: AsyncSession,
    model: type[Base],
    values: dict[str, Any],
    index_elements: list[str],
    update_columns: list[str] | None = None,
) -> None:
    """Insert or update a single row using PostgreSQL ON CONFLICT."""
    stmt = insert(model).values(**values)
    if update_columns:
        update_dict = {col: getattr(stmt.excluded, col) for col in update_columns}
        stmt = stmt.on_conflict_do_update(
            index_elements=index_elements, set_=update_dict
        )
    else:
        stmt = stmt.on_conflict_do_nothing(index_elements=index_elements)
    await session.execute(stmt)


async def bulk_upsert(
    session: AsyncSession,
    model: type[Base],
    rows: list[dict[str, Any]],
    index_elements: list[str],
    update_columns: list[str] | None = None,
) -> int:
    """Bulk upsert; returns number of rows processed."""
    if not rows:
        return 0
    stmt = insert(model).values(rows)
    if update_columns:
        update_dict = {col: getattr(stmt.excluded, col) for col in update_columns}
        stmt = stmt.on_conflict_do_update(
            index_elements=index_elements, set_=update_dict
        )
    else:
        stmt = stmt.on_conflict_do_nothing(index_elements=index_elements)
    await session.execute(stmt)
    return len(rows)


async def get_or_create(
    session: AsyncSession,
    model: type[Base],
    pk_value: Any,
    pk_column: str,
    defaults: dict[str, Any],
) -> tuple[Any, bool]:
    """Return (instance, created). Upserts if not found."""
    col = getattr(model, pk_column)
    result = await session.execute(select(model).where(col == pk_value))
    row = result.scalar_one_or_none()
    if row is not None:
        return row, False
    instance = model(**{pk_column: pk_value, **defaults})
    session.add(instance)
    await session.flush()
    return instance, True


async def health_check(session: AsyncSession) -> bool:
    """Return True if DB is reachable."""
    try:
        await session.execute(text("SELECT 1"))
        return True
    except Exception:
        return False