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

logger = logging.getLogger(__name__)


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 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 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]] = {}
        step_touched: dict[int, datetime] = {}
        if milestones:
            rows = await session.execute(
                select(
                    Note.milestone_id, Note.status, func.count(Note.id),
                    func.max(Note.updated_at),
                )
                .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, 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

    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
        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], *, 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)