"""
Agent context loading.

Loads the active agent for a tenant before message processing. Phase 1 supports
a single agent per tenant. The RLS context variable must be set before calling
any function here so that PostgreSQL RLS filters correctly.
"""

from __future__ import annotations

import logging
import uuid

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from shared.models.tenant import Agent
from shared.rls import current_tenant_id

logger = logging.getLogger(__name__)


async def load_agent_for_tenant(
    tenant_id: str,
    session: AsyncSession,
) -> Agent | None:
    """
    Load the active agent for a tenant.

    Sets the ``current_tenant_id`` ContextVar so that PostgreSQL RLS policies
    correctly filter the agents table to only return rows belonging to this
    tenant.

    Phase 1: Returns the first active agent for the tenant (single-agent model).
    Phase 2+: Will support agent selection based on message content and routing
    rules.

    Args:
        tenant_id:  Konstruct tenant ID as a UUID string.
        session:    Async SQLAlchemy session.

    Returns:
        The active Agent ORM instance, or None if no active agent is configured.
    """
    try:
        tenant_uuid = uuid.UUID(tenant_id)
    except (ValueError, AttributeError):
        logger.error("load_agent_for_tenant: invalid tenant_id=%r", tenant_id)
        return None

    # Set RLS context so the DB query is correctly scoped to this tenant
    token = current_tenant_id.set(tenant_uuid)
    try:
        stmt = (
            select(Agent)
            .where(Agent.tenant_id == tenant_uuid)
            .where(Agent.is_active.is_(True))
            .limit(1)
        )
        result = await session.execute(stmt)
        agent = result.scalars().first()
    except Exception:
        logger.exception(
            "load_agent_for_tenant: DB error for tenant=%s", tenant_id
        )
        return None
    finally:
        # Always reset the RLS context var after DB work completes
        current_tenant_id.reset(token)

    if agent is None:
        logger.warning(
            "load_agent_for_tenant: no active agent for tenant=%s", tenant_id
        )

    return agent
