# Copyright (c) 2023-2024 Arista Networks, Inc.
# Use of this source code is governed by the Apache License 2.0
# that can be found in the LICENSE file.
"""Module related to system-level features and protocols tests."""

# Mypy does not understand AntaTest.Input typing
# mypy: disable-error-code=attr-defined
from __future__ import annotations

import re
from ipaddress import IPv4Address
from typing import TYPE_CHECKING, ClassVar

from pydantic import BaseModel, Field

from anta.custom_types import Hostname, PositiveInteger
from anta.models import AntaCommand, AntaTest
from anta.tools import get_failed_logs, get_value

if TYPE_CHECKING:
    from anta.models import AntaTemplate

CPU_IDLE_THRESHOLD = 25
MEMORY_THRESHOLD = 0.25
DISK_SPACE_THRESHOLD = 75


class VerifyUptime(AntaTest):
    """Verifies if the device uptime is higher than the provided minimum uptime value.

    Expected Results
    ----------------
    * Success: The test will pass if the device uptime is higher than the provided value.
    * Failure: The test will fail if the device uptime is lower than the provided value.

    Examples
    --------
    ```yaml
    anta.tests.system:
      - VerifyUptime:
          minimum: 86400
    ```
    """

    name = "VerifyUptime"
    description = "Verifies the device uptime."
    categories: ClassVar[list[str]] = ["system"]
    commands: ClassVar[list[AntaCommand | AntaTemplate]] = [AntaCommand(command="show uptime", revision=1)]

    class Input(AntaTest.Input):
        """Input model for the VerifyUptime test."""

        minimum: PositiveInteger
        """Minimum uptime in seconds."""

    @AntaTest.anta_test
    def test(self) -> None:
        """Main test function for VerifyUptime."""
        command_output = self.instance_commands[0].json_output
        if command_output["upTime"] > self.inputs.minimum:
            self.result.is_success()
        else:
            self.result.is_failure(f"Device uptime is {command_output['upTime']} seconds")


class VerifyReloadCause(AntaTest):
    """Verifies the last reload cause of the device.

    Expected Results
    ----------------
    * Success: The test will pass if there are NO reload causes or if the last reload was caused by the user or after an FPGA upgrade.
    * Failure: The test will fail if the last reload was NOT caused by the user or after an FPGA upgrade.
    * Error: The test will report an error if the reload cause is NOT available.

    Examples
    --------
    ```yaml
    anta.tests.system:
      - VerifyReloadCause:
    ```
    """

    name = "VerifyReloadCause"
    description = "Verifies the last reload cause of the device."
    categories: ClassVar[list[str]] = ["system"]
    commands: ClassVar[list[AntaCommand | AntaTemplate]] = [AntaCommand(command="show reload cause", revision=1)]

    @AntaTest.anta_test
    def test(self) -> None:
        """Main test function for VerifyReloadCause."""
        command_output = self.instance_commands[0].json_output
        if len(command_output["resetCauses"]) == 0:
            # No reload causes
            self.result.is_success()
            return
        reset_causes = command_output["resetCauses"]
        command_output_data = reset_causes[0].get("description")
        if command_output_data in [
            "Reload requested by the user.",
            "Reload requested after FPGA upgrade",
        ]:
            self.result.is_success()
        else:
            self.result.is_failure(f"Reload cause is: '{command_output_data}'")


class VerifyCoredump(AntaTest):
    """Verifies if there are core dump files in the /var/core directory.

    Expected Results
    ----------------
    * Success: The test will pass if there are NO core dump(s) in /var/core.
    * Failure: The test will fail if there are core dump(s) in /var/core.

    Info
    ----
    * This test will NOT check for minidump(s) generated by certain agents in /var/core/minidump.

    Examples
    --------
    ```yaml
    anta.tests.system:
      - VerifyCoreDump:
    ```
    """

    name = "VerifyCoredump"
    description = "Verifies there are no core dump files."
    categories: ClassVar[list[str]] = ["system"]
    commands: ClassVar[list[AntaCommand | AntaTemplate]] = [AntaCommand(command="show system coredump", revision=1)]

    @AntaTest.anta_test
    def test(self) -> None:
        """Main test function for VerifyCoredump."""
        command_output = self.instance_commands[0].json_output
        core_files = command_output["coreFiles"]
        if "minidump" in core_files:
            core_files.remove("minidump")
        if not core_files:
            self.result.is_success()
        else:
            self.result.is_failure(f"Core dump(s) have been found: {core_files}")


