company-ai-platform/app/modules/subscriptions/services/management.py

from datetime import UTC, datetime
from typing import Any
from uuid import uuid4

from fastapi import HTTPException, status
from sqlalchemy import func, select
from sqlalchemy.orm import Session

from app.core.http.pagination import bounded_limit, bounded_offset
from app.core.utils.time import utc_now
from app.modules.audit.schemas import AuditLogCreate
from app.modules.audit.service import AuditService
from app.modules.feishu_users.constants import (
    FeishuCapability,
    FeishuUserRole,
    FeishuUserStatus,
)
from app.modules.feishu_users.models import FeishuUser
from app.modules.feishu_users.principal import FeishuPrincipal
from app.modules.subscriptions.constants import (
    EMPTY_PROMPT,
    INVALID_GROUP_TARGET,
    INVALID_QUIET_HOURS,
    MAX_ACTIVE_SUBSCRIPTIONS,
    PushSubscriptionStatus,
    SUBSCRIPTION_LIMIT_REACHED,
    SUBSCRIPTION_NOT_FOUND,
    SubscriptionAuditAction,
    SubscriptionScheduleType,
    SubscriptionTargetType,
)
from app.modules.subscriptions.models import PushDelivery, PushSubscription
from app.modules.subscriptions.services.schedule import (
    NormalizedSchedule,
    ScheduleParseError,
    next_occurrence,
    parse_quiet_clock,
    parse_schedule,
    validate_timezone,
)


