FabledGrimoire/src/scribe/services/milestones.py

"""Milestone management service."""
import logging
from datetime import datetime, timezone

from sqlalchemy import func, select

from scribe.models import async_session
from scribe.models.milestone import Milestone
from scribe.models.note import Note

logger = logging.getLogger(__name__)


async def create_milestone(
    user_id: int,
    project_id: int,
    title: str,
    description: str | None = None,
    body: str | None = None,
    order_index: int = 0,
    status: str = "active",
) -> Milestone:
    async with async_session() as session:
        milestone = Milestone(
            user_id=user_id,
            project_id=project_id,
            title=title,
            description=description,
            body=body,
            status=status,
            order_index=order_index,
        )
        session.add(milestone)
        await session.commit()
        await session.refresh(milestone)
        return milestone


async def get_milestone(user_id: int, milestone_id: int) -> Milestone | None:
    async with async_session() as session:
        result = await session.execute(
            select(Milestone).where(
                Milestone.id == milestone_id, Milestone.user_id == user_id,
                Milestone.deleted_at.is_(None),
            )
        )
        return result.scalars().first()


async def get_milestone_in_project(project_id: int, milestone_id: int) -> Milestone | None:
    """Fetch a milestone by id within a project, without a user_id ownership check.
    Callers must verify project access separately before using this."""
    async with async_session() as session:
        result = await session.execute(
            select(Milestone).where(
                Milestone.id == milestone_id,
                Milestone.project_id == project_id,
                Milestone.deleted_at.is_(None),
            )
        )
        return result.scalars().first()


async def get_milestone_by_title(user_id: int, project_id: int, title: str) -> Milestone | None:
    async with async_session() as session:
        result = await session.execute(
            select(Milestone).where(
                Milestone.user_id == user_id,
                Milestone.project_id == project_id,
                func.lower(Milestone.title) == func.lower(title.strip()),
            ).limit(1)
        )
        return result.scalars().first()


async def find_milestone_by_title(user_id: int, title: str) -> Milestone | None:
    """Find a milestone by title across ALL projects for this user (case-insensitive)."""
    async with async_session() as session:
        result = await session.execute(
            select(Milestone).where(
                Milestone.user_id == user_id,
                func.lower(Milestone.title) == func.lower(title.strip()),
            ).order_by(Milestone.id).limit(1)
        )
        return result.scalars().first()


async def get_or_create_milestone(user_id: int, project_id: int, title: str) -> Milestone:
    milestone = await get_milestone_by_title(user_id, project_id, title)
    if milestone:
        return milestone
    return await create_milestone(user_id, project_id, title=title)


async def list_milestones(
    user_id: int, project_id: int, status: str | None = None
) -> list[Milestone]:
    async with async_session() as session:
        query = select(Milestone).where(
            Milestone.user_id == user_id,
            Milestone.project_id == project_id,
            Milestone.deleted_at.is_(None),
        )
        if status:
            query = query.where(Milestone.status == status)
        query = query.order_by(Milestone.order_index.asc(), Milestone.created_at.asc())
        result = await session.execute(query)
        return list(result.scalars().all())


async def update_milestone(user_id: int, milestone_id: int, **fields: object) -> Milestone | None:
    async with async_session() as session:
        result = await session.execute(
            select(Milestone).where(
                Milestone.id == milestone_id, Milestone.user_id == user_id,
                Milestone.deleted_at.is_(None),
            )
        )
        milestone = result.scalars().first()
        if milestone is None:
            return None
        for key, value in fields.items():
            if hasattr(milestone, key):
                setattr(milestone, key, value)
        milestone.updated_at = datetime.now(timezone.utc)
        await session.commit()
        await session.refresh(milestone)
        return milestone


async def delete_milestone(user_id: int, milestone_id: int) -> bool:
    async with async_session() as session:
        result = await session.execute(
            select(Milestone).where(Milestone.id == milestone_id, Milestone.user_id == user_id)
        )
        milestone = result.scalars().first()
        if milestone is None:
            return False
        await session.delete(milestone)
        await session.commit()
        return True


async def get_milestone_progress(milestone_id: int) -> dict:
    """Return task completion stats for a milestone."""
    async with async_session() as session:
        rows = await session.execute(
            select(Note.status, func.count(Note.id))
            .where(
                Note.milestone_id == milestone_id, Note.status.isnot(None),
                Note.deleted_at.is_(None),
            )
            .group_by(Note.status)
        )
        status_counts: dict[str, int] = {}
        for status, count in rows.fetchall():
            status_counts[status] = count

    # Same rule as the batch path, computed in one place so the two cannot
    # drift on the cancelled-exclusion.
    return _progress_from_counts(status_counts)


