social-app/backend/src/v1/automation_jobs/repository.py

from __future__ import annotations

from datetime import datetime, time, timedelta, timezone
from typing import TYPE_CHECKING
from uuid import UUID

from sqlalchemy import func, select, update
from sqlalchemy.ext.asyncio import AsyncSession
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError

from core.db.base_repository import BaseRepository
from models.agent_chat_session import AgentChatSession, SessionType
from models.automation_jobs import AutomationJob, AutomationJobStatus, ScheduleType

if TYPE_CHECKING:
    from v1.automation_jobs.schemas import (
        AutomationJobCreateRequest,
        AutomationJobUpdateRequest,
    )


class AutomationJobsRepository(BaseRepository[AutomationJob]):
    def __init__(self, session: AsyncSession) -> None:
        super().__init__(session=session, model=AutomationJob)

    async def list_due_jobs(
        self,
        *,
        now_utc: datetime,
        limit: int,
    ) -> list[AutomationJob]:
        stmt = (
            select(AutomationJob)
            .where(AutomationJob.deleted_at.is_(None))
            .where(AutomationJob.status == AutomationJobStatus.ACTIVE)
            .where(AutomationJob.next_run_at <= now_utc)
            .order_by(AutomationJob.next_run_at.asc())
            .limit(max(limit, 1))
        )
        rows = (await self._session.execute(stmt)).scalars().all()
        return list(rows)

    async def update_job_schedule(
        self,
        *,
        job_id: UUID,
        next_run_at: datetime,
        last_run_at: datetime,
    ) -> None:
        stmt = (
            update(AutomationJob)
            .where(AutomationJob.id == job_id)
            .where(AutomationJob.deleted_at.is_(None))
            .values(next_run_at=next_run_at, last_run_at=last_run_at)
        )
        await self._session.execute(stmt)
        await self._session.flush()

    async def get_or_create_chat_session(self, *, owner_id: UUID) -> UUID:
        stmt = (
            select(AgentChatSession.id)
            .where(AgentChatSession.user_id == owner_id)
            .where(AgentChatSession.deleted_at.is_(None))
            .where(AgentChatSession.session_type == SessionType.CHAT)
            .order_by(AgentChatSession.last_activity_at.desc())
            .limit(1)
        )
        existing = (await self._session.execute(stmt)).scalar_one_or_none()
        if existing is not None:
            return existing

        from uuid import uuid4

        new_session = AgentChatSession(
            id=uuid4(),
            user_id=owner_id,
            session_type=SessionType.CHAT,
        )
        self._session.add(new_session)
        await self._session.flush()
        return new_session.id

    async def list_by_owner(self, owner_id: UUID) -> list[AutomationJob]:
        stmt = (
            select(AutomationJob)
            .where(AutomationJob.owner_id == owner_id)
            .where(AutomationJob.deleted_at.is_(None))
            .order_by(AutomationJob.created_at.desc())
        )
        rows = (await self._session.execute(stmt)).scalars().all()
        return list(rows)

    async def count_user_jobs(self, owner_id: UUID) -> int:
        stmt = (
            select(func.count())
            .select_from(AutomationJob)
            .where(AutomationJob.owner_id == owner_id)
            .where(AutomationJob.deleted_at.is_(None))
            .where(AutomationJob.bootstrap_key.is_(None))
        )
        result = (await self._session.execute(stmt)).scalar_one()
        return int(result)

    def _resolve_timezone(self, timezone_str: str) -> ZoneInfo:
        try:
            return ZoneInfo(timezone_str)
        except ZoneInfoNotFoundError:
            return ZoneInfo("UTC")

    def _compute_initial_next_run_at(
        self,
        *,
        run_at: time,
        timezone_str: str,
        now_utc: datetime,
        schedule_type: ScheduleType,
    ) -> datetime:
        tz = self._resolve_timezone(timezone_str)
        local_now = now_utc.astimezone(tz)
        run_at_local = datetime.combine(local_now.date(), run_at, tz)
        if run_at_local.tzinfo is None:
            run_at_local = run_at_local.replace(tzinfo=tz)
        next_run_at = run_at_local
        if next_run_at <= local_now:
            if schedule_type == ScheduleType.DAILY:
                next_run_at = next_run_at + timedelta(days=1)
            else:
                next_run_at = next_run_at + timedelta(weeks=1)
        return next_run_at.astimezone(timezone.utc)

    async def create(
        self,
        owner_id: UUID,
        data: AutomationJobCreateRequest,
    ) -> AutomationJob:
        now_utc = datetime.now(tz=timezone.utc)
        timezone_obj = self._resolve_timezone(data.timezone)
        local_now = now_utc.astimezone(timezone_obj)
        date_ref = local_now.date()
        local_dt = datetime.combine(date_ref, data.run_at, timezone_obj)
        run_at_datetime = local_dt.astimezone(timezone.utc)
        next_run_at = self._compute_initial_next_run_at(
            run_at=data.run_at,
            timezone_str=data.timezone,
            now_utc=now_utc,
            schedule_type=data.schedule_type,
        )

        new_job = AutomationJob(
            owner_id=owner_id,
            created_by=owner_id,
            bootstrap_key=None,
            title=data.title,
            schedule_type=data.schedule_type,
            run_at=run_at_datetime,
            timezone=data.timezone,
            status=data.status,
            config=data.config.model_dump(mode="json"),
            next_run_at=next_run_at,
        )
        self._session.add(new_job)
        await self._session.flush()
        return new_job

    async def update(
        self,
        job_id: UUID,
        data: AutomationJobUpdateRequest,
    ) -> AutomationJob | None:
        update_values: dict[str, object] = {}
        existing_job: AutomationJob | None = None

        if data.title is not None:
            update_values["title"] = data.title
        if data.schedule_type is not None:
            update_values["schedule_type"] = data.schedule_type

        should_recompute_schedule = (
            data.run_at is not None
            or data.schedule_type is not None
            or data.timezone is not None
        )
        if should_recompute_schedule:
            now_utc = datetime.now(tz=timezone.utc)
            if existing_job is None:
                existing_job = await self.get_by_id(job_id)
                if existing_job is None:
                    return None

            effective_timezone = data.timezone or existing_job.timezone
            effective_timezone_obj = self._resolve_timezone(effective_timezone)
            effective_schedule_type = data.schedule_type or existing_job.schedule_type

            if data.run_at is not None:
                effective_run_at = data.run_at
            else:
                existing_timezone_obj = self._resolve_timezone(existing_job.timezone)
                effective_run_at = (
                    existing_job.run_at.astimezone(existing_timezone_obj)
                    .time()
                    .replace(microsecond=0)
                )

            local_now = now_utc.astimezone(effective_timezone_obj)
            local_dt = datetime.combine(
                local_now.date(),
                effective_run_at,
                effective_timezone_obj,
            )
            update_values["run_at"] = local_dt.astimezone(timezone.utc)
            update_values["next_run_at"] = self._compute_initial_next_run_at(
                run_at=effective_run_at,
                timezone_str=effective_timezone,
                now_utc=now_utc,
                schedule_type=effective_schedule_type,
            )
        if data.timezone is not None:
            update_values["timezone"] = data.timezone
        if data.status is not None:
            update_values["status"] = data.status
        if data.config is not None:
            if existing_job is None:
                existing_job = await self.get_by_id(job_id)
                if existing_job is None:
                    return None
            merged_config = {
                **existing_job.config,
                **data.config.model_dump(mode="json", exclude_unset=True),
            }
            update_values["config"] = merged_config

        if not update_values:
            return await self.get_by_id(job_id)

        return await self.update_by_id(job_id, update_values)

    async def soft_delete(self, job_id: UUID) -> None:
        await self.soft_delete_by_id(job_id)