/**************************************************************************
*
* Copyright (c) 2026-2026 Diality Inc. - All Rights Reserved.
*
* THIS CODE MAY NOT BE COPIED OR REPRODUCED IN ANY FORM, IN PART OR IN
* WHOLE, WITHOUT THE EXPLICIT PERMISSION OF THE COPYRIGHT OWNER.
*
* @file    ModePostTreat.c
*
* @author  (last)      Praneeth Bunne
* @date    (last)      15-May-2026
*
* @author  (original)  Praneeth Bunne
* @date    (original)  17-Apr-2026
*
***************************************************************************/

#include "AirPump.h"
#include "AirTrap.h"
#include "BloodFlow.h"
#include "Buttons.h"
#include "DDInterface.h"
#include "EjectorMotor.h"
#include "Messaging.h"
#include "ModePostTreat.h"
#include "ModeStandby.h"
#include "OperationModes.h"
#include "Switches.h"
#include "SyringePump.h"
#include "TaskPriority.h"
#include "Timers.h"
#include "TubeSetAutoEject.h"
#include "TxParams.h"
#include "Valve3Way.h"
#include "Valves.h"

/**
 *  @addtogroup TDPostTreatmentMode
 *  @{
 */

#define POST_TREAT_STATE_DATA_PUB_INTERVAL (MS_PER_SECOND / TASK_PRIORITY_INTERVAL)

// ********** private definitions **********

static TD_POST_TREATMENT_STATE_T currentPostTreatmentState;                     ///< Current Post-Treatment sub-state
static BOOL patientDisconnectionConfirmed;                                      ///< Flag indicates user confirms patient disconnection.
static BOOL patientDisconnectHandled;                                           ///< Flag indicates patient disconnect actions have been handled
static BOOL disposableRemovalConfirmed;                                         ///< Flag indicates user confirms disposable removal.
static DISINFECT_SELECTION_T selectedDisinfectionSelection;                     ///< Selected disinfection option from UI.static BOOL disinfectionSelectionReceived;
static BOOL disinfectionSelectionReceived;
static volatile U32 postTreatmentStateBroadcastTimer;                           ///< Timer counter for pending power off request.

// ********** private function prototypes **********

static void setPostTreatStateTransition( void );
static TD_POST_TREATMENT_STATE_T handlePatientDisconnectionState( void );
static TD_POST_TREATMENT_STATE_T handleAutoEjectState( void );
static TD_POST_TREATMENT_STATE_T handleDisposableRemovalState(void );
static TD_POST_TREATMENT_STATE_T handleDisinfectionState( void );

/*********************************************************************//**
* @brief
* The initPostTreatmentMode function initializes the Post-Treatment mode.
* @details \b Inputs: none
* @details \b Outputs: Resets all Post-Treatment sub-state variables.
* @return none
**************************************************************************/
void initPostTreatmentMode( void )
{
    currentPostTreatmentState        = TD_POST_TREATMENT_PATIENT_DISCONNECTION_STATE;
    patientDisconnectionConfirmed    = FALSE;
    patientDisconnectHandled         = FALSE;
    disposableRemovalConfirmed       = FALSE;
    disinfectionSelectionReceived    = FALSE;
    selectedDisinfectionSelection    = DISINFECT_SEL_SKIP;
    postTreatmentStateBroadcastTimer = 0;
}

/*********************************************************************//**
* @brief
* The transitionToPostTreatmentMode function prepares the system for
* entry into Post-Treatment mode.
* @details \b Inputs:  none
* @details \b Outputs: Initializes mode specific state and returns the starting
* Post-Treatment sub-state.
* @return Initial Post-Treatment sub-state
**************************************************************************/
U32  transitionToPostTreatmentMode( void )
{
    initPostTreatmentMode();
    initTubeSetAutoEject();
    doorClosedRequired(TRUE);
    setPostTreatStateTransition();

    return (U32)currentPostTreatmentState;
}

/*********************************************************************//**
 * @brief
 * The execPostTreatmentMode function executes the Post-Treatment
 * mode state machine.
 * @details \b Inputs: stop button status and system conditions via helpers.
 * @details \b Outputs: Advances the Post-Treatment sub-state
 * @return current Post-Treatment sub-state.
 *************************************************************************/
