"""同一 task 内消息序号的统一事务分配入口。"""
from __future__ import annotations

from typing import Optional
from uuid import UUID

from sqlalchemy import func, select
from sqlalchemy.orm import Session as OrmSession

from .models import Message, Task


def allocate_message_idx(
    session: OrmSession,
    task_id: UUID,
    *,
    locked_task: Optional[Task] = None,
) -> int:
    """锁定 task 行并返回唯一递增 idx；max 校准兼容蓝绿旧实例。"""
    task = locked_task
    if task is None:
        task = session.execute(
            select(Task).where(Task.task_id == task_id).with_for_update()
        ).scalar_one()
    elif task.task_id != task_id:
        raise ValueError("locked_task 与 task_id 不一致")

    max_idx = session.execute(
        select(func.max(Message.idx)).where(Message.task_id == task_id)
    ).scalar_one()
    persisted_next = int(getattr(task, "next_message_idx", 0) or 0)
    observed_next = -1 if max_idx is None else int(max_idx)
    next_idx = max(persisted_next, observed_next + 1)
    task.next_message_idx = next_idx + 1
    return next_idx