"""Entity CRUD helpers — create, update, and query MemoryEntity records.

All functions are async and use SQLAlchemy with PostgreSQL ON CONFLICT upsert.
"""

from __future__ import annotations

import logging
from datetime import UTC, datetime

from sqlalchemy import select
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.ext.asyncio import AsyncSession

from arandu.models import MemoryEntity

logger = logging.getLogger(__name__)

__all__ = [
    "create_or_update_entity",
    "get_entities_for_user",
    "get_entity_by_key",
]


async def create_or_update_entity(
    session: AsyncSession,
    user_id: str,
    canonical_key: str,
    display_name: str | None = None,
    entity_type: str = "other",
) -> MemoryEntity:
    """Create a MemoryEntity or update if it exists (ON CONFLICT upsert).

    Args:
        session: Database session.
        user_id: User identifier.
        canonical_key: Canonical entity key.
        display_name: Optional display name.
        entity_type: Entity type (person, pet, place, etc.).

    Returns:
        The created or updated MemoryEntity.
    """
    now = datetime.now(UTC)

    stmt = pg_insert(MemoryEntity).values(
        user_id=user_id,
        canonical_key=canonical_key,
        display_name=display_name,
        entity_type=entity_type,
        first_seen_at=now,
        last_seen_at=now,
        fact_count=0,
        is_active=True,
    )

    update_values: dict = {
        "last_seen_at": now,
        "fact_count": MemoryEntity.fact_count + 1,
    }
    if display_name is not None:
        update_values["display_name"] = display_name

    stmt = stmt.on_conflict_do_update(  # type: ignore[assignment]
        constraint="uq_memory_entity_user_key",
        set_=update_values,
    ).returning(MemoryEntity)

    result = await session.execute(stmt)
    entity: MemoryEntity = result.scalar_one()
    logger.debug(
        "entity_upsert: user=%s key=%s type=%s",
        user_id,
        canonical_key,
        entity_type,
    )
    return entity


async def get_entity_by_key(
    session: AsyncSession,
    user_id: str,
    canonical_key: str,
) -> MemoryEntity | None:
    """Get a single MemoryEntity by user_id and canonical_key.

    Args:
        session: Database session.
        user_id: User identifier.
        canonical_key: Entity key to look up.

    Returns:
        MemoryEntity or None if not found.
    """
    stmt = select(MemoryEntity).where(
        MemoryEntity.user_id == user_id,
        MemoryEntity.canonical_key == canonical_key,
    )
    result = await session.execute(stmt)
    return result.scalar_one_or_none()


async def get_entities_for_user(
    session: AsyncSession,
    user_id: str,
    active_only: bool = True,
) -> list[MemoryEntity]:
    """List all MemoryEntity records for a user.

    Args:
        session: Database session.
        user_id: User identifier.
        active_only: If True, only return active entities.

    Returns:
        List of MemoryEntity records.
    """
    stmt = select(MemoryEntity).where(MemoryEntity.user_id == user_id)
    if active_only:
        stmt = stmt.where(MemoryEntity.is_active.is_(True))
    stmt = stmt.order_by(MemoryEntity.last_seen_at.desc())
    result = await session.execute(stmt)
    return list(result.scalars().all())