class VerifyAgentLogs(AntaTest):
    """Verifies that no agent crash reports are present on the device.

    Expected Results
    ----------------
    * Success: The test will pass if there is NO agent crash reported.
    * Failure: The test will fail if any agent crashes are reported.

    Examples
    --------
    ```yaml
    anta.tests.system:
      - VerifyAgentLogs:
    ```
    """

    name = "VerifyAgentLogs"
    description = "Verifies there are no agent crash reports."
    categories: ClassVar[list[str]] = ["system"]
    commands: ClassVar[list[AntaCommand | AntaTemplate]] = [AntaCommand(command="show agent logs crash", ofmt="text")]

    @AntaTest.anta_test
    def test(self) -> None:
        """Main test function for VerifyAgentLogs."""
        command_output = self.instance_commands[0].text_output
        if len(command_output) == 0:
            self.result.is_success()
        else:
            pattern = re.compile(r"^===> (.*?) <===$", re.MULTILINE)
            agents = "\n * ".join(pattern.findall(command_output))
            self.result.is_failure(f"Device has reported agent crashes:\n * {agents}")


class VerifyCPUUtilization(AntaTest):
    """Verifies whether the CPU utilization is below 75%.

    Expected Results
    ----------------
    * Success: The test will pass if the CPU utilization is below 75%.
    * Failure: The test will fail if the CPU utilization is over 75%.

    Examples
    --------
    ```yaml
    anta.tests.system:
      - VerifyCPUUtilization:
    ```
    """

    name = "VerifyCPUUtilization"
    description = "Verifies whether the CPU utilization is below 75%."
    categories: ClassVar[list[str]] = ["system"]
    commands: ClassVar[list[AntaCommand | AntaTemplate]] = [AntaCommand(command="show processes top once", revision=1)]

    @AntaTest.anta_test
    def test(self) -> None:
        """Main test function for VerifyCPUUtilization."""
        command_output = self.instance_commands[0].json_output
        command_output_data = command_output["cpuInfo"]["%Cpu(s)"]["idle"]
        if command_output_data > CPU_IDLE_THRESHOLD:
            self.result.is_success()
        else:
            self.result.is_failure(f"Device has reported a high CPU utilization: {100 - command_output_data}%")


class VerifyMemoryUtilization(AntaTest):
    """Verifies whether the memory utilization is below 75%.

    Expected Results
    ----------------
    * Success: The test will pass if the memory utilization is below 75%.
    * Failure: The test will fail if the memory utilization is over 75%.

    Examples
    --------
    ```yaml
    anta.tests.system:
      - VerifyMemoryUtilization:
    ```
    """

    name = "VerifyMemoryUtilization"
    description = "Verifies whether the memory utilization is below 75%."
    categories: ClassVar[list[str]] = ["system"]
    commands: ClassVar[list[AntaCommand | AntaTemplate]] = [AntaCommand(command="show version", revision=1)]

    @AntaTest.anta_test
    def test(self) -> None:
        """Main test function for VerifyMemoryUtilization."""
        command_output = self.instance_commands[0].json_output
        memory_usage = command_output["memFree"] / command_output["memTotal"]
        if memory_usage > MEMORY_THRESHOLD:
            self.result.is_success()
        else:
            self.result.is_failure(f"Device has reported a high memory usage: {(1 - memory_usage)*100:.2f}%")


class VerifyFileSystemUtilization(AntaTest):
    """Verifies that no partition is utilizing more than 75% of its disk space.

    Expected Results
    ----------------
    * Success: The test will pass if all partitions are using less than 75% of its disk space.
    * Failure: The test will fail if any partitions are using more than 75% of its disk space.

    Examples
    --------
    ```yaml
    anta.tests.system:
      - VerifyFileSystemUtilization:
    ```
    """

    name = "VerifyFileSystemUtilization"
    description = "Verifies that no partition is utilizing more than 75% of its disk space."
    categories: ClassVar[list[str]] = ["system"]
    commands: ClassVar[list[AntaCommand | AntaTemplate]] = [AntaCommand(command="bash timeout 10 df -h", ofmt="text")]

    @AntaTest.anta_test
    def test(self) -> None:
        """Main test function for VerifyFileSystemUtilization."""
        command_output = self.instance_commands[0].text_output
        self.result.is_success()
        for line in command_output.split("\n")[1:]:
            if "loop" not in line and len(line) > 0 and (percentage := int(line.split()[4].replace("%", ""))) > DISK_SPACE_THRESHOLD:
                self.result.is_failure(f"Mount point {line} is higher than 75%: reported {percentage}%")