U32 execPostTreatmentMode( void )
{
    BOOL stop = isStopButtonPressed();
    TD_POST_TREATMENT_STATE_T prevPostTreatmentState = currentPostTreatmentState;
    POST_TREATMENT_STATES_DATA_PAYLOAD_T payload = { 0 };

    if ( TRUE == stop )
    {
        // Ignore stop button in this mode.
    }

    // Execute mode state machine
    switch ( currentPostTreatmentState )
    {
        case TD_POST_TREATMENT_PATIENT_DISCONNECTION_STATE:
            currentPostTreatmentState = handlePatientDisconnectionState();
            break;

        case TD_POST_TREATMENT_AUTO_EJECT_STATE:
            currentPostTreatmentState = handleAutoEjectState();
            break;

        case TD_POST_TREATMENT_DISPOSABLE_REMOVAL_STATE:
            currentPostTreatmentState = handleDisposableRemovalState();
            break;

        case TD_POST_TREATMENT_DISINFECTION_STATE:
            currentPostTreatmentState = handleDisinfectionState();
            break;

        default:
            SET_ALARM_WITH_2_U32_DATA( ALARM_ID_TD_SOFTWARE_FAULT,
                                       SW_FAULT_ID_MODE_POST_TREATMENT_INVALID_STATE,
                                       (U32)currentPostTreatmentState );
            break;
    }
    postTreatmentStateBroadcastTimer++;

    if (postTreatmentStateBroadcastTimer >= POST_TREAT_STATE_DATA_PUB_INTERVAL)
    {
        postTreatmentStateBroadcastTimer = 0U;
        payload.mSubMode = (U32)currentPostTreatmentState;
        broadcastData( MSG_ID_TD_POST_TREATMENT_STATES_DATA,
                       COMM_BUFFER_OUT_CAN_TD_BROADCAST, (U08 *)&payload,
                       sizeof( POST_TREATMENT_STATES_DATA_PAYLOAD_T ) );
    }

    if ( prevPostTreatmentState != currentPostTreatmentState)
    {
        setPostTreatStateTransition();
    }

    return (U32)currentPostTreatmentState;
}

/*********************************************************************//**
* @brief
* The setPostTreatStateTransition function sets the actuators and variables
* for the state transition in Post-Treatment mode.
* @details \b Inputs: Valve states, Pump speed
* @details \b Outputs: Actuate valves, pumps as desired
* @Param state Post-Treatment state enum
* @return none
**************************************************************************/
static void setPostTreatStateTransition( void )
{
    signalBloodPumpHardStop();
    // Disable air trap level control.
    resetAirTrap();
    // H12 Air Pump - Off.
    setAirPumpState( AIR_PUMP_STATE_OFF, AIR_PUMP_MOTOR_OFF );
    // H13 Air Trap Valve In - Closed.
    set3WayValveState( H13_VALV, VALVE_3WAY_COMMON_TO_CLOSED_STATE );
    // H20 Air Trap Valve Out - Closed.
    set3WayValveState( H20_VALV, VALVE_3WAY_COMMON_TO_CLOSED_STATE );
    //H10 Heparin/Syringe Pump - Off.
    stopSyringePump();
    // TODO: Set H5 ejector Off.
    // Dialyzer - Bypassed.
    // Qdt      - 0.
    // Quf      - 0.
    // Qs       - 0.
    switch ( currentPostTreatmentState )
    {
        case TD_POST_TREATMENT_PATIENT_DISCONNECTION_STATE:
            setValvePosition( H1_VALV, VALVE_POSITION_B_OPEN );
            setValvePosition( H19_VALV, VALVE_POSITION_C_CLOSE );
            break;

        case TD_POST_TREATMENT_AUTO_EJECT_STATE:
            setValvePosition( H1_VALV, VALVE_POSITION_A_INSERT_EJECT );
            setValvePosition( H19_VALV, VALVE_POSITION_A_INSERT_EJECT );
            break;

        case TD_POST_TREATMENT_DISPOSABLE_REMOVAL_STATE:
            setValvePosition( H1_VALV, VALVE_POSITION_A_INSERT_EJECT );
            setValvePosition( H19_VALV, VALVE_POSITION_A_INSERT_EJECT );
            break;

        case TD_POST_TREATMENT_DISINFECTION_STATE:
            setValvePosition( H1_VALV, VALVE_POSITION_A_INSERT_EJECT );
            setValvePosition( H19_VALV, VALVE_POSITION_A_INSERT_EJECT );
            break;

        default:
            SET_ALARM_WITH_2_U32_DATA( ALARM_ID_TD_SOFTWARE_FAULT,
                    SW_FAULT_ID_MODE_POST_TREATMENT_INVALID_STATE1,
                    currentPostTreatmentState );
            break;
    }
}

/*********************************************************************//**
 * @brief
 * The handlePatientDisconnectionState function executes the Post-Treatment
 * mode Patient Disconnection state machine.
 * @details \b Inputs: none
 * @details \b Outputs: processed patient disconnection confirmation
 * @return next post-treatment mode state
 *************************************************************************/
