from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncSession


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

    async def get_rows(
        self,
        year: int,
        sheet: str,
        org_ids: list[int] | None = None,
    ) -> list[tuple]:
        query = text(
            """
            SELECT row_type, depth, sort_order, data
            FROM v3.v_summary_view(
                CAST(:year AS INT),
                CAST(:sheet AS TEXT),
                CAST(:org_ids AS INT[])
            )
            """
        )
        return (
            await self.db.execute(
                query,
                {"year": year, "sheet": sheet, "org_ids": org_ids},
            )
        ).all()

    async def get_svod2_rows(
        self,
        year: int,
        sheet: str,
        org_ids: list[int] | None = None,
        offset: int | None = None,
        limit: int | None = None,
        sort_by: str = "org_name",
        sort_direction: str = "asc",
    ) -> list[tuple]:
        function_name = {
            "GO": "v3.v_svod2_go",
            "RF": "v3.v_svod2_rf",
        }[sheet]
        query = text(
            f"""
            SELECT
                'INPUT'::VARCHAR AS row_type,
                3 AS depth,
                sort_order,
                data
            FROM {function_name}(
                CAST(:year AS INT),
                CAST(:org_ids AS INT[]),
                CAST(:offset AS INT),
                CAST(:limit AS INT),
                CAST(:sort_by AS VARCHAR),
                CAST(:sort_direction AS VARCHAR)
            )
            ORDER BY sort_order
            """
        )
        return (
            await self.db.execute(
                query,
                {
                    "year": year,
                    "org_ids": org_ids,
                    "offset": offset,
                    "limit": limit,
                    "sort_by": sort_by,
                    "sort_direction": sort_direction,
                },
            )
        ).all()