"""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()