from datetime import date from fastapi import APIRouter, Depends, UploadFile from pydantic import BaseModel from sqlmodel import Session, select from app.auth import get_current_user from app.db import get_session from app.models.models import Bank, PensionSnapshot, PensionSnapshotRead from app.services.pension_pdf import parse_pension_pdf router = APIRouter(prefix="/pensions", tags=["pensions"]) class PensionUploadResult(BaseModel): imported: int updated: int duplicates: int errors: list[str] snapshots: list[PensionSnapshotRead] class PensionManualEntry(BaseModel): fund: str period_start: date period_end: date saldo_anterior: float aportes: float rendimientos: float retiros: float traslados: float comision: float correccion: float = 0.0 bonificacion: float = 0.0 saldo_final: float class PensionManualRequest(BaseModel): entries: list[PensionManualEntry] def _upsert_snapshot( session: Session, fund: str, period_start: date, period_end: date, saldo_anterior: float, aportes: float, rendimientos: float, retiros: float, traslados: float, comision: float, correccion: float, bonificacion: float, saldo_final: float, source_filename: str, contract_number: str = "", ) -> tuple[PensionSnapshot, bool]: """Insert or update a pension snapshot. Returns (row, is_new).""" existing = session.exec( select(PensionSnapshot).where( PensionSnapshot.fund == Bank(fund), PensionSnapshot.period_start == period_start, PensionSnapshot.period_end == period_end, ) ).first() if existing: existing.saldo_anterior = saldo_anterior existing.aportes = aportes existing.rendimientos = rendimientos existing.retiros = retiros existing.traslados = traslados existing.comision = comision existing.correccion = correccion existing.bonificacion = bonificacion existing.saldo_final = saldo_final existing.source_filename = source_filename if contract_number: existing.contract_number = contract_number session.add(existing) return existing, False row = PensionSnapshot( fund=Bank(fund), contract_number=contract_number, period_start=period_start, period_end=period_end, saldo_anterior=saldo_anterior, aportes=aportes, rendimientos=rendimientos, retiros=retiros, traslados=traslados, comision=comision, correccion=correccion, bonificacion=bonificacion, saldo_final=saldo_final, source_filename=source_filename, ) session.add(row) return row, True @router.post("/upload", response_model=PensionUploadResult) async def upload_pension_pdfs( files: list[UploadFile], session: Session = Depends(get_session), _user: str = Depends(get_current_user), ): imported = 0 updated = 0 errors: list[str] = [] results: list[PensionSnapshot] = [] for file in files: filename = file.filename or "unknown.pdf" try: pdf_bytes = await file.read() fund_snapshots = parse_pension_pdf(pdf_bytes, filename) except ValueError as e: errors.append(str(e)) continue except Exception as e: errors.append(f"{filename}: {e}") continue for snap in fund_snapshots: row, is_new = _upsert_snapshot( session, fund=snap.fund, period_start=snap.period_start, period_end=snap.period_end, saldo_anterior=snap.saldo_anterior, aportes=snap.aportes, rendimientos=snap.rendimientos, retiros=snap.retiros, traslados=snap.traslados, comision=snap.comision, correccion=snap.correccion, bonificacion=snap.bonificacion, saldo_final=snap.saldo_final, source_filename=filename, contract_number=snap.contract_number, ) results.append(row) if is_new: imported += 1 else: updated += 1 if imported > 0 or updated > 0: session.commit() for row in results: session.refresh(row) return PensionUploadResult( imported=imported, updated=updated, duplicates=0, errors=errors, snapshots=[PensionSnapshotRead.model_validate(r) for r in results], ) @router.post("/manual", response_model=PensionUploadResult) def submit_manual_entries( body: PensionManualRequest, session: Session = Depends(get_session), _user: str = Depends(get_current_user), ): imported = 0 updated = 0 results: list[PensionSnapshot] = [] for entry in body.entries: row, is_new = _upsert_snapshot( session, fund=entry.fund, period_start=entry.period_start, period_end=entry.period_end, saldo_anterior=entry.saldo_anterior, aportes=entry.aportes, rendimientos=entry.rendimientos, retiros=entry.retiros, traslados=entry.traslados, comision=entry.comision, correccion=entry.correccion, bonificacion=entry.bonificacion, saldo_final=entry.saldo_final, source_filename="manual-entry", ) results.append(row) if is_new: imported += 1 else: updated += 1 if imported > 0 or updated > 0: session.commit() for row in results: session.refresh(row) return PensionUploadResult( imported=imported, updated=updated, duplicates=0, errors=[], snapshots=[PensionSnapshotRead.model_validate(r) for r in results], ) @router.get("/snapshots", response_model=list[PensionSnapshotRead]) def get_snapshots( session: Session = Depends(get_session), _user: str = Depends(get_current_user), ): rows = session.exec( select(PensionSnapshot).order_by( PensionSnapshot.period_end.desc(), # type: ignore[union-attr] PensionSnapshot.fund, ) ).all() return rows @router.get("/fund-summary", response_model=list[PensionSnapshotRead]) def get_fund_summary( session: Session = Depends(get_session), _user: str = Depends(get_current_user), ): """Return the latest snapshot per fund (by most recent period_end).""" all_rows = session.exec( select(PensionSnapshot).order_by( PensionSnapshot.period_end.desc(), # type: ignore[union-attr] ) ).all() seen: set[str] = set() latest: list[PensionSnapshot] = [] for row in all_rows: if row.fund.value not in seen: seen.add(row.fund.value) latest.append(row) return latest