# Copyright (c) 2016 Cisco and/or its affiliates.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at:
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""VPP Grub Utility Library."""

import re

from vpplib.VPPUtil import VPPUtil

__all__ = ["VppGrubUtil"]


class VppGrubUtil(object):
    """VPP Grub Utilities."""

    def _get_current_cmdline(self):
        """
        Using /proc/cmdline return the current grub cmdline

        :returns: The current grub cmdline
        :rtype: string
        """

        # Get the memory information using /proc/meminfo
        cmd = "sudo cat /proc/cmdline"
        (ret, stdout, stderr) = VPPUtil.exec_command(cmd)
        if ret != 0:
            raise RuntimeError(
                "{} on node {} {} {}".format(cmd, self._node["host"], stdout, stderr)
            )

        self._current_cmdline = stdout.strip("\n")

    def _get_default_cmdline(self):
        """
        Using /etc/default/grub return the default grub cmdline

        :returns: The default grub cmdline
        :rtype: string
        """

        # Get the default grub cmdline
        rootdir = self._node["rootdir"]
        gfile = self._node["cpu"]["grub_config_file"]
        grubcmdline = self._node["cpu"]["grubcmdline"]
        cmd = "cat {}".format(rootdir + gfile)
        (ret, stdout, stderr) = VPPUtil.exec_command(cmd)
        if ret != 0:
            raise RuntimeError(
                "{} Executing failed on node {} {}".format(
                    cmd, self._node["host"], stderr
                )
            )

        # Get the Default Linux command line, ignoring commented lines
        lines = stdout.split("\n")
        for line in lines:
            if line == "" or line[0] == "#":
                continue
            ldefault = re.findall(r"{}=.+".format(grubcmdline), line)
            if ldefault:
                self._default_cmdline = ldefault[0]
                break

    def get_current_cmdline(self):
        """
        Returns the saved grub cmdline

        :returns: The saved grub cmdline
        :rtype: string
        """
        return self._current_cmdline

    def get_default_cmdline(self):
        """
        Returns the default grub cmdline

        :returns: The default grub cmdline
        :rtype: string
        """
        return self._default_cmdline

    def create_cmdline(self, isolated_cpus):
        """
        Create the new grub cmdline

        :param isolated_cpus: The isolated cpu string
        :type isolated_cpus: string
        :returns: The command line
        :rtype: string
        """
        grubcmdline = self._node["cpu"]["grubcmdline"]
        cmdline = self._default_cmdline
        value = cmdline.split("{}=".format(grubcmdline))[1]
        value = value.rstrip('"').lstrip('"')

        # jadfix intel_pstate=disable sometimes cause networks to
        # hang on reboot
        # iommu = re.findall(r'iommu=\w+', value)
        # pstate = re.findall(r'intel_pstate=\w+', value)
        # If there is already some iommu commands set, leave them,
        # if not use ours
        # if iommu == [] and pstate == []:
        #    value = '{} intel_pstate=disable'.format(value)

        # Replace isolcpus with ours
        isolcpus = re.findall(r"isolcpus=[\w+\-,]+", value)
        if not isolcpus:
            if isolated_cpus != "":
                value = "{} isolcpus={}".format(value, isolated_cpus)
        else:
            if isolated_cpus != "":
                value = re.sub(
                    r"isolcpus=[\w+\-,]+", "isolcpus={}".format(isolated_cpus), value
                )
            else:
                value = re.sub(r"isolcpus=[\w+\-,]+", "", value)

        nohz = re.findall(r"nohz_full=[\w+\-,]+", value)
        if not nohz:
            if isolated_cpus != "":
                value = "{} nohz_full={}".format(value, isolated_cpus)
        else:
            if isolated_cpus != "":
                value = re.sub(
                    r"nohz_full=[\w+\-,]+", "nohz_full={}".format(isolated_cpus), value
                )
            else:
                value = re.sub(r"nohz_full=[\w+\-,]+", "", value)

        rcu = re.findall(r"rcu_nocbs=[\w+\-,]+", value)
        if not rcu:
            if isolated_cpus != "":
                value = "{} rcu_nocbs={}".format(value, isolated_cpus)
        else:
            if isolated_cpus != "":
                value = re.sub(
                    r"rcu_nocbs=[\w+\-,]+", "rcu_nocbs={}".format(isolated_cpus), value
                )
            else:
                value = re.sub(r"rcu_nocbs=[\w+\-,]+", "", value)

        value = value.lstrip(" ").rstrip(" ")
        cmdline = '{}="{}"'.format(grubcmdline, value)
        return cmdline

    def apply_cmdline(self, node, isolated_cpus):
        """
        Apply cmdline to the default grub file

        :param node: Node dictionary with cpuinfo.
        :param isolated_cpus: The isolated cpu string
        :type node: dict
        :type isolated_cpus: string
        :return The vpp cmdline
        :rtype string
        """

        vpp_cmdline = self.create_cmdline(isolated_cpus)
        if len(vpp_cmdline):
            # Update grub
            # Save the original file
            rootdir = node["rootdir"]
            grubcmdline = node["cpu"]["grubcmdline"]
            ofilename = rootdir + node["cpu"]["grub_config_file"] + ".orig"
            filename = rootdir + node["cpu"]["grub_config_file"]

            # Write the output file
            # Does a copy of the original file exist, if not create one
            (ret, stdout, stderr) = VPPUtil.exec_command("ls {}".format(ofilename))
            if ret != 0:
                if stdout.strip("\n") != ofilename:
                    cmd = "sudo cp {} {}".format(filename, ofilename)
                    (ret, stdout, stderr) = VPPUtil.exec_command(cmd)
                    if ret != 0:
                        raise RuntimeError(
                            "{} failed on node {} {}".format(
                                cmd, self._node["host"], stderr
                            )
                        )

            # Get the contents of the current grub config file
            cmd = "cat {}".format(filename)
            (ret, stdout, stderr) = VPPUtil.exec_command(cmd)
            if ret != 0:
                raise RuntimeError(
                    "{} failed on node {} {}".format(cmd, self._node["host"], stderr)
                )

            # Write the new contents
            # Get the Default Linux command line, ignoring commented lines
            content = ""
            lines = stdout.split("\n")
            for line in lines:
                if line == "":
                    content += line + "\n"
                    continue
                if line[0] == "#":
                    content += line + "\n"
                    continue

                ldefault = re.findall(r"{}=.+".format(grubcmdline), line)
                if ldefault:
                    content += vpp_cmdline + "\n"
                else:
                    content += line + "\n"

            content = content.replace(r"`", r"\`")
            content = content.rstrip("\n")
            cmd = "sudo cat > {0} << EOF\n{1}\n".format(filename, content)
            (ret, stdout, stderr) = VPPUtil.exec_command(cmd)
            if ret != 0:
                raise RuntimeError(
                    "{} failed on node {} {}".format(cmd, self._node["host"], stderr)
                )

        return vpp_cmdline

    def __init__(self, node):
        distro = VPPUtil.get_linux_distro()
        if distro[0] == "Ubuntu":
            node["cpu"]["grubcmdline"] = "GRUB_CMDLINE_LINUX_DEFAULT"
        else:
            node["cpu"]["grubcmdline"] = "GRUB_CMDLINE_LINUX"

        self._node = node
        self._current_cmdline = ""
        self._default_cmdline = ""
        self._get_current_cmdline()
        self._get_default_cmdline()