/*
 * Copyright (C) 2014-2025 Intel Corporation.
 * SPDX-License-Identifier: MIT
 */

#ifndef _IALARM_H_
#define _IALARM_H_

#include "pin.H"
#include "control_chain.H"
#include <string.h>
#include <atomic>

namespace CONTROLLER
{
//use this to avoid cache line false sharing on counters
struct alignas(64) CACHELINE_COUNTER
{
    std::atomic_uint64_t _count;
};

//this is the interface class for all the alarms
class IALARM
{
  public:
    IALARM(THREADID tid, UINT64 count, BOOL need_ctxt, ALARM_MANAGER* manager);

    //arms all threads
    VOID Arm();

    //arms only thread id tid
    VOID Arm(THREADID tid) { _armed[tid] = 1; }

    //disarm alarm for thread is tid and init the counter
    VOID Disarm(THREADID tid);

    //disarm alarm for thread is tid and init the counter
    VOID Disarm();

    // Disarm global alarm and return if we did do a reset
    BOOL DisarmGlobalArmed()
    {
        int armed = 1;
        return _global_armed.compare_exchange_strong(armed, 0);
    }

    //set the number of counts to raise the vent after
    VOID SetCount(UINT64 count) { _target_count._count = count; }
    UINT64 GetCount() { return _target_count._count; }
    UINT64 GetGlobalCount() { return _global_count._count; }

    // Get address of the alarm
    ADDRINT GetAddress() { return _address; }

    // Return if this alarms has global counter flag
    BOOL HasGlobalCounter();

    virtual VOID UpdateAlarm(ALARM_MANAGER* alarm_manager, const string& icount_str)
    {
        ASSERT(FALSE, "UpdateAlarm is not supported for this alarm type");
    }
    virtual VOID UpdateAlarm(ALARM_MANAGER* alarm_manager, UINT64 count)
    {
        ASSERT(FALSE, "UpdateAlarm is not supported for this alarm type");
    }

    // Can this alarm support late handler
    virtual BOOL CanSupportLateHandler() { return FALSE; }

    // This function is called from external agent to indicate that one step
    // happen in the execution flow.
    // It fires the controller event when reaching to the defined condition.
    VOID CallExternalStep(ADDRINT ip, UINT8* instruction_bytes);

  protected:
    // Return the call order for delivering the event, use default value
    // as IPOINT_BEFORE to maintain the default behaviour
    UINT32 GetInstrumentOrder(IPOINT ipoint = IPOINT_BEFORE);
    UINT32 GetLateInstrumentOrder();

    //add if analysis function
    static VOID InsertIfCall_Count(IALARM* alarm, INS ins, UINT32 ninst,
                                   IPOINT point = IPOINT_BEFORE);

    //add then analysis function
    static VOID InsertThenCall_Fire(IALARM* alarm, INS ins, IPOINT point = IPOINT_BEFORE);

    //add late fire analysis functions
    static VOID Insert_LateInstrumentation(IALARM* alarm, INS ins);

    //return True if: 1. the alarm is armed
    //                2. we are in the correct tid
    //                3. we reached the target count
    static ADDRINT PIN_FAST_ANALYSIS_CALL Count(IALARM* ialarm, THREADID tid, UINT32 ninst,
                                                BOOL is_rep = FALSE);

    //return True if: 1. the alarm is armed
    //                2. we are in the correct tid
    //                3. we reached the target count
    static ADDRINT PIN_FAST_ANALYSIS_CALL GlobalCount(IALARM* ialarm, THREADID tid,
                                                      UINT32 ninst);

    //do fire
    static VOID Fire(IALARM* ialarm, CONTEXT* ctxt, VOID* ip, THREADID tid);

    // Late handler analysis routines
    static ADDRINT PIN_FAST_ANALYSIS_CALL ActivateLate(IALARM* ialarm, THREADID tid);
    static VOID LateFire(IALARM* ialarm, CONTEXT* ctxt, VOID* ip, THREADID tid);

    // Instrumentation routines for address alarms
    static VOID TraceAddress(TRACE trace, VOID* v);
    static VOID InsertIfCall_Target(IALARM* ialarm, INS ins);
    static VOID InsertIfCall_FirstIp(IALARM* ialarm, INS ins, IPOINT point);
    static VOID InsertIfCall_Rep(IALARM* ialarm, INS ins, IPOINT point);

    // Analysis routines for address alarms
    static ADDRINT PIN_FAST_ANALYSIS_CALL CheckTarget(IALARM* ialarm, THREADID tid,
                                                      ADDRINT branch_target);
    static ADDRINT PIN_FAST_ANALYSIS_CALL CheckTargetGlobal(IALARM* ialarm,
                                                            ADDRINT branch_target);
    static ADDRINT PIN_FAST_ANALYSIS_CALL CheckFirstIp(IALARM* ialarm, THREADID tid,
                                                       ADDRINT addr);
    static ADDRINT PIN_FAST_ANALYSIS_CALL CheckFirstIpGlobal(IALARM* ialarm, THREADID tid,
                                                             ADDRINT addr);
    static ADDRINT PIN_FAST_ANALYSIS_CALL CheckRep(IALARM* ialarm, THREADID tid,
                                                   BOOL is_first_iteration);
    static ADDRINT PIN_FAST_ANALYSIS_CALL CheckRepGlobal(IALARM* ialarm, THREADID tid,
                                                         BOOL is_first_iteration);

    // Thread start callback
    static VOID ThreadStart(THREADID tid, CONTEXT* ctxt, INT32 flags, VOID* v);

    static string ExtractFileName(const string& fullpath);

    static BOOL InsIsApx(INS ins);

    //whether we need context
    BOOL _need_context;

    //the thread Id
    UINT32 _tid;
    BOOL _tid_translated;

    //the count that we need to reach to fire
    CACHELINE_COUNTER _target_count;

    //counter per thread
    CACHELINE_COUNTER _thread_count[CONTROLLER_MAX_THREADS];
    volatile CACHELINE_COUNTER _global_count;

    BOOL _armed[CONTROLLER_MAX_THREADS];
    std::atomic_int _global_armed;

    ALARM_MANAGER* _alarm_manager;

    BOOL _activate_late_handler[CONTROLLER_MAX_THREADS];

    // Address value for address, symbol and image alarms
    ADDRINT _address;

    static set<ADDRINT> _thread_first_ip;
    static ADDRINT _threads_first_ip_vec[CONTROLLER_MAX_THREADS];

  private:
#if defined(TARGET_WINDOWS)
    __declspec(align(64))
#else
    __attribute__((aligned(64)))
#endif
        PIN_LOCK _lock;
};

} // namespace CONTROLLER
#endif
