/*
 * SPDX-FileCopyrightText: NVIDIA CORPORATION & AFFILIATES
 * Copyright (c) 2012-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 *
 * This software is available to you under a choice of one of two
 * licenses.  You may choose to be licensed under the terms of the GNU
 * General Public License (GPL) Version 2, available from the file
 * COPYING in the main directory of this source tree, or the
 * OpenIB.org BSD license below:
 *
 *     Redistribution and use in source and binary forms, with or
 *     without modification, are permitted provided that the following
 *     conditions are met:
 *
 *      - Redistributions of source code must retain the above
 *        copyright notice, this list of conditions and the following
 *        disclaimer.
 *
 *      - Redistributions in binary form must reproduce the above
 *        copyright notice, this list of conditions and the following
 *        disclaimer in the documentation and/or other materials
 *        provided with the distribution.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
 * BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
 * ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 *
 */

#ifndef RESERVATION_MANAGER_H_
#define RESERVATION_MANAGER_H_

#include <cstdint>
#include <map>
#include <unordered_map>
#include <unordered_set>

#include "agg_ib_types.h"
#include "option_manager.h"
#include "smx/smx_api.h"
#include "smx/smx_types.h"

class CommandManager;
class SharpJob;
class ReservationJobInfo;
class AggTree;

// Max number of jobs we can expect at most extreme case is as number of nodes divide by 2
#define MAX_RESERVATION_LIMIT (FABRIC_MAX_VALID_LID / 2)

// Common types
typedef std::map<uint64_t, ReservationJobInfo> JobInfos;   // Client JobID to ReservationJobInfo
typedef std::unordered_set<uint64_t> HashPortKeys;
typedef std::vector<uint64_t> VecPortKeys;

class ReservationJobInfo
{
    uint64_t m_client_job_id_;
    HashPortKeys m_guids;

   public:
    ReservationJobInfo() : m_client_job_id_(0) {};
    ~ReservationJobInfo() {};

    void SetClientJobId(uint64_t client_job_id) { m_client_job_id_ = client_job_id; }

    uint64_t GetClientJobId() const { return m_client_job_id_; }

    bool IsGuidsEmpty() const { return m_guids.empty(); }

    void InsertGuid(uint64_t port_key) { m_guids.insert(port_key); }

    bool CheckAtLeastOneGuidInJob(std::vector<uint64_t>& check_guids);
};

class ReservationInfo
{
    // Declare the reservation manager as a friend class, since it is managing the activity of the reservations
    friend class ReservationManager;

    struct sharp_reservation_info m_info_;
    JobInfos m_job_infos_;
    uint64_t m_reservation_id_;                // temporary member, should be removed when reservation_id removed
    uint32_t m_resource_limit_allowed_trees;   // Max number of trees this reservation is allowed to run
    uint32_t m_current_number_of_trees;        // Number of jobs this reservation is currently running

   public:
    ReservationInfo();

    ReservationInfo(const sharp_reservation_info* ri);

    ~ReservationInfo();

    void SetReservationInfo(const sharp_create_reservation& cr);

    void UpdateReservationInfo(const std::vector<uint64_t>& all_guids);

    void SetResourceLimitation();

    void SetState(enum sharp_reservation_state state) { m_info_.state = state; }

    const char* GetKey() const { return m_info_.reservation_key; }

    uint16_t GetPkey() const
    {
        return m_info_.pkey;   // return only PKEY, without membership bit
    }

    uint16_t GetFullPkey() const
    {
        return m_info_.pkey | 0x8000;   // Return PKEY with full membership. Force using full PKEY for the reservation
    }

    enum sharp_reservation_state GetState() const { return m_info_.state; }

    uint64_t GetGuid(uint64_t index) const { return m_info_.port_guids[index]; }

    uint32_t GetNumGuids() const { return m_info_.num_guids; }

    const struct sharp_reservation_resources& GetInfoResources() const { return m_info_.resource_limitations; }

    const struct sharp_reservation_info& GetInfo() const { return m_info_; }

    bool IsJobEmpty() const { return m_job_infos_.empty(); }

    bool IsJobIdsEnd(JobInfos::iterator& it) const { return (it == m_job_infos_.end()); }

    JobInfos::iterator GetJobIdsIter() { return m_job_infos_.begin(); }

    uint64_t GetReservationId() { return m_reservation_id_; }

    void SetReservationId(uint64_t reservation_id) { m_reservation_id_ = reservation_id; }

    enum sharp_reservation_status CheckModifiedResources(const sharp_create_reservation& create_reservation);
};

// Information that we are keeping per HCA that is part of a reservation
struct ReservationHCAInfo
{
    ReservationInfo* m_ri;       // Pointer to the reservation info this HCA belongs to
    uint16_t m_number_of_jobs;   // Number of jobs that use this HCA

    ReservationHCAInfo(ReservationInfo* ri) : m_ri(ri), m_number_of_jobs(0) {}
};

