import os
import subprocess
import functools
import json
import re
import pytest
from enum import IntEnum
from collections import deque
from common import SshConnectionError, is_ssh_connection_error, has_ssh_connection_err_msg, ClientServerTest
from retrying import retry

# EFA-specific message size lists for @pytest.mark.message_sizes decorator.
# Generic (shared) size lists live in fabtests/pytest/common.py.
DIRECT_SIZES = ["r:0,4,32", "r:0,1024,8192"]
# RMA variant: starts at 1 because fabtests RMA benchmarks reject 0-byte.
DIRECT_RMA_SIZES = ["r:1,4,32", "r:1,1024,8192"]
REMOTE_EXIT_SIZES = [65536, 131072, 1048576]
# Use the default behavior of the test, so pass None.
DGRAM_DEFAULT = [None]
DGRAM_PR_CI = ["l:16,128,8192"]


@functools.lru_cache(2)
@retry(retry_on_exception=is_ssh_connection_error, stop_max_attempt_number=3, wait_fixed=5000)
def parse_lspci_tree(server_id):
    """
    Function parse the lspci output and construct a tree of the PCIe hierarchy

    Snippet of lspci -tv output on p5en instance type

    $ lspci -tv
    +-[0000:44]---00.0-[45-5c]----00.0-[46-5c]--+-00.0-[47]----00.0  Amazon.com, Inc. Elastic Network Adapter (ENA)
    |                                           +-00.1-[48]----00.0  Amazon.com, Inc. Elastic Network Adapter (ENA)
    |                                           +-00.2-[49]----00.0  Amazon.com, Inc. Elastic Network Adapter (ENA)
    |                                           +-00.3-[4a]----00.0  Amazon.com, Inc. Elastic Network Adapter (ENA)
    |                                           +-00.4-[4b]--
    |                                           +-00.5-[4c]--
    |                                           +-00.6-[4d]--
    |                                           +-00.7-[4e]--
    |                                           +-01.0-[4f]--
    |                                           +-01.1-[50]--
    |                                           +-01.2-[51]--
    |                                           +-01.3-[52]--
    |                                           +-01.4-[53]--
    |                                           +-01.5-[54]--
    |                                           +-01.6-[55]----00.0  Amazon.com, Inc. Elastic Fabric Adapter (EFA)
    |                                           +-01.7-[56]----00.0  Amazon.com, Inc. Elastic Fabric Adapter (EFA)
    |                                           +-02.0-[57]----00.0  Amazon.com, Inc. Elastic Fabric Adapter (EFA)
    |                                           +-02.1-[58]----00.0  Amazon.com, Inc. Elastic Fabric Adapter (EFA)
    |                                           +-02.2-[59]----00.0  NVIDIA Corporation GH100 [H200 SXM 141GB]
    |                                           +-02.3-[5a]----00.0  NVIDIA Corporation GH100 [H200 SXM 141GB]
    |                                           +-02.4-[5b]----00.0  Amazon.com, Inc. NVMe SSD Controller
    |                                           \-02.5-[5c]----00.0  Amazon.com, Inc. NVMe SSD Controller

    In the output above, "[0000:44]" is a bus downstream of the root complex
    "00.0-[45-5c]" is a bridge and "00.0-[46-5c]" is a switch upstream of buses 0x45 to 0x5c and 0x46 to 0x5c respectively

    Looking at the line with the EFA NIC,
    "01.6-[55]" is the bridge on bus 1 device 6 on bus 46 and is upstream of bus 55
    "00.0  Amazon.com, Inc. Elastic Fabric Adapter (EFA)" is the EFA NIC on bus 55, device function 00.0
    """
    timeout = 60
    result = subprocess.run(f"ssh {server_id} lspci -tv",
            shell=True,
            stdout=subprocess.PIPE, stderr=subprocess.PIPE,
            encoding="utf-8", timeout=timeout)

    if has_ssh_connection_err_msg(result.stderr):
        raise SshConnectionError()

    lines = result.stdout.strip().split('\n')

    tree = {}
    indent_stack = []

    for line in lines:
        indent = len(line) - len(line.lstrip(' '))

        bridge_matches = list(re.finditer(r'([0-9a-f]{2}\.[0-9a-f])-\[([0-9a-f]{2,4})-[0-9a-f]+\]', line))

        if bridge_matches:
            while indent_stack and indent_stack[-1][0] >= indent:
                indent_stack.pop()

            for bridge_match in bridge_matches:
                dev_func = bridge_match.group(1)
                bus_start = bridge_match.group(2).zfill(2)
                bridge_id = f'bridge_{bus_start}_{dev_func}'

                parent = indent_stack[-1][1] if indent_stack else None
                tree[bridge_id] = {'parent': parent, 'children': [], 'is_efa': False, 'slot': None}

                if parent:
                    if parent not in tree:
                        tree[parent] = {'parent': None, 'children': [], 'is_efa': False, 'slot': None}
                    tree[parent]['children'].append(bridge_id)

                indent_stack.append((indent, bridge_id))

        leaf_match = re.search(r'[+\\]-([0-9a-f]{2}\.[0-9a-f])-\[([0-9a-f]{2,4})\]----([0-9a-f]{2}\.[0-9a-f])', line)
        if leaf_match:
            slot = leaf_match.group(1)
            bus = leaf_match.group(2).zfill(2)
            dev_func = leaf_match.group(3)
            bdf = f'0000:{bus}:{dev_func}'
            is_efa = 'Amazon' in line and ('EFA' in line or 'efa' in line)

            parent = indent_stack[-1][1] if indent_stack else None
            tree[bdf] = {'parent': parent, 'children': [], 'is_efa': is_efa, 'slot': slot}

            if parent:
                if parent not in tree:
                    tree[parent] = {'parent': None, 'children': [], 'is_efa': False, 'slot': None}
                tree[parent]['children'].append(bdf)

    return tree


