# Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0
"""Helper functions for testing CPU identification functionality."""

import subprocess
from enum import Enum, auto

from framework.utils import run_cmd
import host_tools.network as net_tools


class CpuVendor(Enum):
    """CPU vendors enum."""

    AMD = auto()
    INTEL = auto()


def get_cpu_vendor():
    """Return the CPU vendor."""
    brand_str = subprocess.check_output("lscpu", shell=True).strip().decode()
    if 'AuthenticAMD' in brand_str:
        return CpuVendor.AMD
    return CpuVendor.INTEL


def get_cpu_model_name():
    """Return the CPU model name."""
    _, stdout, _ = run_cmd("cat /proc/cpuinfo | grep 'model name' | uniq")
    info = stdout.strip().split(sep=":")
    assert len(info) == 2
    return info[1].strip()


def check_guest_cpuid_output(vm, guest_cmd, expected_header,
                             expected_separator,
                             expected_key_value_store):
    """Parse cpuid output inside guest and match with expected one."""
    ssh_connection = net_tools.SSHConnection(vm.ssh_config)
    _, stdout, stderr = ssh_connection.execute_command(guest_cmd)

    assert stderr.read() == ''
    while True:
        line = stdout.readline()
        if line != '':
            # All the keys have been matched. Stop.
            if not expected_key_value_store:
                break

            # Try to match the header if needed.
            if expected_header not in (None, ''):
                if line.strip() == expected_header:
                    expected_header = None
                continue

            # See if any key matches.
            # We Use a try-catch block here since line.split() may fail.
            try:
                [key, value] = list(
                    map(lambda x: x.strip(), line.split(expected_separator)))
            except ValueError:
                continue

            if key in expected_key_value_store.keys():
                assert value == expected_key_value_store[key], \
                    "%s does not have the expected value" % key
                del expected_key_value_store[key]
            else:
                for given_key in expected_key_value_store.keys():
                    if given_key in key:
                        assert value == expected_key_value_store[given_key], \
                            "%s does not have the expected value" % given_key
                        del expected_key_value_store[given_key]
                        break
        else:
            break

    assert not expected_key_value_store, \
        "some keys in dictionary have not been found in the output: %s" \
        % expected_key_value_store

@@ -175,16 +175,32 @@
 ty.%22%22%22%0A%0A
+import platform%0A
 import s
@@ -314,16 +314,61 @@
 _tools%0A%0A
+ARM_CPU_DICT = %7B%220xd0c%22: %22ARM_NEOVERSE_N1%22%7D%0A%0A
 %0Aclass C
@@ -736,24 +736,155 @@
 el name.%22%22%22%0A
+    if platform.machine() == %22aarch64%22:%0A        _, stdout, _ = run_cmd(%22cat /proc/cpuinfo %7C grep 'CPU part' %7C uniq%22)%0A    else:%0A    
     _, stdou
@@ -1022,29 +1022,349 @@
    r
-eturn info%5B1%5D.strip()
+aw_cpu_model = info%5B1%5D.strip()%0A    if platform.machine() == %22x86_64%22:%0A        return raw_cpu_model%0A    return ARM_CPU_DICT%5Braw_cpu_model%5D%0A%0A%0Adef get_instance_type():%0A    %22%22%22Get the instance type through IMDS.%22%22%22%0A    imds_cmd = %22curl http://169.254.169.254/latest/meta-data/instance-type%22%0A    _, stdout, _ = run_cmd(imds_cmd)%0A    return stdout
 %0A%0A%0Ad

enable perf tests on arm