static TD_POST_TREATMENT_STATE_T handlePatientDisconnectionState( void )
{
    TD_POST_TREATMENT_STATE_T state = TD_POST_TREATMENT_PATIENT_DISCONNECTION_STATE;

    VALVE_T valve;

    if ( ( TRUE == patientDisconnectionConfirmed ) && ( FALSE == patientDisconnectHandled ) )
    {
        // home valves and pumps while front door is still closed.
        homeBloodPump();
        for ( valve = H1_VALV; valve < NUM_OF_VALVES; valve++ )
        {
            setValvePosition( valve, VALVE_POSITION_A_INSERT_EJECT );
        }
        if ( ( getTreatmentParameterF32( TREATMENT_PARAM_HEPARIN_BOLUS_VOLUME ) > 0.0 ) ||
             ( getTreatmentParameterF32( TREATMENT_PARAM_HEPARIN_DELIVERY_RATE ) > 0.0 ) )
        {
            retractSyringePump();
        }
        doorClosedRequired( FALSE );
        // NOTE: patientDisconnectionConfirmed is NOT reset because
        // Patient Disconnect State needs it to move to next state
        patientDisconnectHandled = TRUE;
    }
    if ( TRUE == patientDisconnectionConfirmed )
    {
        patientDisconnectionConfirmed = FALSE;
        patientDisconnectHandled      = FALSE;
        state                         = TD_POST_TREATMENT_AUTO_EJECT_STATE;
    }

    return state;
}

/*********************************************************************//**
 * @brief
 * The handleAutoEjectState function calls Tube Set Auto-Eject Service and
 * advances to Disposable removal state once the service signals
 * completion.
 * @details \b Inputs: none
 * @details \b Outputs: Advances Post-Treatment state when eject is complete.
 * @return next Post-Treatment mode state
 *************************************************************************/
static TD_POST_TREATMENT_STATE_T handleAutoEjectState( void )
{
    TD_POST_TREATMENT_STATE_T state = TD_POST_TREATMENT_AUTO_EJECT_STATE;

    // call tube set eject service to auto-eject tube set
    execTubeSetAutoEject();

    if( TRUE == isTubeSetAutoEjectComplete() )
    {
        state = TD_POST_TREATMENT_DISPOSABLE_REMOVAL_STATE;
    }

    return state;
}

/*********************************************************************//**
 * @brief
 * The handleDisposableRemovalState function executes the Post-Treatment
 * mode Disposable Removal state machine.
 * @details \b Inputs: disposableRemovalConfirmed
 * @details \b Outputs: processed disposable removal confirmation
 * @return next post-treatment mode state
 *************************************************************************/
static TD_POST_TREATMENT_STATE_T handleDisposableRemovalState( void )
{
    TD_POST_TREATMENT_STATE_T state = TD_POST_TREATMENT_DISPOSABLE_REMOVAL_STATE;

    if ( TRUE == disposableRemovalConfirmed )
    {
        disposableRemovalConfirmed = FALSE;
        state                      = TD_POST_TREATMENT_DISINFECTION_STATE;
    }

    return state;
}

/*********************************************************************//**
 * @brief
 * The handleDisinfectionState function executes the Post-Treatment
 * Disinfection state.
 * @details \b Inputs: User disinfection selection from UI.
 * @details \b Outputs: selectedDisinfectionSelection
 * @return Current Post-Treatment state.
 *************************************************************************/
static TD_POST_TREATMENT_STATE_T handleDisinfectionState( void )
{
    TD_POST_TREATMENT_STATE_T state = TD_POST_TREATMENT_DISINFECTION_STATE;

    DD_COMMAND_T command = DD_CMD_NONE;

    if ( TRUE == disinfectionSelectionReceived )
    {
        switch ( selectedDisinfectionSelection )
        {
        case DISINFECT_SEL_SKIP:
            command = DD_CMD_START_DISINFECT_SKIP;
            break;
        case DISINFECT_SEL_FLUSH:
            command = DD_CMD_START_DISINFECT_FLUSH;
            break;
        case DISINFECT_SEL_HEAT:
            command = DD_CMD_START_DISINFECT_HEAT;
            break;
        case DISINFECT_SEL_HEAT_CITRIC:
            command = DD_CMD_START_DISINFECT_HEAT_CITRIC;
            break;
        case DISINFECT_SEL_SUBSTITUTION_PORT:
            command = DD_CMD_START_DISINFECT_SUBSTITUTION_PORT;
            break;
        default:
            command = DD_CMD_NONE;
            break;
        }

        if ( DD_CMD_NONE != command )
        {
            // Save selected disinfection option in Standby.
            signalStandbyDisinfectionSelection( command );

            // Request DD transition to Standby.
            cmdStopGenerateDialysate();

            // Transition TD to Standby.
            requestNewOperationMode( MODE_STAN );
        }

        disinfectionSelectionReceived = FALSE;
    }

    return state;
}

