src/crud/peer.py

from logging import getLogger
from typing import Any

from cashews import NOT_NONE
from sqlalchemy import Select, select
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import make_transient_to_detached

from src import models, schemas
from src.cache.client import cache, get_cache_namespace, safe_cache_delete
from src.config import settings
from src.crud.workspace import get_or_create_workspace
from src.exceptions import ConflictException, ResourceNotFoundException
from src.models import Peer
from src.utils.filter import apply_filter
from src.utils.types import GetOrCreateResult

logger = getLogger(__name__)

PEER_CACHE_KEY_TEMPLATE = "v2:workspace:{workspace_name}:peer:{peer_name}"
PEER_LOCK_PREFIX = f"{get_cache_namespace()}:lock:v2"


def peer_cache_key(workspace_name: str, peer_name: str) -> str:
    """Generate cache key for peer."""
    return (
        get_cache_namespace()
        + ":"
        + PEER_CACHE_KEY_TEMPLATE.format(
            workspace_name=workspace_name,
            peer_name=peer_name,
        )
    )


async def get_or_create_peers(
    db: AsyncSession,
    workspace_name: str,
    peers: list[schemas.PeerCreate],
    *,
    _retry: bool = False,
) -> GetOrCreateResult[list[models.Peer]]:
    """
    Get an existing list of peers or create new peers if they don't exist.
    Updates existing peers with metadata and configuration if provided.

    Args:
        db: Database session
        workspace_name: Name of the workspace
        peers: List of peer creation schemas
        _retry: Whether to retry the operation

    Returns:
        GetOrCreateResult containing the list of peers and whether any were created

    Raises:
        ConflictException: If we fail to get or create the peers
    """

    await get_or_create_workspace(db, schemas.WorkspaceCreate(name=workspace_name))
    peer_names = [p.name for p in peers]
    stmt = (
        select(models.Peer)
        .where(models.Peer.workspace_name == workspace_name)
        .where(models.Peer.name.in_(peer_names))
    )
    result = await db.execute(stmt)
    existing_peers: list[Peer] = list(result.scalars().all())

    # Create a mapping of peer names to peer schemas for easy lookup
    peer_schema_map = {p.name: p for p in peers}

    # Track which peers actually changed
    changed_peers: list[Peer] = []

    # Update existing peers with metadata and configuration if provided
    for existing_peer in existing_peers:
        peer_schema = peer_schema_map[existing_peer.name]
        changed = False

        # Update with metadata if provided AND different
        if (
            peer_schema.metadata is not None
            and existing_peer.h_metadata != peer_schema.metadata
        ):
            existing_peer.h_metadata = peer_schema.metadata
            changed = True

        # Update with configuration if provided AND different
        if (
            peer_schema.configuration is not None
            and existing_peer.configuration != peer_schema.configuration
        ):
            existing_peer.configuration = peer_schema.configuration
            changed = True

        if changed:
            changed_peers.append(existing_peer)

    # Find which peers need to be created
    existing_names = {p.name for p in existing_peers}
    peers_to_create = [p for p in peers if p.name not in existing_names]

    # Create new peers
    new_peers = [
        models.Peer(
            workspace_name=workspace_name,
            name=p.name,
            h_metadata=p.metadata or {},
            configuration=p.configuration or {},
        )
        for p in peers_to_create
    ]
    try:
        async with db.begin_nested():
            db.add_all(new_peers)
    except IntegrityError:
        if _retry:
            raise ConflictException(
                f"Unable to create or get peers: {peer_names}"
            ) from None
        return await get_or_create_peers(db, workspace_name, peers, _retry=True)

    # Capture peer names eagerly so the closure holds plain strings, not ORM objects
    _cache_keys_to_invalidate = [
        peer_cache_key(workspace_name, p.name) for p in changed_peers + new_peers
    ]

    async def _invalidate_peer_cache():
        for cache_key in _cache_keys_to_invalidate:
            await safe_cache_delete(cache_key)

    # Return combined list of existing and new peers
    # created=True if any new peers were created
    return GetOrCreateResult(
        existing_peers + new_peers,
        created=len(new_peers) > 0,
        on_commit=_invalidate_peer_cache if _cache_keys_to_invalidate else None,
    )


