FabledScribe/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
from scribe.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 (milestone 415).

    At the service, so every path that writes a milestone gets it — the lesson
    embed_note records (#2056): 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 (#3262).
    """
    try:
        import asyncio

        from scribe.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 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 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 (#4154).
    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 (rule 78).
            # 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.
            # #2384 was exactly this listing fanned out per milestone, and it
            # drained 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 (#4045). 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, and that is how a done step was reported as the
        # open one (#4154). 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. 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", "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. A project whose
    roadmap was written that way ended up with every later plan opened as a
    new milestone beside the one that already described it (milestone 415).

    `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)