from __future__ import annotations

from datetime import datetime
from typing import TYPE_CHECKING
from uuid import UUID

from sqlalchemy import select, update
from sqlalchemy.ext.asyncio import AsyncSession

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

if TYPE_CHECKING:
    pass


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 == "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