/*********************************************************************//**
 * @brief
 * The signalUserConfirmDisinfection function signals Post-Treatment mode
 * that the user has selected a disinfection option.
 * @details \b Inputs: current operation mode, current Post-Treatment state,
 * and selected disinfection option.
 * @details \b Outputs: disinfectionSelectionReceived,
 * selectedDisinfectionSelection.
 * @param options Disinfection option selected by the user.
 * @return TRUE if the selection is accepted, otherwise FALSE.
 *************************************************************************/
BOOL signalUserConfirmDisinfection( DISINFECT_SELECTION_T options )
{
    BOOL result = FALSE;

    if ( ( TD_POST_TREATMENT_DISINFECTION_STATE == currentPostTreatmentState ) &&
         ( (U32)options < (U32)NUM_OF_DISINFECT_SELECTIONS ) )
    {
        selectedDisinfectionSelection = options;
        disinfectionSelectionReceived = TRUE;
        result = TRUE;
    }

    return result;
}

/*********************************************************************//**
 * @brief
 * The handlePatientDisconnectionConfirmCmd function handles user
 * confirmation of patient disconnection.
 * @details \b Inputs: Patient disconnection confirmation message.
 * @details \b Outputs: Acceptance response and patient disconnection signal.
 * @param message Pointer to the message to handle.
 * @return TRUE if the confirmation is accepted, FALSE otherwise.
 *************************************************************************/
BOOL handlePatientDisconnectionConfirmCmd( MESSAGE_T *message )
{
    TD_OP_MODE_T mode = getCurrentOperationMode();
    TD_POST_TREATMENT_STATE_T postTxState = currentPostTreatmentState;
    U32 payloadLen = message->hdr.payloadLen;
    BOOL accepted = FALSE;
    UI_RESPONSE_PAYLOAD_T response;

    response.accepted = FALSE;
    response.rejectionReason = REQUEST_REJECT_REASON_NOT_ALLOWED_IN_CURRENT_MODE;

    if ( ( 0U == payloadLen ) && ( MODE_POST == mode ) &&
         ( TD_POST_TREATMENT_PATIENT_DISCONNECTION_STATE == postTxState ) )
    {
        patientDisconnectionConfirmed = TRUE;

        accepted = TRUE;
        response.accepted = TRUE;
        response.rejectionReason = REQUEST_REJECT_REASON_NONE;
    }
    else if ( 0U != payloadLen )
    {
        response.rejectionReason = REQUEST_REJECT_REASON_INVALID_REQUEST_FORMAT;
    }

    sendMessage( MSG_ID_TD_PATIENT_DISCONNECT_CONFIRM_RESPONSE,
                 COMM_BUFFER_OUT_CAN_TD_2_UI, (U08 *)&response,
                 sizeof( UI_RESPONSE_PAYLOAD_T ) );

    return accepted;
}

/*********************************************************************//**
 * @brief
 * The handleDisposableRemovalConfirmCmd function handles user confirmation
 * of disposable removal.
 * @details \b Inputs: Disposable removal confirmation message.
 * @details \b Outputs: Acceptance response and disposable removal signal.
 * @param message Pointer to the message to handle.
 * @return TRUE if the confirmation is accepted, FALSE otherwise.
 *************************************************************************/
BOOL handleDisposableRemovalConfirmCmd( MESSAGE_T *message )
{
    TD_OP_MODE_T mode = getCurrentOperationMode();
    U32 payloadLen = message->hdr.payloadLen;
    BOOL accepted = FALSE;
    UI_RESPONSE_PAYLOAD_T response;

    response.accepted = FALSE;
    response.rejectionReason = REQUEST_REJECT_REASON_NOT_ALLOWED_IN_CURRENT_MODE;

    if ( ( 0U == payloadLen ) && ( MODE_POST == mode ) &&
         ( TD_POST_TREATMENT_DISPOSABLE_REMOVAL_STATE == currentPostTreatmentState ) )
    {
        disposableRemovalConfirmed = TRUE;

        accepted = TRUE;
        response.accepted = TRUE;
        response.rejectionReason = REQUEST_REJECT_REASON_NONE;
    }
    else if ( 0U != payloadLen )
    {
        response.rejectionReason = REQUEST_REJECT_REASON_INVALID_REQUEST_FORMAT;
    }

    sendMessage( MSG_ID_TD_POST_TREATMENT_DISPOSABLE_REMOVAL_CONFIRM_RESPONSE,
            COMM_BUFFER_OUT_CAN_TD_2_UI, (U08 *)&response,
            sizeof( UI_RESPONSE_PAYLOAD_T ) );

    return accepted;
}


/**@}*/