WealthySmart/backend/app/api/v1/endpoints/pensions.py

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