class VerifyNTP(AntaTest):
    """Verifies that the Network Time Protocol (NTP) is synchronized.

    Expected Results
    ----------------
    * Success: The test will pass if the NTP is synchronised.
    * Failure: The test will fail if the NTP is NOT synchronised.

    Examples
    --------
    ```yaml
    anta.tests.system:
      - VerifyNTP:
    ```
    """

    name = "VerifyNTP"
    description = "Verifies if NTP is synchronised."
    categories: ClassVar[list[str]] = ["system"]
    commands: ClassVar[list[AntaCommand | AntaTemplate]] = [AntaCommand(command="show ntp status", ofmt="text")]

    @AntaTest.anta_test
    def test(self) -> None:
        """Main test function for VerifyNTP."""
        command_output = self.instance_commands[0].text_output
        if command_output.split("\n")[0].split(" ")[0] == "synchronised":
            self.result.is_success()
        else:
            data = command_output.split("\n")[0]
            self.result.is_failure(f"The device is not synchronized with the configured NTP server(s): '{data}'")


class VerifyNTPAssociations(AntaTest):
    """Verifies the Network Time Protocol (NTP) associations.

    Expected Results
    ----------------
    * Success: The test will pass if the Primary NTP server (marked as preferred) has the condition 'sys.peer' and
    all other NTP servers have the condition 'candidate'.
    * Failure: The test will fail if the Primary NTP server (marked as preferred) does not have the condition 'sys.peer' or
    if any other NTP server does not have the condition 'candidate'.

    Examples
    --------
    ```yaml
    anta.tests.system:
      - VerifyNTPAssociations:
          ntp_servers:
            - server_address: 1.1.1.1
              preferred: True
              stratum: 1
            - server_address: 2.2.2.2
              stratum: 2
            - server_address: 3.3.3.3
              stratum: 2
    ```
    """

    name = "VerifyNTPAssociations"
    description = "Verifies the Network Time Protocol (NTP) associations."
    categories: ClassVar[list[str]] = ["system"]
    commands: ClassVar[list[AntaCommand | AntaTemplate]] = [AntaCommand(command="show ntp associations")]

    class Input(AntaTest.Input):
        """Input model for the VerifyNTPAssociations test."""

        ntp_servers: list[NTPServer]
        """List of NTP servers."""

        class NTPServer(BaseModel):
            """Model for a NTP server."""

            server_address: Hostname | IPv4Address
            """The NTP server address as an IPv4 address or hostname. The NTP server name defined in the running configuration
            of the device may change during DNS resolution, which is not handled in ANTA. Please provide the DNS-resolved server name.
            For example, 'ntp.example.com' in the configuration might resolve to 'ntp3.example.com' in the device output."""
            preferred: bool = False
            """Optional preferred for NTP server. If not provided, it defaults to `False`."""
            stratum: int = Field(ge=0, le=16)
            """NTP stratum level (0 to 15) where 0 is the reference clock and 16 indicates unsynchronized.
            Values should be between 0 and 15 for valid synchronization and 16 represents an out-of-sync state."""

    @AntaTest.anta_test
    def test(self) -> None:
        """Main test function for VerifyNTPAssociations."""
        failures: str = ""

        if not (peer_details := get_value(self.instance_commands[0].json_output, "peers")):
            self.result.is_failure("None of NTP peers are not configured.")
            return

        # Iterate over each NTP server.
        for ntp_server in self.inputs.ntp_servers:
            server_address = str(ntp_server.server_address)
            preferred = ntp_server.preferred
            stratum = ntp_server.stratum

            # Check if NTP server details exists.
            if (peer_detail := get_value(peer_details, server_address, separator="..")) is None:
                failures += f"NTP peer {server_address} is not configured.\n"
                continue

            # Collecting the expected NTP peer details.
            expected_peer_details = {"condition": "candidate", "stratum": stratum}
            if preferred:
                expected_peer_details["condition"] = "sys.peer"

            # Collecting the actual NTP peer details.
            actual_peer_details = {"condition": get_value(peer_detail, "condition"), "stratum": get_value(peer_detail, "stratumLevel")}

            # Collecting failures logs if any.
            failure_logs = get_failed_logs(expected_peer_details, actual_peer_details)
            if failure_logs:
                failures += f"For NTP peer {server_address}:{failure_logs}\n"

        # Check if there are any failures.
        if not failures:
            self.result.is_success()
        else:
            self.result.is_failure(failures)
