state-hub/api/routers/messages.py

from datetime import datetime, timezone

from fastapi import APIRouter, Depends, HTTPException, status
from sqlalchemy import or_, select
from sqlalchemy.ext.asyncio import AsyncSession

from api.database import get_session
from api.models.agent_message import AgentMessage
from api.schemas.agent_message import MessageCreate, MessageRead, MessageReply
from api.services.repository_aliases import (
    canonicalize_repository_slug,
    resolve_repository_slug,
)
from hub_core.message_identity import resolve_message_reference
from hub_core.models.message_identity_alias import MessageIdentityAlias

router = APIRouter(prefix="/messages", tags=["messages"])


async def _get_message(reference: str, session: AsyncSession) -> AgentMessage:
    message_id = await resolve_message_reference(
        session, reference, alias_model=MessageIdentityAlias
    )
    if message_id is None:
        raise HTTPException(
            status_code=404, detail=f"Message reference {reference!r} not found"
        )
    message = await session.get(AgentMessage, message_id)
    if message is None:
        raise HTTPException(
            status_code=404, detail=f"Message reference {reference!r} not found"
        )
    return message


@router.post("/", response_model=MessageRead, status_code=status.HTTP_201_CREATED)
async def send_message(
    body: MessageCreate,
    session: AsyncSession = Depends(get_session),
) -> AgentMessage:
    if body.thread_id and await session.get(AgentMessage, body.thread_id) is None:
        raise HTTPException(status_code=404, detail=f"Thread root {body.thread_id} not found")
    payload = body.model_dump()
    payload["from_agent"] = await canonicalize_repository_slug(session, body.from_agent)
    payload["to_agent"] = await canonicalize_repository_slug(session, body.to_agent)
    message = AgentMessage(**payload)
    session.add(message)
    await session.commit()
    await session.refresh(message)
    return message


@router.get("/", response_model=list[MessageRead])
async def list_messages(
    to_agent: str | None = None,
    from_agent: str | None = None,
    unread_only: bool = False,
    limit: int = 50,
    session: AsyncSession = Depends(get_session),
) -> list[AgentMessage]:
    query = select(AgentMessage).where(AgentMessage.archived_at.is_(None))
    if to_agent:
        resolution = await resolve_repository_slug(session, to_agent, required=False)
        values = resolution.slug_values if resolution else (to_agent,)
        query = query.where(
            or_(AgentMessage.to_agent.in_(values), AgentMessage.to_agent == "broadcast")
        )
    if from_agent:
        resolution = await resolve_repository_slug(session, from_agent, required=False)
        values = resolution.slug_values if resolution else (from_agent,)
        query = query.where(AgentMessage.from_agent.in_(values))
    if unread_only:
        query = query.where(AgentMessage.read_at.is_(None))
    result = await session.execute(
        query.order_by(AgentMessage.created_at.desc()).limit(limit)
    )
    return list(result.scalars().all())


@router.get("/thread/{thread_id}", response_model=list[MessageRead])
async def get_thread(
    thread_id: str,
    session: AsyncSession = Depends(get_session),
) -> list[AgentMessage]:
    resolved = await resolve_message_reference(
        session, thread_id, alias_model=MessageIdentityAlias
    )
    if resolved is None:
        raise HTTPException(
            status_code=404, detail=f"Message reference {thread_id!r} not found"
        )
    result = await session.execute(
        select(AgentMessage)
        .where(or_(AgentMessage.id == resolved, AgentMessage.thread_id == resolved))
        .order_by(AgentMessage.created_at)
    )
    return list(result.scalars().all())


@router.patch("/{message_id}/read", response_model=MessageRead)
async def mark_read(
    message_id: str,
    session: AsyncSession = Depends(get_session),
) -> AgentMessage:
    message = await _get_message(message_id, session)
    if message.read_at is None:
        message.read_at = datetime.now(timezone.utc)
        await session.commit()
        await session.refresh(message)
    return message


@router.patch("/{message_id}/archive", response_model=MessageRead)
async def archive_message(
    message_id: str,
    session: AsyncSession = Depends(get_session),
) -> AgentMessage:
    message = await _get_message(message_id, session)
    message.archived_at = datetime.now(timezone.utc)
    if message.read_at is None:
        message.read_at = message.archived_at
    await session.commit()
    await session.refresh(message)
    return message


@router.post(
    "/{message_id}/reply",
    response_model=MessageRead,
    status_code=status.HTTP_201_CREATED,
)
async def reply_to_message(
    message_id: str,
    body: MessageReply,
    session: AsyncSession = Depends(get_session),
) -> AgentMessage:
    original = await _get_message(message_id, session)
    if original.read_at is None:
        original.read_at = datetime.now(timezone.utc)
    reply = AgentMessage(
        from_agent=await canonicalize_repository_slug(session, body.from_agent),
        to_agent=await canonicalize_repository_slug(session, original.from_agent),
        subject=f"Re: {original.subject}",
        body=body.body,
        thread_id=original.thread_id or original.id,
    )
    session.add(reply)
    await session.commit()
    await session.refresh(reply)
    return reply


__all__ = ["router"]