/*
 * 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 JOB_MANAGER_H_
#define JOB_MANAGER_H_

#include <chrono>
#include <unordered_map>
#include "agg_types.h"
#include "parallel_tree_finder.h"
#include "reservation_manager.h"
#include "sharp_job.h"

#define MAX_SHARP_RAILS_NUM 4

using SetJobPtr = std::set<SharpJob*>;
using SetJobPtrInsertRes = std::pair<SetJobPtr::iterator, bool>;
using ExtJobIdToJob = std::unordered_map<SharpExtJobId, SharpJob*>;
using ExtJobIdToJobPair = std::pair<SharpExtJobId, SharpJob*>;
using SharpJobIdToJob = std::map<sharp_job_id_t, SharpJob*>;
using PortKeyToJobId = std::unordered_map<port_key_t, SharpExtJobId>;
using SharpJobIdToPortKeyToJobId = std::unordered_map<sharp_job_id_t, PortKeyToJobId>;
using ListJobs = std::list<SharpJob*>;
using ListOfJobsPerReservationKey = std::map<string, ListJobs>;   // Each reservation_key can have multiples job
using ConnectionIdToJob = std::map<int, SharpJob*>;
using ListJobIds = std::list<sharp_job_id_t>;

class JobManager
{
    friend class ReservationManager;
    friend class JobManagerStartJobTests_InitializeJobManager_MakeSureThatMembersWereInitializedCorrectly_Test;

    ParallelTreeFinder m_tree_finder_;
    ReservationManager& m_reservation_manager_;

    // Containers of current jobs
    ExtJobIdToJob m_jobs_;
    SharpJobIdToJob m_sharp_job_id_to_job_;   // DS to be removed at second phase
    SharpJobIdToPortKeyToJobId m_sharp_job_id_to_an_to_job_id;
    ConnectionIdToJob m_connection_id_to_job_;
    ListOfJobsPerReservationKey m_reservation_key_list_of_jobs;   // Hash table. Key is reservation_key, value is a list of jobs operating
                                                                  // under that reservation_key

    // List of available ids (that can still be used)
    ListJobIds m_available_job_ids_;
    CommandManager* m_command_manager_;

    sharp_job_id_t m_pool_min_sharp_id_;   // min sharp_job_id that is taken from a pool
    sharp_job_id_t m_max_sharp_id_;        // max allowed sharp_job_id
    bool m_job_id_solely_from_pool;        // allocate job id only from pool of ids, regadless of job type (e.g. multi-rail or not)
    bool m_is_init_ = false;

    std::unique_ptr<file_utils::DumpFile> m_resource_alloc_dump_file_;
    std::unique_ptr<file_utils::DumpFile> m_job_trees_dump_;

   public:
    JobManager(CommandManager* p_command_manager, ReservationManager& reservation_manager)
        : m_reservation_manager_(reservation_manager),
          m_command_manager_(p_command_manager),
          m_pool_min_sharp_id_(0),
          m_max_sharp_id_(0),
          m_job_id_solely_from_pool(false),
          m_resource_alloc_dump_file_{file_utils::GetDumpFileIfEnabled("Resource Alloc", "sharp_am_resource_alloc.dump")}
    {}

    ~JobManager();

    void Init(const sharp_job_id_t max_jobs_number);
    void InitTreeFinder() { m_tree_finder_.Init(); };
    void ReCreateTreeFinderTasks() { m_tree_finder_.ReCreateTasks(); };
    inline std::size_t GetJobsCount() const { return m_jobs_.size(); };

    void BeginJob(const sharp_begin_job& begin_job_data, const smx_ep* ep, uint64_t tid);
    void EndJob(const sharp_end_job& end_job_data);
    void ConnectionDisconnected(int conn_id);
    void RefreshJobsConnection();
    void StartJobReconnection(const bool should_reconnect_to_clients = false);
    void CheckJobsActiveness(const bool should_reconnect_to_clients = false);
    void JobInfoReplyReceived(const sharp_mgmt_job_info& job_info);

    SharpJob* FindJob(uint64_t client_job_id, uint64_t reservation_id, const char* reservation_key);
    SharpJob* GetSharpJobMod(sharp_job_id_t sharp_job_id, port_key_t an_port_key);
    const SharpJob* GetSharpJob(const SharpExtJobId& job_id) const;
    SharpJob* GetSharpJobMod(const SharpExtJobId& job_id);

    void HandleJobError(SharpJob* p_job);

    void GetJobsOnTree(SetJobPtr& job_ids, sharp_trees_t tree_id);
    void GetJobsOnSwitch(SetJobPtr& job_ids, uint64_t agg_node_guid);
    void GetJobsOnLink(SetJobPtr& job_ids, uint64_t agg_node_guid, uint32_t port_num);
    // const ListJobs* GetJobsByReservationId(const char* reservation_key) const;
    void GetJobsByReservationKey(SetJobPtr& job_ids, const char* reservation_key) const;

    int RestoreFromNetworkAllJobIdsOfRunningJobs();
    bool IsJobFileSupportSeamlessRestart(persistent_job_info* data);
    int UpdateRunningJobsData(persistent_job_info& data, string job_file_path);
    int RestoreFromNetworkAllJobQuotasOfRunningJobs();
    int RestoreFromNetworkTreeNodesConfigurations();
    int UpdateJobsWithRestoredJobsData();
    int UpdateRunningJobsWithTreesInfo();
    int UpdateRunningJobsQuota();
    int RestoreJobsTrees();
    int AddRestoredJobsToTreeNodes();
    int UpdatePendingTreeEdgesOnRunningJobs();
    void CleanRestoredJobs();
    void CleanJobsDataRestoredDuringSeamlessRestart();
    void CleanEndedJobs();
    SharpJob* CreateJobFromPersistentJobInfo(persistent_job_info* job_info);

    void FreeJobId(const sharp_job_id_t sharp_job_id);

    void RemoveJobIdFromPool(const sharp_job_id_t sharp_job_id);

    void HandleJobInfoRequest(const sharp_jobs_request& job_info_request, const smx_ep* ep, uint64_t tid);

    uint32_t GetInternalJobIdForJob(const uint64_t external_job_id) const;

    // Provide a way to execute a callback on all jobs
    inline void ExecCallbackOnJobs(const std::function<void(const ExtJobIdToJobPair&)>& callback)
    {
        std::for_each(m_jobs_.begin(), m_jobs_.end(), callback);
    }

    bool CleanJob(SharpJob* job, const bool is_job_end, std::string reason_msg);

    void UpdateRecoverdJobsMetrics();

   private:
    SharpJob const* GetSharpJobFromExternalJobId(const uint64_t external_job_id) const;

    sharp_am_status BeginJobInternal(const sharp_begin_job& begin_job_data,
                                     const smx_ep* ep,
                                     sharp_job_data& job_data,
                                     SharpExtJobId& external_job_id,
                                     SharpJob*& p_job);

    sharp_am_status CheckJobReservation(const sharp_begin_job& begin_job_data,
                                        ReservationInfo*& reservation,
                                        SharpExtJobId& external_job_id);

    SharpJob* AllocateJob(const sharp_begin_job& begin_job_data,
                          const smx_ep* ep,
                          sharp_am_status& status,
                          const uint64_t job_key,
                          SharpExtJobId& external_job_id);

    void ReCleanJob(SharpJob* job);
    bool FreeJob(SharpJob*& p_job);
    bool IsSharpJobIdFromPool(sharp_job_id_t sharp_job_id) const
    {
        return sharp_job_id >= m_pool_min_sharp_id_ && sharp_job_id <= m_max_sharp_id_;
    }

    // allocate sharp_job_id from a pool of allowed numbers.
    // the mechanism is legacy: it results in a single sharp_job_id along all subnet switches at a given time.
    sharp_job_id_t AllocateJobIdFromPool();

    // apply sharp_job_id to a given job, either by using legacy approach or newer one.
    sharp_am_status AssignSharpJobId(SharpJob* job);

    void SetJobResource(JobResource& job_resource, const sharp_begin_job& begin_job_data, SharpJob* p_job, const ReservationInfo* ri);

    bool CreateJobConnection(SharpJob* const job, const bool should_disconnect_before_creating_connection = false);
    void EndJobConnection(SharpJob* const job);
    bool IsPeriodicRefreshDisabled();

    // AN guids are stored per job, helping to uniquely identify it for a given sharp_job_id.
    void RegisterJobANs(const SharpJob* job);
    // Use when the AN guids to be registered are held by some iterable.
    template <class InputIter>
    void RegisterJobANs(sharp_job_id_t sharp_job_id, InputIter begin, InputIter end, const SharpExtJobId& external_job_id);
    void UnregisterJobANs(const SharpJob* job);

    // Reconstruct JobTrees on seamless restart
    int CreateJobTreeNodes();
    int CreateJobMulticastTreeNodes();
    int ReconstructJobTrees();
    int ReconstructJobMulticastTrees();

    // Metrics update
    void UpdateBeginJobMetrics(sharp_am_status job_status, bool is_sat, uint32_t num_hcas, const SharpJob* job);
    void UpdateCleanJobMetrics(std::chrono::seconds job_duration, bool is_sat, uint32_t num_hcas, const SharpJob& job);
    void UpdateEndJobMetrics(bool job_ended_successfully);
    void UpdateTreesLevelHistogram(const SharpJob& job, bool is_begin_job);
};
#endif   // JOB_MANAGER_H_
