from typing import 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 from src.models import StandardResponse from .flows import (create_calculation, get_create_calculation_parameters, get_or_create_scope_equipment_calculation) from .schema import (CalculationResultsRead, CalculationSelectedEquipmentUpdate, CalculationTimeConstrainsCreate, CalculationTimeConstrainsParametersCreate, CalculationTimeConstrainsParametersRead, CalculationTimeConstrainsParametersRetrive, CalculationTimeConstrainsRead, EquipmentResult) from .service import (bulk_update_equipment, get_calculation_result, get_calculation_result_by_day, get_calculation_by_assetnum) from src.database.core import CollectorDbSession 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, scope_calculation_id: Optional[str] = Query(None), with_results: Optional[int] = Query(0), simulation_id = Query(None) ): """Save calculation time constrains Here""" 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( "/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): if calculation_id == 'default': calculation_id = DEFAULT_TC_ID results = await get_calculation_result( db_session=db_session, calculation_id=calculation_id, token=token ) 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.put("/{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", )