# Neuron device BDF to EFA BDF mapping for trn2 instance types.
# Each neuron device index maps to a specific EFA NIC BDF.
_TRN2_ND_TO_EFA_BDF = {
    0: "0000:c9:00.0",
    1: "0000:b4:00.0",
    2: "0000:b3:00.0",
    3: "0000:ca:00.0",
    4: "0000:6c:00.0",
    5: "0000:57:00.0",
    6: "0000:56:00.0",
    7: "0000:6d:00.0",
    8: "0000:98:00.0",
    9: "0000:83:00.0",
    10: "0000:82:00.0",
    11: "0000:99:00.0",
    12: "0000:f5:00.0",
    13: "0000:e0:00.0",
    14: "0000:df:00.0",
    15: "0000:f6:00.0",
}


@functools.lru_cache(10)
@retry(retry_on_exception=is_ssh_connection_error, stop_max_attempt_number=3, wait_fixed=5000)
def get_neuron_ls_output(ip):
    """
    Run neuron-ls -j on the remote host and return the parsed JSON output.
    Results are cached per ip to avoid repeated SSH calls.
    """
    proc = subprocess.run("ssh {} /opt/aws/neuron/bin/neuron-ls -j".format(ip),
                          shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                          timeout=60, encoding="utf-8")

    if has_ssh_connection_err_msg(proc.stderr):
        raise SshConnectionError()

    if proc.returncode != 0:
        return []

    return json.loads(proc.stdout)


