import json from sqlalchemy import exists, false, func, select, text from sqlalchemy.ext.asyncio import AsyncSession from src.db.models.org_unit import OrgUnit from src.db.models.project import Project from src.db.models.rf_project_report import RfProjectReport from src.db.models.rf_project_report_line import RfProjectReportLine class ProjectRepository: def __init__(self, db: AsyncSession): self.db = db @staticmethod def _serialize_project(project: Project, org_unit_name: str | None, report_count: int, years: list[InterruptedError] | None = None) -> dict: return { "id": project.id, "name": project.name, "level": project.level, "parent_id": project.parent_id, "project_type": project.project_type, "vsp_format": project.vsp_format, "placement_type": project.placement_type, "object_address": project.object_address, "staff_count": project.staff_count, "total_area": float(project.total_area) if project.total_area is not None else None, "org_unit_id": project.org_unit_id, "org_unit_name": org_unit_name, "report_count": int(report_count), "years": years, } @staticmethod def _apply_project_filters(query, year: int | None, branch_id: int | None, org_unit_ids: list[int] | None): if year is not None: query = query.where( exists( select(1).where( RfProjectReport.project_id == Project.id, RfProjectReport.year == year, ) ) ) if branch_id is not None: query = query.where(Project.org_unit_id == branch_id) if org_unit_ids is not None: if len(org_unit_ids) == 0: query = query.where(false()) else: query = query.where(Project.org_unit_id.in_(org_unit_ids)) return query async def get_list( self, year: int | None = None, branch_id: int | None = None, offset: int | None = None, limit: int | None = None, with_count: bool = False, org_unit_ids: list[int] | None = None, ) -> list[dict] | tuple[int, list[dict]]: report_count_subq = ( select(func.count(RfProjectReport.id)) .where(RfProjectReport.project_id == Project.id) .scalar_subquery() ) report_years_subq = ( select(RfProjectReport.year) .where(RfProjectReport.project_id == Project.id) .distinct() .scalar_subquery() ) query = ( select( Project, OrgUnit.title.label("org_unit_name"), report_count_subq.label("report_count"), func.array(report_years_subq).label("years"), ) .outerjoin(OrgUnit, OrgUnit.id == Project.org_unit_id) .order_by(Project.id) ) query = self._apply_project_filters(query, year, branch_id, org_unit_ids) if offset is not None: query = query.offset(offset) if limit is not None: query = query.limit(limit) rows = (await self.db.execute(query)).all() payload = [self._serialize_project(row[0], row[1], row[2], row[3]) for row in rows] if not with_count: return payload count_query = select(func.count(Project.id)) count_query = self._apply_project_filters(count_query, year, branch_id, org_unit_ids) total = int((await self.db.execute(count_query)).scalar() or 0) return total, payload async def get( self, project_id: int, org_unit_ids: list[int] | None = None, ) -> dict | None: report_count_subq = ( select(func.count(RfProjectReport.id)) .where(RfProjectReport.project_id == Project.id) .scalar_subquery() ) report_years_subq = ( select(RfProjectReport.year) .where(RfProjectReport.project_id == Project.id) .distinct() .scalar_subquery() ) query = ( select( Project, OrgUnit.title.label("org_unit_name"), report_count_subq.label("report_count"), func.array(report_years_subq).label("years"), ) .outerjoin(OrgUnit, OrgUnit.id == Project.org_unit_id) .where(Project.id == project_id) ) if org_unit_ids is not None: if len(org_unit_ids) == 0: return None query = query.where(Project.org_unit_id.in_(org_unit_ids)) row = (await self.db.execute(query)).first() if not row: return None return self._serialize_project(row[0], row[1], row[2], row[3]) async def get_reports(self, project_id: int) -> list[dict]: line_count_subq = ( select(func.count(RfProjectReportLine.id)) .where(RfProjectReportLine.rf_project_report_id == RfProjectReport.id) .scalar_subquery() ) query = ( select( RfProjectReport.id, RfProjectReport.year, RfProjectReport.report_type, RfProjectReport.created_by, RfProjectReport.created_at, RfProjectReport.updated_by, RfProjectReport.updated_at, line_count_subq.label("line_count"), ) .where(RfProjectReport.project_id == project_id) .order_by(RfProjectReport.year, RfProjectReport.report_type) ) rows = (await self.db.execute(query)).all() reports: list[dict] = [] for row in rows: reports.append( { "id": row[0], "year": row[1], "report_type": row[2], "created_by": row[3], "created_at": row[4].isoformat() if row[4] is not None else None, "updated_by": row[5], "updated_at": row[6].isoformat() if row[6] is not None else None, "line_count": int(row[7]), } ) return reports async def resolve_report_id( self, project_id: int, year: int, report_type: str, org_unit_ids: list[int] | None = None, ) -> int | None: query = ( select(RfProjectReport.id) .join(Project, Project.id == RfProjectReport.project_id) .where( Project.id == project_id, RfProjectReport.year == year, RfProjectReport.report_type == report_type, ) ) if org_unit_ids is not None: if len(org_unit_ids) == 0: return None query = query.where(Project.org_unit_id.in_(org_unit_ids)) return (await self.db.execute(query)).scalar_one_or_none() async def get_report_rows(self, report_id: int, sections: list[str] | None = None) -> list[tuple]: query = text( """ SELECT row_type, depth, sort_order, data FROM v3.v_form3_report_jsonb(CAST(:report_id AS INT), CAST(:sections AS TEXT[])) """ ) return (await self.db.execute(query, {"report_id": report_id, "sections": sections})).all() async def get_rf_rollup_rows( self, branch_id: int, year: int, sections: list[str] | None = None, ) -> list[tuple]: query = text( """ SELECT row_type, depth, sort_order, data FROM v3.v_form3_rf_rollup_jsonb(CAST(:branch_id AS INT), CAST(:year AS INT), CAST(:sections AS TEXT[])) """ ) return ( await self.db.execute( query, {"branch_id": branch_id, "year": year, "sections": sections}, ) ).all() async def upd_form3_cell( self, report_id: int, line_id: int, column: str, value, ) -> list[tuple]: query = text( """ SELECT row_type, depth, sort_order, data FROM v3.upd_form3_cell( CAST(:report_id AS INT), CAST(:line_id AS INT), CAST(:column AS TEXT), CAST(:value AS JSONB) ) """ ) rows = ( await self.db.execute( query, { "report_id": report_id, "line_id": line_id, "column": column, "value": json.dumps(value), }, ) ).all() return [tuple(r) for r in rows] async def upd_form3_cells(self, report_id: int, changes: list[dict]) -> list[tuple]: query = text( """ SELECT row_type, depth, sort_order, data FROM v3.upd_form3_cells(CAST(:report_id AS INT), CAST(:changes AS JSONB)) """ ) rows = ( await self.db.execute( query, { "report_id": report_id, "changes": json.dumps(changes), }, ) ).all() return [tuple(r) for r in rows] async def add_form3_line(self, report_id: int, expense_item_id: int) -> list[tuple]: query = text( """ SELECT row_type, depth, sort_order, data FROM v3.add_form3_line(CAST(:report_id AS INT), CAST(:expense_item_id AS INT)) """ ) rows = ( await self.db.execute( query, { "report_id": report_id, "expense_item_id": expense_item_id, }, ) ).all() return [tuple(r) for r in rows] async def del_form3_line(self, line_id: int) -> list[tuple]: query = text( """ SELECT row_type, depth, sort_order, data FROM v3.del_form3_line(CAST(:line_id AS INT)) """ ) rows = (await self.db.execute(query, {"line_id": line_id})).all() return [tuple(r) for r in rows] async def upd_project(self, project_id: int, column: str, value): query = text( """ SELECT v3.upd_project(CAST(:project_id AS INT), CAST(:column AS TEXT), CAST(:value AS JSONB)) """ ) result = ( await self.db.execute( query, { "project_id": project_id, "column": column, "value": json.dumps(value), }, ) ).scalar_one() return result async def add_project( self, name: str, year: int, org_unit_id: int, level: str = "project", parent_id: int | None = None, project_type: str | None = None, vsp_format: str | None = None, placement_type: str | None = None, object_address: str | None = None, staff_count: int | None = None, total_area: float | None = None, ) -> tuple[int, int, int]: query = text( """ SELECT project_id, limit_report_id, current_expenses_report_id FROM v3.add_project( CAST(:name AS VARCHAR), CAST(:year AS INT), CAST(:org_unit_id AS INT), CAST(:level AS VARCHAR), CAST(:parent_id AS INT), CAST(:project_type AS VARCHAR), CAST(:vsp_format AS VARCHAR), CAST(:placement_type AS VARCHAR), CAST(:object_address AS VARCHAR), CAST(:staff_count AS INT), CAST(:total_area AS NUMERIC) ) """ ) row = ( await self.db.execute( query, { "name": name, "year": year, "org_unit_id": org_unit_id, "level": level, "parent_id": parent_id, "project_type": project_type, "vsp_format": vsp_format, "placement_type": placement_type, "object_address": object_address, "staff_count": staff_count, "total_area": total_area, }, ) ).first() return int(row[0]), int(row[1]), int(row[2])