be-optimumoh/src/calculation_time_constrains/router.py

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