def get_efa_device_name_for_neuron_core(ip, neuron_core_id):
    """
    Return the EFA device name for a given neuron core ID on trn2 instances.

    Uses the neuron-ls output to find which neuron device the given core is in,
    then looks up the corresponding EFA NIC BDF from the static mapping table
    and resolves it to an rdma device name.

    Falls back to round-robin assignment for non-trn2 instance types.
    """
    neuron_devices = get_neuron_ls_output(ip)
    if not neuron_devices:
        return get_efa_device_name_for_hmem_device(ip, 0, 1)

    instance_type = neuron_devices[0].get("instance_type", "")

    if instance_type.startswith("trn2"):
        # For trn2, use the static BDF mapping
        for device in neuron_devices:
            if neuron_core_id in device.get("neuroncore_ids", []):
                nd_index = device["neuron_device"]
                efa_bdf = _TRN2_ND_TO_EFA_BDF.get(nd_index)
                assert(efa_bdf is not None)
                rdma_name = get_rdma_core_name_for_efa_nic(ip, efa_bdf)
                assert(rdma_name)
                return rdma_name

    # For non-trn2, fall back to round-robin based on device index
    # TODO: Implement topology aware NIC selection similar to CUDA
    for device in neuron_devices:
        if neuron_core_id in device.get("neuroncore_ids", []):
            return get_efa_device_name_for_hmem_device(
                ip, device["neuron_device"], len(neuron_devices))
    return get_efa_device_name_for_hmem_device(ip, 0, len(neuron_devices))


class CudaMemorySupport(IntEnum):
    NOT_INITIALIZED = -1
    NOT_SUPPORTED = 0
    DMA_BUF_ONLY = 1
    GDR_ONLY = 2
    DMABUF_GDR_BOTH = 3

    def __str__(self):
        return self.name

@retry(retry_on_exception=is_ssh_connection_error, stop_max_attempt_number=3, wait_fixed=5000)
def get_cuda_memory_support(cmdline_args, ip):
    """
    Execute check_dmabuf binary to determine CUDA memory support capabilities.
    
    Args:
        cmdline_args: Command line arguments containing binpath, timeout, provider, and environments.
        ip: IP address or hostname of the target machine.
        
    Returns:
        CudaMemorySupport: Enum value indicating hardware CUDA memory support type.
        
    Notes:
        - Executes check_dmabuf binary remotely via SSH with timeout
        - Maps return code directly to CudaMemorySupport enum values
        - Returns UNKNOWN for negative return codes indicating errors
        - Retries on SSH connection errors up to 3 times
    """
    binpath = cmdline_args.binpath or ""
    cmd = "timeout " + str(cmdline_args.timeout) \
          + " " + os.path.join(binpath, "cuda_check_dmabuf") \
          + " -p " + cmdline_args.provider
    if cmdline_args.environments:
        cmd = cmdline_args.environments + " " + cmd

    proc = subprocess.run("ssh {} {}".format(ip, cmd),
               stdout=subprocess.PIPE,
               stderr=subprocess.STDOUT,
               shell=True,
               universal_newlines=True)

    if has_ssh_connection_err_msg(proc.stdout):
        raise SshConnectionError()

    if proc.returncode < 0:
        return CudaMemorySupport.NOT_SUPPORTED

    print(f"The ssh return is {proc}")
    rc = proc.returncode
    if rc not in (CudaMemorySupport.NOT_SUPPORTED,
                  CudaMemorySupport.DMA_BUF_ONLY,
                  CudaMemorySupport.GDR_ONLY,
                  CudaMemorySupport.DMABUF_GDR_BOTH):
        print(f"[warn] check_dmabuf returned unexpected code {rc}, treating as NOT_INITIALIZED")
        return CudaMemorySupport.NOT_INITIALIZED

    return CudaMemorySupport(rc)

def efa_run_client_server_test(cmdline_args, executable, iteration_type,
                               completion_semantic, memory_type, message_size,
                               warmup_iteration_type=None, timeout=None,
                               completion_type="queue", fabric=None,
                               additional_env='',
                               might_fail=False):
    if timeout is None:
        timeout = cmdline_args.timeout

    # It is observed that cuda tests requires larger time-out limit to test all
    # message sizes (especailly when running with multiple workers).
    if "cuda" in memory_type:
        timeout = max(1000, timeout)

    # For host only test, we do not need to initialize non-system hmem ifaces
    if memory_type == "host_to_host":
        additional_env += ' FI_HMEM=system'

    test = ClientServerTest(cmdline_args, executable, iteration_type,
                            completion_semantic=completion_semantic,
                            datacheck_type="with_datacheck",
                            message_size=message_size,
                            memory_type=memory_type,
                            timeout=timeout,
                            warmup_iteration_type=warmup_iteration_type,
                            completion_type=completion_type, fabric=fabric,
                            additional_env=additional_env,
                            might_fail=might_fail)
    test.run()

