from typing import Annotated, List, Optional, Union from fastapi import APIRouter from fastapi.params import Query from src.auth.service import CurrentUser, InternalKey, Token from src.config import DEFAULT_TC_ID from src.database.core import DbSession, CollectorDbSession from src.models import StandardResponse from src.optimum_time_constraint.flows import ( create_calculation, get_create_calculation_parameters, get_or_create_scope_equipment_calculation, recalculate_calculation, ) from src.calculation_time_constrains.schema import ( CalculationResultsRead, CalculationSelectedEquipmentUpdate, CalculationTimeConstrainsCreate, CalculationTimeConstrainsParametersCreate, CalculationTimeConstrainsParametersRead, CalculationTimeConstrainsParametersRetrive, CalculationTimeConstrainsRead, CreateCalculationQuery, EquipmentResult, CalculationTimeConstrainsReadNoResult, ) from src.optimum_time_constraint.service import ( bulk_update_equipment, get_calculation_result, get_calculation_result_by_day, get_calculation_by_assetnum, get_all_calculations, refresh_spareparts_service, ) router = APIRouter() get_calculation = APIRouter() @router.post( "", response_model=StandardResponse[Union[dict, CalculationTimeConstrainsRead]] ) async def create_calculation_time_constrains( token: Token, db_session: DbSession, collector_db_session: CollectorDbSession, current_user: CurrentUser, calculation_time_constrains_in: CalculationTimeConstrainsParametersCreate, params: Annotated[CreateCalculationQuery, Query()], ): """Save calculation time constrains Here""" scope_calculation_id = params.scope_calculation_id simulation_id = params.simulation_id if scope_calculation_id: results = await get_or_create_scope_equipment_calculation( db_session=db_session, scope_calculation_id=scope_calculation_id, calculation_time_constrains_in=calculation_time_constrains_in, ) else: results = await create_calculation( token=token, db_session=db_session, collector_db_session=collector_db_session, calculation_time_constrains_in=calculation_time_constrains_in, created_by=current_user.name, simulation_id=simulation_id, ) return StandardResponse(data=results, message="Data created successfully") @router.get( "", response_model=StandardResponse[List[CalculationTimeConstrainsReadNoResult]] ) async def get_all_simulation_calculations( db_session: DbSession, token: Token, current_user: CurrentUser, ): """Get all calculation time constrains Here""" calculations = await get_all_calculations(db_session=db_session) return StandardResponse( data=calculations, message="Data retrieved successfully", ) @router.get( "/parameters", response_model=StandardResponse[ Union[ CalculationTimeConstrainsParametersRetrive, CalculationTimeConstrainsParametersRead, ] ], ) async def get_calculation_parameters( db_session: DbSession, calculation_id: Optional[str] = Query(default=None) ): """Get all calculation parameter.""" parameters = await get_create_calculation_parameters( db_session=db_session, calculation_id=calculation_id ) return StandardResponse( data=parameters, message="Data retrieved successfully", ) @get_calculation.get( "/{calculation_id}", response_model=StandardResponse[CalculationTimeConstrainsRead] ) async def get_calculation_results( db_session: DbSession, calculation_id, token: InternalKey, include_risk_cost: int = Query(1, alias="risk_cost"), ): if calculation_id == "default": calculation_id = DEFAULT_TC_ID results = await get_calculation_result( db_session=db_session, calculation_id=calculation_id, token=token, include_risk_cost=include_risk_cost, ) return StandardResponse( data=results, message="Data retrieved successfully", ) @router.get( "/{calculation_id}/{assetnum}", response_model=StandardResponse[EquipmentResult] ) async def get_calculation_per_equipment(db_session: DbSession, calculation_id, assetnum): results = await get_calculation_by_assetnum( db_session=db_session, assetnum=assetnum, calculation_id=calculation_id ) return StandardResponse( data=results, message="Data retrieved successfully", ) @router.post( "/{calculation_id}/simulation", response_model=StandardResponse[CalculationResultsRead], ) async def get_simulation_result( db_session: DbSession, calculation_id, calculation_simuation_in: CalculationTimeConstrainsCreate, ): simulation_result = await get_calculation_result_by_day( db_session=db_session, calculation_id=calculation_id, simulation_day=calculation_simuation_in.intervalDays, ) return StandardResponse( data=simulation_result, message="Data retrieved successfully" ) @router.post("/update/{calculation_id}", response_model=StandardResponse[List[str]]) async def update_selected_equipment( db_session: DbSession, calculation_id, calculation_time_constrains_in: List[CalculationSelectedEquipmentUpdate], ): if calculation_id == "default": calculation_id = "3b9a73a2-bde6-418c-9e2f-19046f501a05" results = await bulk_update_equipment( db=db_session, selected_equipments=calculation_time_constrains_in, calculation_data_id=calculation_id, ) return StandardResponse( data=results, message="Data retrieved successfully", ) @router.post("/{calculation_id}/refresh-spareparts", response_model=StandardResponse[dict]) async def refresh_spareparts( db_session: DbSession, collector_db_session: CollectorDbSession, calculation_id: str, current_user: CurrentUser, ): """Refresh sparepart availability for an existing calculation""" await refresh_spareparts_service( db_session=db_session, collector_db_session=collector_db_session, calculation_id=calculation_id, ) return StandardResponse(data={}, message="Spareparts refreshed successfully") @router.post("/{calculation_id}/recalculate", response_model=StandardResponse[CalculationTimeConstrainsRead]) async def recalculate_calculation_api( db_session: DbSession, collector_db_session: CollectorDbSession, calculation_id: str, token: Token, current_user: CurrentUser, ): """Recalculate an existing simulation with fresh data""" results = await recalculate_calculation( token=token, db_session=db_session, collector_db_session=collector_db_session, calculation_id=calculation_id, ) return StandardResponse(data=results, message="Calculation updated with fresh data")