from datetime import date

from sqlalchemy import func, select
from sqlalchemy.orm import selectinload

from src.models import UsersORM
from src.models.tasks import TasksORM
from src.repository.base import BaseRepo


class UsersRepo(BaseRepo[UsersORM]):
    model: type[UsersORM] = UsersORM

    async def get_one_with_load(
        self,
        user_id: int,
        status: str | None = None,
        priority: str | None = None,
        tasks_limit: int | None = None,
        tasks_offset: int | None = 0,
        date_to: date | None = None,
        date_from: date | None = None,
    ) -> UsersORM | None:
        filters_sq: dict = {"user_id": user_id}
        if status:
            filters_sq["status"] = status
        if priority:
            filters_sq["priority"] = priority
        tasks_subquery = self._tasks_subquary(
            date_from=date_from, date_to=date_to, **filters_sq
        )

        if tasks_limit is not None:
            tasks_subquery = tasks_subquery.limit(tasks_limit)
        if tasks_offset is not None and tasks_offset > 0:
            tasks_subquery = tasks_subquery.offset(tasks_offset)
        query = (
            select(self.model)
            .where(self.model.id == user_id)
            .options(
                selectinload(
                    self.model.tasks.and_(TasksORM.id.in_(tasks_subquery))
                ).load_only(
                    TasksORM.id,
                    TasksORM.title,
                    TasksORM.due_date,
                    TasksORM.priority,
                    TasksORM.status,
                )
            )
        )
        result = await self.session.execute(query)
        obj = result.scalar_one_or_none()
        if obj and obj.tasks:
            obj.tasks.sort(
                key=lambda task: (
                    task.due_date or date.max,
                    -self._priority(task.priority),
                )
            )
        return obj

    async def get_tasks_count(
        self, date_from: date | None, date_to: date | None, **filter_by
    ) -> int:
        subq = self._tasks_subquary(date_from, date_to, **filter_by).subquery()
        stmt = select(func.count()).select_from(subq)
        result = await self.session.execute(stmt)
        return result.scalar_one()

    def _priority(self, priority: str):
        priority_map = {"low": 1, "medium": 2, "high": 3, "critical": 4}
        return priority_map.get(priority, 0)

    def _tasks_subquary(
        self, date_from: date | None, date_to: date | None, **filter_by
    ):
        subq = select(TasksORM.id).filter_by(**filter_by)
        if date_from is not None:
            subq = subq.filter(TasksORM.due_date >= date_from)
        if date_to is not None:
            subq = subq.filter(TasksORM.due_date <= date_to)
        return subq
