from typing import Annotated, List, Optional, Union from fastapi import APIRouter from fastapi.params import Query import requests from src import config 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, CreateCalculationQuery, EquipmentResult, CalculationTimeConstrainsReadNoResult) from .service import (bulk_update_equipment, get_calculation_result, get_calculation_result_by_day, get_calculation_by_assetnum, get_all_calculations) 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, params: Annotated[CreateCalculationQuery, Query()], # scope_calculation_id: Optional[str] = Query(None), # with_results: Optional[int] = Query(0), # simulation_id = Query(None) ): """Save calculation time constrains Here""" scope_calculation_id = params.scope_calculation_id with_results = params.with_results 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 ) # requests.post(f"{config.AUTH_SERVICE_API}/sign-out", headers={ # "Authorization": f"Bearer {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.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", )