def _progress_from_counts(status_counts: dict[str, int]) -> dict:
    """The progress shape, computed from already-fetched counts.

    Split out of get_milestone_progress so the batch path can reuse the rule
    rather than restate it — the cancelled-exclusion below is easy to get
    subtly different in a second copy, and then two screens disagree about
    whether a milestone is finished.
    """
    total = sum(status_counts.values())
    cancelled = status_counts.get("cancelled", 0)
    completed = status_counts.get("done", 0)
    # Cancelled tasks are resolved work, not pending — excluded from the
    # denominator so a milestone whose only open task was cancelled reaches
    # 100% instead of stalling.
    active_total = total - cancelled
    return {
        "total": total,
        "completed": completed,
        "pct": round(completed / active_total * 100, 1) if active_total > 0 else 0.0,
        "status_counts": {
            "todo": status_counts.get("todo", 0),
            "in_progress": status_counts.get("in_progress", 0),
            "done": status_counts.get("done", 0),
            "cancelled": cancelled,
        },
    }


async def get_project_milestone_summaries(
    user_id: int, project_ids: list[int]
) -> dict[int, list[dict]]:
    """Milestone summaries for MANY projects in two queries total.

    The per-project version below is a nested fan-out: one query to list a
    project's milestones, then one more per milestone for its progress. Called
    for 25 projects concurrently it asked for ~250 pooled connections against a
    pool of 15, and every one of them waited out the 30-second checkout timeout
    (#2384). This does the same work in two queries and one session.
    """
    if not project_ids:
        return {}

    async with async_session() as session:
        milestones = list((await session.execute(
            select(Milestone).where(
                Milestone.user_id == user_id,
                Milestone.project_id.in_(project_ids),
                Milestone.deleted_at.is_(None),
            ).order_by(Milestone.order_index.asc(), Milestone.created_at.asc())
        )).scalars().all())

        counts: dict[int, dict[str, int]] = {}
        if milestones:
            rows = await session.execute(
                select(Note.milestone_id, Note.status, func.count(Note.id))
                .where(
                    Note.milestone_id.in_([m.id for m in milestones]),
                    Note.status.isnot(None),
                    Note.deleted_at.is_(None),
                )
                .group_by(Note.milestone_id, Note.status)
            )
            for milestone_id, status, count in rows.fetchall():
                counts.setdefault(milestone_id, {})[status] = count

    out: dict[int, list[dict]] = {pid: [] for pid in project_ids}
    for m in milestones:
        entry = m.to_dict()
        entry.update(_progress_from_counts(counts.get(m.id, {})))
        out.setdefault(m.project_id, []).append(entry)
    return out


async def get_project_milestone_summary(user_id: int, project_id: int) -> list[dict]:
    """Ordered milestones with progress — the one-project view of
    get_project_milestone_summaries (two queries, not N+1)."""
    return (await get_project_milestone_summaries(user_id, [project_id])).get(project_id, [])


# What a milestone LISTING needs: enough to say what each plan is and how far
# along it is. The plan itself (`body`) is get_milestone's job. Summaries once
# carried it, and on a project with 39 milestones enter_project came to ~222k
# characters, 168k of them plan bodies. That is past what an MCP client will
# accept as a tool result, so the session handshake arrived as a file to page
# through (#4045). user_id / project_id / timestamps repeat what the caller
# already knows.
_BRIEF_FIELDS = (
    "id", "title", "description", "status", "order_index",
    "total", "completed", "pct", "status_counts",
)


def brief_milestone_summary(
    rows: list[dict], *, done_kept: int | None = None,
) -> tuple[list[dict], int]:
    """Trim summary rows to the listing fields, optionally capping done ones.

    `done_kept` keeps only the N most recently updated done milestones (every
    open one stays), in the original order_index order. A done plan is
    history, and a handshake that grows with a project's whole history grows
    past what a client accepts. None keeps every row. Returns (rows, omitted),
    where `omitted` counts the done milestones left out.
    """
    if done_kept is None:
        kept = rows
    else:
        done = [r for r in rows if r.get("status") == "done"]
        recent = sorted(done, key=lambda r: r.get("updated_at") or "", reverse=True)
        keep_ids = {r.get("id") for r in recent[:done_kept]}
        kept = [r for r in rows if r.get("status") != "done" or r.get("id") in keep_ids]
    brief = [{k: r[k] for k in _BRIEF_FIELDS if k in r} for r in kept]
    return brief, len(rows) - len(kept)