class ReservationManager
{
    typedef std::map<string, ReservationInfo*> MapKeyToReservationInfo;
    typedef std::map<uint64_t, ReservationHCAInfo> MapHCAReservationInfo;
    typedef std::unordered_map<uint64_t, string> MapIdToReservationKey;
    typedef std::pair<uint64_t, ReservationHCAInfo> PairHCAReservationInfo;
    typedef std::pair<string, ReservationInfo*> PairKeyToReservationInfo;
    typedef std::pair<uint64_t, string> PairIdToReservationKey;

    CommandManager* const m_command_manager_;

    MapKeyToReservationInfo m_reservations_;
    MapHCAReservationInfo m_hca_reservation_info_;
    MapIdToReservationKey m_id_to_reservation_key_;

   public:
    ReservationManager(CommandManager* const p_command_manager) : m_command_manager_(p_command_manager) {}

    ~ReservationManager();

    void CreateReservation(const sharp_create_reservation& create_reservation, const smx_ep* ep, uint64_t tid);

    void DeleteReservation(const sharp_delete_reservation& delete_reservation, const smx_ep* ep, uint64_t tid);

    void UpdateReservationByDelta(const sharp_update_reservation_by_delta& update_reservation_by_delta, const smx_ep* ep, uint64_t tid);

    void ReservationInfoRequest(const sharp_reservation_info_request& request, const smx_ep* ep, uint64_t tid);

    ReservationInfo* GetReservationByJob(const SharpJob* p_job);

    ReservationInfo* GetReservationByGuids(uint64_t* guids, uint32_t num_guids);

    ReservationInfo* GetReservationInfoByKey(const std::string& reservation_key);

    void ReadPersistentReservationInfoFiles();
    void DeleteAllPersistentReservationFiles() const;

    sharp_am_status ValidateReservationResources(ReservationInfo* ri, JobResource& job_resource, const SharpJob* p_job);

    void AddJob(ReservationInfo* ri, SharpJob* job);

    void RemoveJob(ReservationInfo* ri, const SharpJob* job);

    void RemoveJob(ReservationInfo* ri, JobInfos::iterator& it, const SharpJob* job);

    void UpdateReservationWithGuidPkeys(hca_port_key_t port_key, const Pkeys_t& pkeys);

    // Method to handle already running jobs for all reservations
    void HandleAlreadyRunningJobsForAllReservations();

   private:
    enum sharp_reservation_status InternalUpdateReservation(ReservationInfo* ri, const sharp_create_reservation& create_reservation);
    void InternalCreateNewReservation(ReservationInfo* ri);
    void InternalDeleteReservation(ReservationInfo* ri);
    enum sharp_reservation_status HandleRemovedGuidsAndActiveJobs(ReservationInfo* ri, std::vector<uint64_t>& deleted_guids);
    enum sharp_reservation_status HandleAddedGuidsAndOtherReservations(const char* reservation_key,
                                                                       uint32_t num_guids,
                                                                       uint64_t* port_guids,
                                                                       std::unordered_set<ReservationInfo*>& reservations_to_delete);

    enum sharp_reservation_status PerformUpdateAfterDataPreparation(ReservationInfo* ri,
                                                                    std::vector<uint64_t>& all_guids,
                                                                    std::vector<uint64_t>& added_guids,
                                                                    std::vector<uint64_t>& removed_guids);

    void PostUpdateReservation(ReservationInfo* ri, const std::unordered_set<ReservationInfo*>& reservations_to_delete);

    void RemoveReservationJobs(ReservationInfo* ri);
    void CreateReservationId(ReservationInfo& ri);
    void SetReservationId(ReservationInfo& ri, uint64_t ri_id);

    void AddReservationToDB(ReservationInfo* reservation);

    void CreateReservationInfoFile(const sharp_reservation_info& reservation_info, const string& persistent_path);

    void DeleteReservationInfoFile(const sharp_reservation_info& reservation_info, const string& persistent_path);

    void DeleteFile(const string& file_path) const;

    void UpdateReservationInfoFile(const sharp_reservation_info& reservation_info, const string& persistent_path);

    int ChechIfDirExists(const string& dir_name) const;

    int BuildDataFromReservationFile(const char* job_file, sharp_smx_msg*& smx_msg) const;

    void ApplyJobReservationResourcesChange(ReservationInfo* ri, bool is_added_job, const SharpJob* p_job);

    // Internal function to update reservation without sending reply
    enum sharp_reservation_status InternalUpdateReservationByDelta(const sharp_update_reservation_by_delta& update_reservation_by_delta);

    // Helper methods for UpdateReservationWithGuidPkeys
    void RemoveGuidFromReservation(ReservationInfo* ri, uint64_t guid);
    void AddOrCreateReservationForGuid(uint64_t guid, uint16_t port_pkey);

    // Helper method to handle already running jobs during reservation creation
    void HandleAlreadyRunningJobs(ReservationInfo* ri);
};
#endif   // RESERVATION_MANAGER_H_