@retry(retry_on_exception=is_ssh_connection_error, stop_max_attempt_number=3, wait_fixed=5000)
def efa_retrieve_hw_counter_value(hostname, hw_counter_name, efa_device_name=None):
    """
    retrieve the value of EFA's hardware counter
    hostname: a host that has efa
    hw_counter_name: EFA hardware counter name. Options are: lifespan, rdma_read_resp_bytes, rdma_read_wrs,recv_wrs,
                     rx_drops, send_bytes, tx_bytes, rdma_read_bytes,  rdma_read_wr_err, recv_bytes, rx_bytes, rx_pkts, send_wrs, tx_pkts
    efa_device_name: Name of the EFA device. Corresponds to the name of the EFA device's directory
    return: an integer that is sum of all EFA device's counter
    """

    if efa_device_name:
        efa_device_dir = efa_device_name
    else:
        efa_device_dir = '*'

    command = 'ssh {} cat "/sys/class/infiniband/{}/ports/*/hw_counters/{}"'.format(hostname, efa_device_dir, hw_counter_name)
    process = subprocess.run(command, shell=True, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE, encoding="utf-8")
    if process.returncode != 0:
        if process.stderr and has_ssh_connection_err_msg(process.stderr):
            print("encountered ssh connection issue")
            raise SshConnectionError()
        # this can happen when OS is using older version of EFA kernel module
        return None

    linelist = process.stdout.split()
    sumvalue = 0
    for strvalue in linelist:
        sumvalue += int(strvalue)
    return sumvalue

def has_gdrcopy(hostname):
    """
    determine whether a host has gdrcopy installed
    hostname: a host
    return: a boolean
    """
    command = "ssh {} /bin/bash --login -c lsmod | grep gdrdrv".format(hostname)
    process = subprocess.run(command, shell=True, check=False, stdout=subprocess.PIPE)
    return process.returncode == 0

@retry(retry_on_exception=is_ssh_connection_error, stop_max_attempt_number=3, wait_fixed=5000)
def has_rdma(cmdline_args, operation):
    """
    determine whether both client/server sides have rdma <operation> enabled in efa device
    cmdline_args: command line argument object, including server/client id, and timeout
    operation: rdma operation name, allowed values are read and write
    return: a boolean
    """
    return _has_rdma_cached(cmdline_args.server_id, cmdline_args.client_id,
                            cmdline_args.binpath, cmdline_args.timeout,
                            cmdline_args.environments, operation)


@functools.lru_cache(maxsize=None)
def _has_rdma_cached(server_id, client_id, binpath, timeout, environments, operation):
    assert operation in ["read", "write", "writedata"]
    binpath = binpath or ""
    cmd = "timeout " + str(timeout) \
          + " " + os.path.join(binpath, f"fi_efa_rdma_checker -o {operation}")
    if environments:
        cmd = environments + " " + cmd
    server_proc = subprocess.run("ssh {} {}".format(server_id, cmd),
               stdout=subprocess.PIPE,
               stderr=subprocess.STDOUT,
               shell=True,
               universal_newlines=True)
    if has_ssh_connection_err_msg(server_proc.stdout):
        raise SshConnectionError()

    client_proc = subprocess.run("ssh {} {}".format(client_id, cmd),
               stdout=subprocess.PIPE,
               stderr=subprocess.STDOUT,
               shell=True,
               universal_newlines=True)
    if has_ssh_connection_err_msg(client_proc.stdout):
        raise SshConnectionError()

    return (server_proc.returncode == 0 and client_proc.returncode == 0)

