"""Write assistant rows to ``run_chat_messages`` from worker context."""
from __future__ import annotations
from typing import Any, Optional
from uuid import UUID

from sqlalchemy.ext.asyncio import AsyncSession


async def resolve_session_id_for_run(db_factory, run) -> Optional[UUID]:
    if not run.prompt_id:
        return None
    from src.domain.instructions.repository import PromptRepository

    async with db_factory() as db:
        p = await PromptRepository(db).get_by_id(run.prompt_id)
        return p.session_id if p else None


async def persist_assistant_message(
    db_factory,
    *,
    project_id: UUID,
    session_id: UUID,
    run_id: UUID,
    message_type: str,
    content: str,
    payload: Optional[dict[str, Any]] = None,
) -> None:
    from src.domain.codegen.run_chat_repository import RunChatMessageRepository

    async with db_factory() as db:
        await RunChatMessageRepository(db).create(
            project_id=project_id,
            session_id=session_id,
            run_id=run_id,
            role="assistant",
            message_type=message_type,
            content=content,
            payload=payload,
        )
        await db.commit()


async def persist_assistant_for_run(
    db_factory,
    run,
    *,
    message_type: str,
    content: str,
    payload: Optional[dict[str, Any]] = None,
) -> None:
    sid = await resolve_session_id_for_run(db_factory, run)
    if not sid:
        return
    await persist_assistant_message(
        db_factory,
        project_id=run.project_id,
        session_id=sid,
        run_id=run.id,
        message_type=message_type,
        content=content,
        payload=payload,
    )


async def persist_assistant_for_run_session(
    db: AsyncSession,
    run,
    *,
    message_type: str,
    content: str,
    payload: Optional[dict[str, Any]] = None,
) -> None:
    """Append assistant row using an existing worker DB session (caller commits)."""
    if not run.prompt_id:
        return
    from src.domain.instructions.repository import PromptRepository
    from src.domain.codegen.run_chat_repository import RunChatMessageRepository

    p = await PromptRepository(db).get_by_id(run.prompt_id)
    if not p:
        return
    await RunChatMessageRepository(db).create(
        project_id=run.project_id,
        session_id=p.session_id,
        run_id=run.id,
        role="assistant",
        message_type=message_type,
        content=content or "",
        payload=payload,
    )