FabledScribe/src/fabledassistant/services/recurrence.py

"""Recurring task rule validation, date calculation, and scheduled spawning."""
import asyncio
import calendar
import logging
from datetime import date, datetime, timedelta, timezone

from sqlalchemy import and_, select

from fabledassistant.models import async_session
from fabledassistant.models.note import Note, TaskStatus

logger = logging.getLogger(__name__)

_VALID_UNITS = {"day", "week", "month", "year"}


def validate_recurrence_rule(rule: dict) -> None:
    """Raise ValueError if rule is structurally invalid."""
    if not isinstance(rule, dict):
        raise ValueError("recurrence_rule must be an object")
    rule_type = rule.get("type")
    if rule_type not in ("interval", "calendar"):
        raise ValueError("recurrence_rule.type must be 'interval' or 'calendar'")
    unit = rule.get("unit")
    if unit not in _VALID_UNITS:
        raise ValueError(f"recurrence_rule.unit must be one of {_VALID_UNITS}")
    if rule_type == "interval":
        every = rule.get("every")
        if not isinstance(every, int) or every < 1:
            raise ValueError("recurrence_rule.every must be a positive integer")
    elif rule_type == "calendar":
        day = rule.get("day_of_month")
        if not isinstance(day, int) or not 1 <= day <= 31:
            raise ValueError("recurrence_rule.day_of_month must be an integer between 1 and 31")
        if unit == "year":
            month = rule.get("month")
            if not isinstance(month, int) or not 1 <= month <= 12:
                raise ValueError("recurrence_rule.month must be an integer between 1 and 12")
        elif unit != "month":
            raise ValueError("calendar recurrence only supports unit 'month' or 'year'")


def _add_months(d: date, months: int) -> date:
    month = d.month - 1 + months
    year = d.year + month // 12
    month = month % 12 + 1
    day = min(d.day, calendar.monthrange(year, month)[1])
    return d.replace(year=year, month=month, day=day)


def _next_day_of_month(from_date: date, day: int) -> date:
    """Return the next occurrence of day-of-month strictly after from_date."""
    year, month = from_date.year, from_date.month
    max_day = calendar.monthrange(year, month)[1]
    clamped = min(day, max_day)
    if clamped > from_date.day:
        return from_date.replace(day=clamped)
    # Advance one month
    if month == 12:
        year, month = year + 1, 1
    else:
        month += 1
    return date(year, month, min(day, calendar.monthrange(year, month)[1]))


def _next_annual_date(from_date: date, month: int, day: int) -> date:
    """Return the next occurrence of month/day after from_date."""
    max_day = calendar.monthrange(from_date.year, month)[1]
    candidate = date(from_date.year, month, min(day, max_day))
    if candidate > from_date:
        return candidate
    return date(
        from_date.year + 1,
        month,
        min(day, calendar.monthrange(from_date.year + 1, month)[1]),
    )


def calculate_next_due(rule: dict, from_date: date) -> date:
    """Return the next due date after from_date according to rule."""
    rule_type = rule["type"]
    unit = rule["unit"]
    if rule_type == "interval":
        every = rule["every"]
        if unit == "day":
            return from_date + timedelta(days=every)
        if unit == "week":
            return from_date + timedelta(weeks=every)
        if unit == "month":
            return _add_months(from_date, every)
        if unit == "year":
            return _add_months(from_date, every * 12)
    elif rule_type == "calendar":
        if unit == "month":
            return _next_day_of_month(from_date, rule["day_of_month"])
        if unit == "year":
            return _next_annual_date(from_date, rule["month"], rule["day_of_month"])
    raise ValueError(f"Unhandled recurrence rule: {rule!r}")


async def spawn_recurring_tasks() -> int:
    """Create child tasks for all recurring tasks whose spawn date has arrived.

    Returns the number of tasks spawned.
    """
    from fabledassistant.services.embeddings import upsert_note_embedding
    from fabledassistant.services.notes import create_note

    now = datetime.now(timezone.utc)
    spawned = 0

    async with async_session() as session:
        result = await session.execute(
            select(Note).where(
                and_(
                    Note.recurrence_rule.isnot(None),
                    Note.recurrence_next_spawn_at <= now,
                )
            )
        )
        tasks = result.scalars().all()

    for task in tasks:
        base_date = task.due_date or now.date()
        next_due = calculate_next_due(task.recurrence_rule, base_date)
        try:
            child = await create_note(
                task.user_id,
                title=task.title,
                body=task.body or "",
                status=TaskStatus.todo.value,
                priority=task.priority or "none",
                due_date=next_due,
                tags=list(task.tags or []),
                project_id=task.project_id,
                milestone_id=task.milestone_id,
                recurrence_rule=task.recurrence_rule,
            )
            text = f"{child.title}\n{child.body}".strip() if child.body else (child.title or "")
            if text:
                asyncio.create_task(upsert_note_embedding(child.id, task.user_id, text))
        except Exception:
            logger.exception("Failed to spawn recurring task %d", task.id)
            continue

        async with async_session() as session:
            parent = await session.get(Note, task.id)
            if parent:
                parent.recurrence_next_spawn_at = None
                await session.commit()

        spawned += 1
        logger.info("Spawned recurring task %d → child due %s", task.id, next_due)

    return spawned