@cache(
    key=PEER_CACHE_KEY_TEMPLATE,
    ttl=f"{settings.CACHE.DEFAULT_TTL_SECONDS}s",
    prefix=get_cache_namespace(),
    condition=NOT_NONE,
)
@cache.locked(
    key=PEER_CACHE_KEY_TEMPLATE,
    ttl=f"{settings.CACHE.DEFAULT_LOCK_TTL_SECONDS}s",
    prefix=PEER_LOCK_PREFIX,
)
async def _fetch_peer(
    db: AsyncSession,
    workspace_name: str,
    peer_name: str,
) -> dict[str, Any] | None:
    """Fetch a peer from the database and return as a plain dict for safe caching."""
    obj = await db.scalar(
        select(models.Peer)
        .where(models.Peer.workspace_name == workspace_name)
        .where(models.Peer.name == peer_name)
    )
    if obj is None:
        return None
    return {
        "id": obj.id,
        "name": obj.name,
        "workspace_name": obj.workspace_name,
        "h_metadata": obj.h_metadata,
        "internal_metadata": obj.internal_metadata,
        "configuration": obj.configuration,
        "created_at": obj.created_at,
    }


async def get_peer(
    db: AsyncSession,
    workspace_name: str,
    peer: schemas.PeerCreate,
) -> models.Peer:
    """
    Get an existing peer.

    Args:
        db: Database session
        workspace_name: Name of the workspace
        peer: Peer creation schema

    Returns:
        The peer if found

    Raises:
        ResourceNotFoundException: If the peer does not exist
    """
    data = await _fetch_peer(db, workspace_name, peer.name)
    if data is None:
        raise ResourceNotFoundException(
            f"Peer {peer.name} not found in workspace {workspace_name}"
        )

    # Reconstruct ORM object from cached dict and merge into session
    obj = models.Peer(**data)
    make_transient_to_detached(obj)
    existing_peer = await db.merge(obj, load=False)

    return existing_peer


async def get_peers(
    workspace_name: str,
    filters: dict[str, str] | None = None,
) -> Select[tuple[models.Peer]]:
    stmt = select(models.Peer).where(models.Peer.workspace_name == workspace_name)

    stmt = apply_filter(stmt, models.Peer, filters)

    return stmt.order_by(models.Peer.created_at)


async def update_peer(
    db: AsyncSession, workspace_name: str, peer_name: str, peer: schemas.PeerUpdate
) -> models.Peer:
    """
    Get or create a peer, then apply metadata and configuration updates.

    If the peer does not exist, the workspace and peer are created first.
    Provided metadata and configuration replace the existing values when
    present.

    Args:
        db: Database session
        workspace_name: Name of the workspace
        peer_name: Name of the peer
        peer: Peer update schema

    Returns:
        The updated peer

    Raises:
        ConflictException: If concurrent creation prevents fetching or creating
            the peer
    """
    peers_result = await get_or_create_peers(
        db, workspace_name, [schemas.PeerCreate(name=peer_name)]
    )
    honcho_peer = peers_result.resource[0]

    needs_update = False

    if peer.metadata is not None and honcho_peer.h_metadata != peer.metadata:
        honcho_peer.h_metadata = peer.metadata
        needs_update = True

    if (
        peer.configuration is not None
        and honcho_peer.configuration != peer.configuration
    ):
        honcho_peer.configuration = peer.configuration
        needs_update = True

    # Early exit if unchanged
    if not needs_update:
        await db.commit()
        await peers_result.post_commit()
        logger.debug(
            "Peer %s unchanged in workspace %s, skipping update",
            peer_name,
            workspace_name,
        )
        return honcho_peer

    await db.commit()
    await peers_result.post_commit()

    cache_key = peer_cache_key(workspace_name, honcho_peer.name)
    await safe_cache_delete(cache_key)

    logger.debug("Peer %s updated successfully", peer_name)
    return honcho_peer


async def get_sessions_for_peer(
    workspace_name: str,
    peer_name: str,
    filters: dict[str, Any] | None = None,
) -> Select[tuple[models.Session]]:
    """
    Get all sessions for a peer through the session_peers relationship.

    Args:
        workspace_name: Name of the workspace
        peer_name: Name of the peer
        filters: Filter sessions by metadata

    Returns:
        SQLAlchemy Select statement
    """
    stmt = (
        select(models.Session)
        .join(
            models.SessionPeer,
            (models.Session.name == models.SessionPeer.session_name)
            & (models.Session.workspace_name == models.SessionPeer.workspace_name),
        )
        .where(models.SessionPeer.peer_name == peer_name)
        .where(models.Session.workspace_name == workspace_name)
    )

    stmt = apply_filter(stmt, models.Session, filters)

    stmt: Select[tuple[models.Session]] = stmt.order_by(models.Session.created_at)

    return stmt