class SubscriptionManagementService:
    """Manage subscriptions only through authenticated Feishu principals."""

    def __init__(self, db: Session):
        self.db = db

    def create_private(
        self,
        principal: FeishuPrincipal,
        schedule_expression: str,
        prompt: str,
        *,
        now: datetime | None = None,
    ) -> tuple[PushSubscription, NormalizedSchedule]:
        principal.require_capability(FeishuCapability.PRIVATE_SUBSCRIPTION)
        owner = self._active_owner(principal, for_update=True)
        return self._create(
            owner=owner,
            target_type=SubscriptionTargetType.USER,
            target_id=owner.open_id,
            schedule_expression=schedule_expression,
            prompt=prompt,
            now=now,
        )

    def create_group(
        self,
        principal: FeishuPrincipal,
        schedule_expression: str,
        prompt: str,
        *,
        now: datetime | None = None,
    ) -> tuple[PushSubscription, NormalizedSchedule]:
        principal.require_capability(FeishuCapability.GROUP_SUBSCRIPTION)
        owner = self._active_owner(principal, for_update=True)
        if (
            owner.role != FeishuUserRole.ADMIN
            or not principal.chat_id
            or principal.chat_type not in {"group", "group_chat"}
        ):
            raise HTTPException(
                status_code=status.HTTP_403_FORBIDDEN,
                detail=INVALID_GROUP_TARGET,
            )
        return self._create(
            owner=owner,
            target_type=SubscriptionTargetType.CHAT,
            target_id=principal.chat_id,
            schedule_expression=schedule_expression,
            prompt=prompt,
            now=now,
        )

    def list_for_owner(self, principal: FeishuPrincipal) -> list[PushSubscription]:
        principal.require_active()
        return list(
            self.db.execute(
                select(PushSubscription)
                .where(PushSubscription.owner_id == principal.owner_id)
                .order_by(PushSubscription.id.desc())
            ).scalars()
        )

    def latest_deliveries_for_owner(
        self,
        principal: FeishuPrincipal,
    ) -> dict[int, PushDelivery]:
        """Return at most one latest delivery per owner-scoped subscription."""

        principal.require_active()
        latest_ids = (
            select(func.max(PushDelivery.id))
            .join(PushSubscription)
            .where(PushSubscription.owner_id == principal.owner_id)
            .group_by(PushDelivery.subscription_id)
        )
        records = self.db.execute(
            select(PushDelivery).where(PushDelivery.id.in_(latest_ids))
        ).scalars()
        return {record.subscription_id: record for record in records}

    def pause(self, principal: FeishuPrincipal, code: str) -> PushSubscription:
        principal.require_capability(FeishuCapability.PRIVATE_SUBSCRIPTION)
        record = self._owned_subscription(principal.owner_id, code, for_update=True)
        if record.status == PushSubscriptionStatus.ACTIVE:
            record.status = PushSubscriptionStatus.PAUSED
            self._audit(principal, SubscriptionAuditAction.PAUSE, record)
            self.db.commit()
            self.db.refresh(record)
        return record

    def resume(
        self,
        principal: FeishuPrincipal,
        code: str,
        *,
        now: datetime | None = None,
    ) -> PushSubscription:
        principal.require_capability(FeishuCapability.PRIVATE_SUBSCRIPTION)
        self._active_owner(principal, for_update=True)
        record = self._owned_subscription(principal.owner_id, code, for_update=True)
        if record.status != PushSubscriptionStatus.PAUSED:
            return record
        self._ensure_active_capacity(principal.owner_id)
        current = _naive_utc(now or utc_now())
        if record.next_run_at is None or record.next_run_at <= current:
            next_run = next_occurrence(
                record.schedule_type,
                record.schedule_config,
                record.timezone,
                after=current,
            )
            if next_run is None:
                raise HTTPException(
                    status_code=status.HTTP_409_CONFLICT,
                    detail="Expired one-time subscriptions cannot be resumed",
                )
            record.next_run_at = next_run
        record.status = PushSubscriptionStatus.ACTIVE
        self._audit(principal, SubscriptionAuditAction.RESUME, record)
        self.db.commit()
        self.db.refresh(record)
        return record

    def cancel(self, principal: FeishuPrincipal, code: str) -> PushSubscription:
        principal.require_capability(FeishuCapability.PRIVATE_SUBSCRIPTION)
        record = self._owned_subscription(principal.owner_id, code, for_update=True)
        if record.status != PushSubscriptionStatus.CANCELLED:
            record.status = PushSubscriptionStatus.CANCELLED
            record.next_run_at = None
            self._audit(principal, SubscriptionAuditAction.CANCEL, record)
            self.db.commit()
            self.db.refresh(record)
        return record

    def set_timezone(
        self,
        principal: FeishuPrincipal,
        timezone_name: str,
        *,
        now: datetime | None = None,
    ) -> FeishuUser:
        principal.require_active()
        try:
            validate_timezone(timezone_name)
        except ScheduleParseError as exc:
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
                detail=str(exc),
            ) from exc
        owner = self._active_owner(principal, for_update=True)
        owner.timezone = timezone_name
        current = _naive_utc(now or utc_now())
        subscriptions = list(
            self.db.execute(
                select(PushSubscription).where(
                    PushSubscription.owner_id == owner.id,
                    PushSubscription.status.in_(
                        [
                            PushSubscriptionStatus.ACTIVE,
                            PushSubscriptionStatus.PAUSED,
                        ]
                    ),
                )
            ).scalars()
        )
        for subscription in subscriptions:
            subscription.timezone = timezone_name
            if (
                subscription.status == PushSubscriptionStatus.ACTIVE
                and subscription.schedule_type
                not in {
                    SubscriptionScheduleType.ONCE,
                    SubscriptionScheduleType.INTERVAL,
                }
            ):
                subscription.next_run_at = next_occurrence(
                    subscription.schedule_type,
                    subscription.schedule_config,
                    timezone_name,
                    after=current,
                )
        self._audit_user(
            principal,
            SubscriptionAuditAction.UPDATE_TIMEZONE,
            {"timezone": timezone_name},
        )
        self.db.commit()
        self.db.refresh(owner)
        return owner

    def set_quiet_hours(
        self,
        principal: FeishuPrincipal,
        start: str,
        end: str,
    ) -> FeishuUser:
        principal.require_active()
        try:
            quiet_start = parse_quiet_clock(start)
            quiet_end = parse_quiet_clock(end)
        except ScheduleParseError as exc:
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
                detail=str(exc),
            ) from exc
        if quiet_start == quiet_end:
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
                detail=INVALID_QUIET_HOURS,
            )
        owner = self._active_owner(principal, for_update=True)
        owner.quiet_hours_start = quiet_start
        owner.quiet_hours_end = quiet_end
        self._audit_user(
            principal,
            SubscriptionAuditAction.UPDATE_QUIET_HOURS,
            {"enabled": True},
        )
        self.db.commit()
        self.db.refresh(owner)
        return owner

    def clear_quiet_hours(self, principal: FeishuPrincipal) -> FeishuUser:
        principal.require_active()
        owner = self._active_owner(principal, for_update=True)
        owner.quiet_hours_start = None
        owner.quiet_hours_end = None
        self._audit_user(
            principal,
            SubscriptionAuditAction.UPDATE_QUIET_HOURS,
            {"enabled": False},
        )
        self.db.commit()
        self.db.refresh(owner)
        return owner

    def list_all(
        self,
        *,
        status_filter: str | None = None,
        owner_id: int | None = None,
        limit: int = 100,
        offset: int = 0,
    ) -> tuple[int, list[PushSubscription]]:
        stmt = select(PushSubscription)
        count_stmt = select(func.count()).select_from(PushSubscription)
        if status_filter:
            stmt = stmt.where(PushSubscription.status == status_filter)
            count_stmt = count_stmt.where(PushSubscription.status == status_filter)
        if owner_id is not None:
            stmt = stmt.where(PushSubscription.owner_id == owner_id)
            count_stmt = count_stmt.where(PushSubscription.owner_id == owner_id)
        stmt = (
            stmt.order_by(PushSubscription.id.desc())
            .limit(bounded_limit(limit))
            .offset(bounded_offset(offset))
        )
        total = int(self.db.scalar(count_stmt) or 0)
        return total, list(self.db.execute(stmt).scalars())

    def list_deliveries(
        self,
        *,
        status_filter: str | None = None,
        subscription_code: str | None = None,
        limit: int = 100,
        offset: int = 0,
    ) -> tuple[int, list[PushDelivery]]:
        stmt = select(PushDelivery).join(PushSubscription)
        count_stmt = (
            select(func.count())
            .select_from(PushDelivery)
            .join(PushSubscription)
        )
        if status_filter:
            stmt = stmt.where(PushDelivery.status == status_filter)
            count_stmt = count_stmt.where(PushDelivery.status == status_filter)
        if subscription_code:
            stmt = stmt.where(PushSubscription.code == subscription_code)
            count_stmt = count_stmt.where(PushSubscription.code == subscription_code)
        stmt = (
            stmt.order_by(PushDelivery.id.desc())
            .limit(bounded_limit(limit))
            .offset(bounded_offset(offset))
        )
        total = int(self.db.scalar(count_stmt) or 0)
        return total, list(self.db.execute(stmt).scalars())

    def _create(
        self,
        *,
        owner: FeishuUser,
        target_type: str,
        target_id: str,
        schedule_expression: str,
        prompt: str,
        now: datetime | None,
    ) -> tuple[PushSubscription, NormalizedSchedule]:
        clean_prompt = str(prompt or "").strip()
        if not clean_prompt:
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
                detail=EMPTY_PROMPT,
            )
        self._ensure_active_capacity(owner.id)
        try:
            schedule = parse_schedule(
                schedule_expression,
                owner.timezone,
                now=now,
            )
        except ScheduleParseError as exc:
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
                detail=str(exc),
            ) from exc
        record = PushSubscription(
            code=f"SUB-{uuid4().hex.upper()}",
            owner_id=owner.id,
            target_type=target_type,
            target_id=target_id,
            prompt=clean_prompt,
            schedule_type=schedule.schedule_type,
            schedule_config=schedule.schedule_config,
            timezone=schedule.timezone,
            next_run_at=schedule.next_run_at,
            status=PushSubscriptionStatus.ACTIVE,
            consented_at=_naive_utc(now or utc_now()),
        )
        self.db.add(record)
        self.db.flush()
        self._audit_values(
            actor=owner.code,
            action=SubscriptionAuditAction.CREATE,
            target_id=record.code,
            response={
                "target_type": target_type,
                "schedule_type": schedule.schedule_type,
            },
        )
        self.db.commit()
        self.db.refresh(record)
        return record, schedule

    def _ensure_active_capacity(self, owner_id: int) -> None:
        count = int(
            self.db.scalar(
                select(func.count())
                .select_from(PushSubscription)
                .where(
                    PushSubscription.owner_id == owner_id,
                    PushSubscription.status == PushSubscriptionStatus.ACTIVE,
                )
            )
            or 0
        )
        if count >= MAX_ACTIVE_SUBSCRIPTIONS:
            raise HTTPException(
                status_code=status.HTTP_409_CONFLICT,
                detail=SUBSCRIPTION_LIMIT_REACHED,
            )

    def _active_owner(
        self,
        principal: FeishuPrincipal,
        *,
        for_update: bool = False,
    ) -> FeishuUser:
        stmt = select(FeishuUser).where(FeishuUser.id == principal.owner_id)
        if for_update:
            stmt = stmt.with_for_update()
        owner = self.db.execute(stmt).scalar_one_or_none()
        if owner is None or owner.status != FeishuUserStatus.ACTIVE:
            raise HTTPException(
                status_code=status.HTTP_403_FORBIDDEN,
                detail="Feishu user is disabled",
            )
        if owner.tenant_key != principal.tenant_key or owner.open_id != principal.open_id:
            raise HTTPException(
                status_code=status.HTTP_403_FORBIDDEN,
                detail="Feishu identity mismatch",
            )
        return owner

    def _owned_subscription(
        self,
        owner_id: int,
        code: str,
        *,
        for_update: bool = False,
    ) -> PushSubscription:
        stmt = select(PushSubscription).where(
            PushSubscription.owner_id == owner_id,
            PushSubscription.code == code,
        )
        if for_update:
            stmt = stmt.with_for_update()
        record = self.db.execute(stmt).scalar_one_or_none()
        if record is None:
            raise HTTPException(
                status_code=status.HTTP_404_NOT_FOUND,
                detail=SUBSCRIPTION_NOT_FOUND,
            )
        return record

    def _audit(
        self,
        principal: FeishuPrincipal,
        action: str,
        record: PushSubscription,
    ) -> None:
        self._audit_values(
            actor=principal.user_code,
            action=action,
            target_id=record.code,
            response={"status": record.status},
        )

    def _audit_user(
        self,
        principal: FeishuPrincipal,
        action: str,
        response: dict[str, Any],
    ) -> None:
        self._audit_values(
            actor=principal.user_code,
            action=action,
            target_id=principal.user_code,
            response=response,
        )

    def _audit_values(
        self,
        *,
        actor: str,
        action: str,
        target_id: str,
        response: dict[str, Any],
    ) -> None:
        AuditService(self.db).record(
            AuditLogCreate(
                actor=actor,
                source="subscriptions",
                action=action,
                target_type="subscription",
                target_id=target_id,
                response_payload=response,
            )
        )


def _naive_utc(value: datetime) -> datetime:
    if value.tzinfo is None:
        return value
    return value.astimezone(UTC).replace(tzinfo=None)