def efa_retrieve_gid(hostname):
    """
    return the GID of efa device on a host
    hostname: a host
    return: a string if the host has efa device,
            None otherwise
    """
    command = "ssh {} ibv_devinfo  -v | grep GID | awk '{{print $NF}}' | head -n 1".format(hostname)
    try:
        process = subprocess.run(command, shell=True, check=True, stdout=subprocess.PIPE)
    except subprocess.CalledProcessError:
        # this can happen on instance without EFA device
        return None

    return process.stdout.decode("utf-8").strip()

@functools.lru_cache(10)
@retry(retry_on_exception=is_ssh_connection_error, stop_max_attempt_number=3, wait_fixed=5000)
def get_efa_device_names(server_id):
    timeout = 60

    # This command returns a list of EFA devices names
    command = "ssh {} 'fi_info -p efa -t FI_EP_RDM -f efa || /opt/amazon/efa/bin/fi_info -p efa -t FI_EP_RDM -f efa' | grep domain".format(server_id)
    proc = subprocess.run(command, shell=True,
                          stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                          encoding="utf-8", timeout=timeout)

    if has_ssh_connection_err_msg(proc.stderr):
        raise SshConnectionError()

    devices = []
    stdouts = proc.stdout.strip().split("\n")
    #
    # Example out of fi_info -p efa -t FI_EP_RDM -f efa | grep domain are like the following:
    #     domain: rdmap0s31-rdm
    #     ...
    #
    # Extract the device name from the second column of the stdout
    # by removing the -rdm suffix
    for line in stdouts:
        parts = line.split()
        if len(parts) > 1 and "rdm" in parts[1]:
            devices.append(parts[1].replace("-rdm", ""))
    return devices

@functools.lru_cache(16)
@retry(retry_on_exception=is_ssh_connection_error, stop_max_attempt_number=3, wait_fixed=5000)
def get_gpu_bdf(server_id, gpu_index):
    timeout = 60

    result = subprocess.run(
        f"ssh {server_id} nvidia-smi --query-gpu=pci.bus_id --format=csv,noheader --id {str(gpu_index)}",
            shell=True,
            stdout=subprocess.PIPE, stderr=subprocess.PIPE,
            encoding="utf-8", timeout=timeout)

    if has_ssh_connection_err_msg(result.stderr):
        raise SshConnectionError()

    bdf_raw = result.stdout.strip().lower()
    bdf_parts = bdf_raw.split(':')
    if len(bdf_parts) != 3:
        return None
    domain = bdf_parts[0]
    if len(domain) < 4:
        domain = domain.zfill(4)
    else:
        domain = domain[-4:].zfill(4)
    bus = bdf_parts[1].zfill(2)
    dev_func = bdf_parts[2]
    return f'{domain}:{bus}:{dev_func}'

@functools.lru_cache(32)
@retry(retry_on_exception=is_ssh_connection_error, stop_max_attempt_number=3, wait_fixed=5000)
def get_rdma_core_name_for_efa_nic(server_id, bdf):
    timeout = 60
    result = subprocess.run(
        f"ssh {server_id} ls /sys/bus/pci/devices/{bdf}/infiniband",
        shell=True,
        stdout=subprocess.PIPE, stderr=subprocess.PIPE,
        encoding="utf-8", timeout=timeout)

    if has_ssh_connection_err_msg(result.stderr):
        raise SshConnectionError()

    if result.returncode == 0 and result.stdout.strip():
        return result.stdout.strip().split()[0]
    return None

