galaxis-po/backend/app/api/pension.py

"""
Pension account API endpoints.
"""
from typing import List

from fastapi import APIRouter, Depends, HTTPException, Query
from sqlalchemy.orm import Session, joinedload

from app.core.database import get_db
from app.api.deps import CurrentUser
from app.models.pension import PensionAccount, PensionHolding
from app.schemas.pension import (
    PensionAccountCreate,
    PensionAccountUpdate,
    PensionAccountResponse,
    AllocationResult,
    RecommendationResult,
)
from app.services.pension_allocation import calculate_allocation, get_recommendation

router = APIRouter(prefix="/api/pension", tags=["pension"])


@router.post("/accounts", response_model=PensionAccountResponse, status_code=201)
async def create_account(
    data: PensionAccountCreate,
    current_user: CurrentUser,
    db: Session = Depends(get_db),
):
    account = PensionAccount(
        user_id=current_user.id,
        **data.model_dump(),
    )
    db.add(account)
    db.commit()
    db.refresh(account)
    return account


@router.get("/accounts", response_model=List[PensionAccountResponse])
async def list_accounts(
    current_user: CurrentUser,
    db: Session = Depends(get_db),
    skip: int = Query(0, ge=0),
    limit: int = Query(50, ge=1, le=200),
):
    accounts = (
        db.query(PensionAccount)
        .options(joinedload(PensionAccount.holdings))
        .filter(PensionAccount.user_id == current_user.id)
        .order_by(PensionAccount.created_at.desc())
        .offset(skip)
        .limit(limit)
        .all()
    )
    return accounts


@router.get("/accounts/{account_id}", response_model=PensionAccountResponse)
async def get_account(
    account_id: int,
    current_user: CurrentUser,
    db: Session = Depends(get_db),
):
    account = (
        db.query(PensionAccount)
        .options(joinedload(PensionAccount.holdings))
        .filter(PensionAccount.id == account_id, PensionAccount.user_id == current_user.id)
        .first()
    )
    if not account:
        raise HTTPException(status_code=404, detail="Pension account not found")
    return account


@router.put("/accounts/{account_id}", response_model=PensionAccountResponse)
async def update_account(
    account_id: int,
    data: PensionAccountUpdate,
    current_user: CurrentUser,
    db: Session = Depends(get_db),
):
    account = (
        db.query(PensionAccount)
        .options(joinedload(PensionAccount.holdings))
        .filter(PensionAccount.id == account_id, PensionAccount.user_id == current_user.id)
        .first()
    )
    if not account:
        raise HTTPException(status_code=404, detail="Pension account not found")

    update_data = data.model_dump(exclude_unset=True)
    for field, value in update_data.items():
        setattr(account, field, value)

    db.commit()
    db.refresh(account)
    return account


@router.post("/accounts/{account_id}/allocate", response_model=AllocationResult)
async def allocate_assets(
    account_id: int,
    current_user: CurrentUser,
    db: Session = Depends(get_db),
):
    account = (
        db.query(PensionAccount)
        .filter(PensionAccount.id == account_id, PensionAccount.user_id == current_user.id)
        .first()
    )
    if not account:
        raise HTTPException(status_code=404, detail="Pension account not found")

    result = calculate_allocation(
        account_id=account.id,
        account_type=account.account_type.value,
        total_amount=account.total_amount,
        birth_year=account.birth_year,
        target_retirement_age=account.target_retirement_age,
    )

    # Save allocation as holdings
    db.query(PensionHolding).filter(PensionHolding.account_id == account_id).delete()
    for alloc in result.allocations:
        holding = PensionHolding(
            account_id=account_id,
            asset_name=alloc.asset_name,
            asset_type=alloc.asset_type,
            amount=alloc.amount,
            ratio=alloc.ratio,
        )
        db.add(holding)
    db.commit()

    return result


@router.get("/accounts/{account_id}/recommendation", response_model=RecommendationResult)
async def get_account_recommendation(
    account_id: int,
    current_user: CurrentUser,
    db: Session = Depends(get_db),
):
    account = (
        db.query(PensionAccount)
        .filter(PensionAccount.id == account_id, PensionAccount.user_id == current_user.id)
        .first()
    )
    if not account:
        raise HTTPException(status_code=404, detail="Pension account not found")

    return get_recommendation(
        account_id=account.id,
        birth_year=account.birth_year,
        target_retirement_age=account.target_retirement_age,
    )