#!/usr/share/venvs/netq-agent/bin/python
#
# Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# LicenseRef-NvidiaProprietary
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited
#

NETQ_AGENT_SERVICE_FILE = "/lib/systemd/system/netq-agent.service"
NETQ_AGENT_CONFIG_FILE = "/etc/netq/netq.yml"


def get_cpu_limit_from_netq_config():
    res = None
    try:
        with open(NETQ_AGENT_CONFIG_FILE, 'r') as file:
            lines = file.readlines()
        for i, line in enumerate(lines):
            if 'cpu-limit' in line:
                res = line.split(':')[1].strip()
                break
    except IOError:
        return res
    return res


def modify_service_file(cpu_quota):
    try:
        with open(NETQ_AGENT_SERVICE_FILE, 'r') as file:
            lines = file.readlines()

        # Check if 'CPUQuota' exists in the [Service] section
        service_section_found = False
        cpu_quota_exists = False

        for i, line in enumerate(lines):
            if '[Service]' in line:
                service_section_found = True
            if service_section_found and 'CPUQuota=' in line:
                cpu_quota_exists = True
                lines[i] = f'CPUQuota={cpu_quota}%\n'
                break

        if not cpu_quota_exists:
            if service_section_found:
                # Add 'CPUQuota=<value>' under the [Service] section
                for i, line in enumerate(lines):
                    if '[Service]' in line:
                        lines.insert(i + 1, f'CPUQuota={cpu_quota}%\n')
                        break
        with open(NETQ_AGENT_SERVICE_FILE, 'w') as file:
            file.writelines(lines)

        print(f"Updated {NETQ_AGENT_SERVICE_FILE} with CPUQuota={cpu_quota}%")

    except IOError as e:
        print(f"Couldn't read or write the file: {e}")


if __name__ == '__main__':
    cpu_limit = get_cpu_limit_from_netq_config()
    if cpu_limit:
        modify_service_file(cpu_limit)