def get_closest_efa_nics_for_gpu(server_id, gpu_index, pcie_tree):
    """
    BFS traversal to find the closest EFA NICs for a given GPU
    """
    gpu_bdf = get_gpu_bdf(server_id, gpu_index)

    if gpu_bdf not in pcie_tree:
        return []

    visited = set()
    queue = deque([(gpu_bdf, 0)])
    closest_distance = None
    closest_efas = []

    while queue:
        node, distance = queue.popleft()

        if node in visited:
            continue
        visited.add(node)

        if closest_distance is not None and distance > closest_distance:
            break

        if node in pcie_tree:
            if pcie_tree[node]['is_efa'] and node != gpu_bdf:
                rdma_name = get_rdma_core_name_for_efa_nic(server_id, node)
                if rdma_name:
                    if closest_distance is None:
                        closest_distance = distance
                    if distance == closest_distance:
                        closest_efas.append(rdma_name)

            # Add parent
            if pcie_tree[node]['parent'] and pcie_tree[node]['parent'] not in visited:
                queue.append((pcie_tree[node]['parent'], distance + 1))

            # Add children
            for child in pcie_tree[node]['children']:
                if child not in visited:
                    queue.append((child, distance + 1))

    return closest_efas

def get_efa_device_name_for_hmem_device(ip, hmem_device_id, num_hmem_devices):
    # this function implemented a simple way to find the closest EFA device for a given
    # hmem device. It assumes EFA devices names are in order (which is usually true but not always)
    #
    # For example, one a system with 8 accelerators and 4 EFA devices, this function would
    # for accelerator 0 and 1, return EFA device 0
    # for accelerator 2 and 3, return EFA device 1
    # for accelerator 4 and 5, return EFA device 2
    # for accelerator 6 and 7, return EFA device 3
    efa_devices = get_efa_device_names(ip)
    num_efa = len(efa_devices)
    return efa_devices[(hmem_device_id * num_efa) // num_hmem_devices]

@functools.lru_cache(10)
def get_efa_device_name_for_cuda_device(ip, cuda_device_id, num_cuda_devices):
    """
    Traverse the PCIe hierarchy and find the closest EFA NIC for a given GPU
    If the PCIe hierarchy traversal fails, fallback to a simple round robin
    """
    efa_devices = get_efa_device_names(ip)
    num_efa = len(efa_devices)

    pcie_tree = parse_lspci_tree(ip)
    closest_nics = get_closest_efa_nics_for_gpu(ip, cuda_device_id, pcie_tree)

    if not closest_nics:
        return get_efa_device_name_for_hmem_device(ip, cuda_device_id, num_cuda_devices)

    # Each GPU can use a different EFA NIC
    if num_efa >= num_cuda_devices:
        return closest_nics[0]

    # Multiple GPUs share the same EFA NIC
    return closest_nics[(cuda_device_id * num_efa) // num_cuda_devices]

@retry(retry_on_exception=is_ssh_connection_error, stop_max_attempt_number=3, wait_fixed=5000)
def support_cq_interrupts(hostname):
    """
    determine whether an EFA device supports fi_cq_sread
    This is based on checking if the device has more than 1 completion vector,
    which indicates support for CQ interrupts
    hostname: a host that has efa
    return: True if sread is supported, False otherwise, None if unable to determine
    """
    try:
        efa_devices = get_efa_device_names(hostname)
        if not efa_devices:
            return None
        efa_device_name = efa_devices[0]
    except Exception:
        return None

    command = 'ssh {} ibv_devinfo -d {} -v | grep "num_comp_vectors"'.format(hostname, efa_device_name)
    process = subprocess.run(command, shell=True, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE, encoding="utf-8")

    if process.returncode != 0:
        if process.stderr and has_ssh_connection_err_msg(process.stderr):
            print("encountered ssh connection issue")
            raise SshConnectionError()
        return None

    # Example output of ibv_devinfo -d rdmap85s0 -v | grep "num_comp_vectors"
    # num_comp_vectors:		32
    lines = process.stdout.split('\n')
    for line in lines:
        line = line.strip()
        if 'num_comp_vectors' in line:
            parts = line.split(':')
            if len(parts) >= 2:
                try:
                    comp_vectors = int(parts[1].strip())
                    # Support sread if num_comp_vectors > 1
                    return comp_vectors > 1
                except ValueError:
                    continue

    return None
