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

from src.db.models.form4_project import Form4Project, Form4ProjectLevelEnum


class Form4ProjectRepository:
    def __init__(self, db: AsyncSession):
        self.db = db

    async def get(
        self,
        project_id: int,
        level: Form4ProjectLevelEnum | str | None = None,
    ) -> Form4Project | None:
        where = [Form4Project.id == project_id]
        if level is not None:
            where.append(Form4Project.level == level)

        return (
            (await self.db.execute(select(Form4Project).where(*where).limit(1)))
            .scalars()
            .first()
        )

    async def get_list(
        self,
        level: Form4ProjectLevelEnum | None = None,
        form_id: int | None = None,
        parent_id: int | None = None,
        offset: int | None = None,
        limit: int | None = None,
        with_count: bool = False,
    ) -> list[Form4Project] | tuple[int, list[Form4Project]]:
        if with_count:
            query = select(func.count().over().label("total_count"), Form4Project)
        else:
            query = select(Form4Project)

        where = []
        if level is not None:
            where.append(Form4Project.level == level)
        if form_id is not None:
            where.append(Form4Project.form_id == form_id)
        if parent_id is not None:
            where.append(Form4Project.parent_id == parent_id)

        query = query.where(*where).order_by(Form4Project.id)

        if offset is not None:
            query = query.offset(offset)
        if limit is not None:
            query = query.limit(limit)

        if with_count:
            result = (await self.db.execute(query)).all()
            if result:
                return result[0][0], [res[1] for res in result]
            else:
                count_query = select(func.count(Form4Project.id))
                if query.whereclause is not None:
                    count_query = count_query.where(query.whereclause)
                return (await self.db.execute(count_query)).scalar(), []
        else:
            return (await self.db.execute(query)).scalars().all()
