DFiP_Budget_planing/api/src/repository/summary_repository.py

from typing import Any

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

from src.repository.row_json import ROW_COLUMNS, fetch_row_texts, row_json_select


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

    @staticmethod
    def _svod1_source() -> str:
        return (
            "FROM v3.v_summary_view("
            "CAST(:year AS INT), "
            "CAST(:sheet AS TEXT), "
            "CAST(:org_ids AS INT[]))"
        )

    @staticmethod
    def _svod2_source(sheet: str) -> str:
        function_name = {
            "GO": "v3.v_svod2_go",
            "RF": "v3.v_svod2_rf",
        }[sheet]
        return (
            f"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))"
        )

    @staticmethod
    def _svod2_params(
        year: int,
        org_ids: list[int] | None,
        offset: int | None,
        limit: int | None,
        sort_by: str,
        sort_direction: str,
    ) -> dict:
        return {
            "year": year,
            "org_ids": org_ids,
            "offset": offset,
            "limit": limit,
            "sort_by": sort_by,
            "sort_direction": sort_direction,
        }

    async def get_svod1_rows(
        self,
        year: int,
        sheet: str,
        org_ids: list[int] | None = None,
    ) -> list[tuple]:
        query = text(f"SELECT {ROW_COLUMNS} {self._svod1_source()}")
        return (
            await self.db.execute(
                query,
                {"year": year, "sheet": sheet, "org_ids": org_ids},
            )
        ).all()

    async def get_svod1_json(
        self,
        year: int,
        sheet: str,
        org_ids: list[int] | None = None,
        strip_nulls: bool = False,
    ) -> list[str]:
        sql = f"{row_json_select(strip_nulls=strip_nulls)} {self._svod1_source()}"
        return await fetch_row_texts(
            self.db,
            sql,
            {"year": year, "sheet": sheet, "org_ids": org_ids},
        )

    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]:
        query = text(
            f"SELECT 'INPUT'::VARCHAR AS row_type, 3 AS depth, sort_order, data "
            f"{self._svod2_source(sheet)}"
        )
        return (
            await self.db.execute(
                query,
                self._svod2_params(year, org_ids, offset, limit, sort_by, sort_direction),
            )
        ).all()

    async def get_svod2_json(
        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",
        strip_nulls: bool = False,
    ) -> list[str]:
        select = row_json_select(row_type="'INPUT'::VARCHAR", depth="3", strip_nulls=strip_nulls)
        return await fetch_row_texts(
            self.db,
            f"{select} {self._svod2_source(sheet)}",
            self._svod2_params(year, org_ids, offset, limit, sort_by, sort_direction),
        )

    async def get_svod2_go_export_rows(
        self,
        year: int,
        org_ids: list[int] | None,
        sort_by: str,
        sort_direction: str,
    ) -> list[tuple]:
        """Load the compact flat representation produced for XLSX export."""
        common_quarter_fields = (
            "adj_current", "adj_ssp", "adj_rf", "adj_reserve",
            "payment_date", "payment_amount", "payment_amount_ho",
            "payment_amount_rf", "payment_act", "booking_amount",
            "actual_m1", "actual_m2", "actual_m3",
        )
        export_keys = [
            "header.smeta_type",
            "header.smeta_direction",
            "header.org_name",
            "header.name",
            "header.item_id",
            "header.num_group",
            "header.justification",
            *(f"plan.plan_q{quarter}" for quarter in range(1, 5)),
            "allocation.internal_order",
            "allocation.property_object",
            *(f"sequestration.DFIP.adj_q{quarter}" for quarter in range(1, 5)),
            *(f"reserve.amount_q{quarter}" for quarter in range(1, 5)),
            *(f"approved.approved_q{quarter}" for quarter in range(1, 5)),
            *(
                f"contract_detail.{field}"
                for field in (
                    "counterparty", "reference", "subject", "currency",
                    "ceiling_amount", "expenses_q1", "expenses_q2",
                    "expenses_q3", "expenses_q4", "vat_rate", "deadline",
                )
            ),
        ]
        for quarter in range(1, 5):
            quarter_fields = list(common_quarter_fields)
            if quarter == 4:
                quarter_fields.append("actual_spod")
            if quarter > 1:
                quarter_fields.extend(
                    (
                        "plan_revision_eco_change",
                        "plan_revision_item_adj",
                        "plan_revision_increase",
                        "plan_revision_sequester",
                    )
                )
            quarter_fields.extend(
                f"transfer_to_q{target_quarter}"
                for target_quarter in range(quarter + 1, 5)
            )
            if quarter < 4:
                quarter_fields.append("transfer_to_economy")
            export_keys.extend(f"q{quarter}.{field}" for field in quarter_fields)

        query = text(
            """
            SELECT sort_order, row_values
            FROM v3.v_svod2_go_export(
                CAST(:year AS INT),
                CAST(:org_ids AS INT[]),
                CAST(:sort_by AS VARCHAR),
                CAST(:sort_direction AS VARCHAR)
            )
            """
        )
        result = await self.db.execute(
            query,
            {
                "year": year,
                "org_ids": org_ids,
                "sort_by": sort_by,
                "sort_direction": sort_direction.lower(),
            },
        )
        rows = []
        for row in result.mappings():
            data = {
                key: self._parse_svod2_export_value(key, value)
                for key, value in zip(export_keys, row["row_values"])
            }
            rows.append(("INPUT", 3, row["sort_order"], data))
        return rows

    @staticmethod
    def _parse_svod2_export_value(key: str, value: str | None) -> Any:
        if value is None:
            return None
        numeric_prefixes = (
            "plan.",
            "sequestration.",
            "reserve.",
            "approved.",
        )
        numeric_contract_fields = {
            "contract_detail.ceiling_amount",
            "contract_detail.expenses_q1",
            "contract_detail.expenses_q2",
            "contract_detail.expenses_q3",
            "contract_detail.expenses_q4",
        }
        quarter_text_suffixes = (".payment_date", ".payment_act")
        is_numeric_quarter = (
            key.startswith(("q1.", "q2.", "q3.", "q4."))
            and not key.endswith(quarter_text_suffixes)
        )
        if (
            key.startswith(numeric_prefixes)
            or key in numeric_contract_fields
            or is_numeric_quarter
        ):
            return float(value)
        return value