be-optimumoh/src/optimum_time_constraint/router.py

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