FabledGrimoire/src/grimoire/services/milestones.py

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

from sqlalchemy import func, select

from grimoire.models import async_session
from grimoire.models.milestone import Milestone
from grimoire.models.note import Note
from grimoire.services import access as access_svc

logger = logging.getLogger(__name__)

# The statuses that make a step OPEN work, defined here because this is the
# lower layer: services/placement.py imports it rather than restating it. Two
# surfaces that both answer "what is next" and disagree about what counts as
# next is worse than one of them staying silent.
OPEN_STEP_STATUSES = ("todo", "in_progress")


def embed_milestone(milestone: Milestone) -> None:
    """Refresh a milestone's vectors, fire-and-forget.

    At the service, so every path that writes a milestone gets it: a record
    written through a door that forgot the call stays out of search until a
    restart. Exceptions are swallowed because
    a milestone that saved must not fail on its index refresh; no running loop
    (a script, a unit test) is ordinary. A delete racing the refresh wins: the
    upsert claims the milestone's row first.
    """
    try:
        import asyncio

        from grimoire.services.embeddings import upsert_milestone_embedding

        asyncio.create_task(upsert_milestone_embedding(
            milestone.id, milestone.title, milestone.description, milestone.body,
        ))
    except RuntimeError:
        pass
    except Exception:  # noqa: BLE001 - never let indexing break a write
        logger.exception("embedding refresh failed for milestone %s", milestone.id)


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)
        embed_milestone(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 titles_for(milestone_ids: set[int]) -> dict[int, str]:
    """{id: title} for the given milestones, for rows that name where a record sits.

    No ownership filter, deliberately: the callers are listings whose rows the
    caller could already read, and a milestone title is part of "where does
    this record sit". Filtering here would blank the placement of a shared
    task in someone else's plan while still showing the task.
    """
    ids = {i for i in milestone_ids if i}
    if not ids:
        return {}
    async with async_session() as session:
        rows = (await session.execute(
            select(Milestone.id, Milestone.title).where(
                Milestone.id.in_(ids), Milestone.deleted_at.is_(None),
            )
        )).all()
    return {mid: title for mid, title in rows}


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)
        if {"title", "description", "body"} & set(fields):
            embed_milestone(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 three 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 many projects concurrently it asks for far more pooled connections than
    the pool holds, and every one of them waits out the checkout timeout. This
    does the same work in a fixed number of queries and one
    session: the milestones, their step counts, and their open steps.

    Each row carries `next_step` — the earliest open step, or None.
    That is NOT the same question `services/placement.py` answers: placement
    knows which step you are on and names the next one AFTER it, while a
    listing has no current step, so the earliest open one is the whole answer.
    The two share `OPEN_STEP_STATUSES` and the creation ordering so they can
    never disagree about which steps are candidates.
    """
    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]] = {}
        step_touched: dict[int, datetime] = {}
        next_step: dict[int, dict] = {}
        if milestones:
            milestone_ids = [m.id for m in milestones]
            # Both step queries below take the SAME visibility clause.
            # They have to: `next_step` names a step and the counts beside it
            # say how many there are, so a row that could name a step its own
            # progress excludes would be reporting two different milestones.
            readable = access_svc.readable_notes_clause(user_id)
            rows = await session.execute(
                select(
                    Note.milestone_id, Note.status, func.count(Note.id),
                    func.max(Note.updated_at),
                )
                .where(
                    Note.milestone_id.in_(milestone_ids),
                    Note.status.isnot(None),
                    Note.deleted_at.is_(None),
                    readable,
                )
                .group_by(Note.milestone_id, Note.status)
            )
            for milestone_id, status, count, latest in rows.fetchall():
                counts.setdefault(milestone_id, {})[status] = count
                if latest and (milestone_id not in step_touched
                               or latest > step_touched[milestone_id]):
                    step_touched[milestone_id] = latest

            # ONE query for every milestone in the batch, not one per row.
            # Fanned out per milestone, this listing drains the connection
            # pool; a third flat query keeps the cost
            # constant in the number of plans. Ordered the way
            # services/placement.py orders steps — creation order, the order a
            # plan is written and a batch create inserts — so the first row per
            # milestone IS its next open step.
            open_rows = await session.execute(
                select(Note.milestone_id, Note.id, Note.title, Note.status)
                .where(
                    Note.milestone_id.in_(milestone_ids),
                    Note.status.in_(OPEN_STEP_STATUSES),
                    Note.deleted_at.is_(None),
                    readable,
                )
                .order_by(Note.created_at.asc(), Note.id.asc())
            )
            for milestone_id, note_id, title, status in open_rows.fetchall():
                next_step.setdefault(
                    milestone_id, {"id": note_id, "title": title, "status": status},
                )

    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, {})))
        # A milestone's own updated_at doesn't move when its steps do, so a
        # plan whose steps closed today would read as untouched since it was
        # written. Touched is the later of the two.
        touched = [t for t in (m.updated_at, step_touched.get(m.id)) if t]
        entry["last_touched_at"] = max(touched).isoformat() if touched else None
        # Always present, None included. "8 of 9 done" tells a reader there is
        # an open step and not WHICH, and a gap that shape gets filled from
        # whatever id is nearest to hand — a retrieval hint carries an id and a
        # title and no status, so a done step can be reported as the open one.
        # A listing that names it leaves nothing to guess.
        entry["next_step"] = next_step.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, how far
# along it is, and which step is next. The plan itself (`body`) is
# get_milestone's job: on a project with a few dozen milestones, plan bodies
# would push enter_project past what an MCP client accepts as a tool result,
# and the session handshake would arrive as a file to page through.
# user_id / project_id / timestamps repeat what the caller already knows.
_BRIEF_FIELDS = (
    "id", "title", "description", "status", "order_index",
    "total", "completed", "pct", "status_counts", "next_step",
)


def brief_milestone_summary(
    rows: list[dict], *, limit: int | None = None,
) -> tuple[list[dict], int]:
    """Trim summary rows to the listing fields, optionally keeping only the
    most recently touched.

    With `limit`, keeps the N rows with the latest `last_touched_at`, whatever
    their status, most recent first: a handshake says what was worked on
    lately, and "active" alone doesn't (a plan can sit active for months).
    Without it, every row stays in its original order. Returns (rows, omitted).
    """
    kept = rows
    if limit is not None:
        kept = sorted(
            rows, key=lambda r: r.get("last_touched_at") or "", reverse=True,
        )[:limit]
    brief = [{k: r[k] for k in _BRIEF_FIELDS if k in r} for r in kept]
    return brief, len(rows) - len(kept)


def unplanned_milestones(
    rows: list[dict], *, exclude_ids: set[int] = frozenset(), limit: int | None = None,
) -> tuple[list[dict], int]:
    """Active milestones with no steps yet, as (rows, omitted).

    A plan written as a milestone with a description and no steps is open work
    that nothing else names. It is never "touched" — touching is a step
    changing — so the recency list that brief_milestone_summary(limit=) builds
    can never reach it, and progress reads 0% either way. Unlisted, a roadmap
    written that way gets every later plan opened as a new milestone beside the
    one that already described it.

    `exclude_ids` drops milestones a caller already listed. Kept in roadmap
    order (order_index, then creation), the order they were written in.
    Rows are id, title and description: what a reader needs to recognise the
    plan, and not its body, which get_milestone reads.
    """
    found = [
        {"id": r["id"], "title": r.get("title"), "description": r.get("description")}
        for r in rows
        if r.get("status") == "active" and not r.get("total") and r["id"] not in exclude_ids
    ]
    kept = found if limit is None else found[:limit]
    return kept, len(found) - len(kept)