fire-planner/fire_planner/api/income_streams.py

"""Income-stream CRUD nested under a scenario.

Streams are typed (salary / dividend / rental / pension / social_security
/ rsu / other) so the simulator can route the after-tax cash through the
jurisdiction tax engine using `tax_treatment` as the bucket.
"""
from __future__ import annotations

from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy import delete, select
from sqlalchemy.ext.asyncio import AsyncSession

from fire_planner.api.dependencies import get_session
from fire_planner.api.schemas import (
    IncomeStreamCreate,
    IncomeStreamOut,
    IncomeStreamPatch,
)
from fire_planner.db import IncomeStream, Scenario

router = APIRouter(tags=["income-streams"])


@router.get(
    "/scenarios/{scenario_id}/income-streams",
    response_model=list[IncomeStreamOut],
)
async def list_streams(
    scenario_id: int,
    session: AsyncSession = Depends(get_session),
) -> list[IncomeStream]:
    scen = await session.get(Scenario, scenario_id)
    if scen is None:
        raise HTTPException(status_code=404, detail="Scenario not found")
    rows = (await session.execute(
        select(IncomeStream).where(IncomeStream.scenario_id == scenario_id).order_by(
            IncomeStream.start_year, IncomeStream.id))).scalars().all()
    return list(rows)


@router.post(
    "/scenarios/{scenario_id}/income-streams",
    response_model=IncomeStreamOut,
    status_code=201,
)
async def create_stream(
    scenario_id: int,
    payload: IncomeStreamCreate,
    session: AsyncSession = Depends(get_session),
) -> IncomeStream:
    scen = await session.get(Scenario, scenario_id)
    if scen is None:
        raise HTTPException(status_code=404, detail="Scenario not found")
    if payload.end_year is not None and payload.end_year < payload.start_year:
        raise HTTPException(status_code=400, detail="end_year < start_year")
    stream = IncomeStream(scenario_id=scenario_id, **payload.model_dump())
    session.add(stream)
    await session.commit()
    await session.refresh(stream)
    return stream


@router.patch(
    "/income-streams/{stream_id}",
    response_model=IncomeStreamOut,
)
async def patch_stream(
    stream_id: int,
    payload: IncomeStreamPatch,
    session: AsyncSession = Depends(get_session),
) -> IncomeStream:
    stream = await session.get(IncomeStream, stream_id)
    if stream is None:
        raise HTTPException(status_code=404, detail="Income stream not found")
    updates = payload.model_dump(exclude_unset=True)
    for k, v in updates.items():
        setattr(stream, k, v)
    if stream.end_year is not None and stream.end_year < stream.start_year:
        raise HTTPException(status_code=400, detail="end_year < start_year")
    await session.commit()
    await session.refresh(stream)
    return stream


@router.delete(
    "/income-streams/{stream_id}",
    status_code=204,
    response_model=None,
)
async def delete_stream(
    stream_id: int,
    session: AsyncSession = Depends(get_session),
) -> None:
    stream = await session.get(IncomeStream, stream_id)
    if stream is None:
        raise HTTPException(status_code=404, detail="Income stream not found")
    await session.execute(delete(IncomeStream).where(IncomeStream.id == stream_id))
    await session.commit()