#!/usr/bin/env python
# SPDX-License-Identifier: ISC
#
# common_check.py
#
# Copyright 2024 6WIND S.A.

#
import json
from lib import topotest


def ip_check_path_selection(router, ipaddr_str, expected, vrf_name=None):
    if vrf_name:
        cmdstr = f'show ip route vrf {vrf_name} {ipaddr_str} json'
    else:
        cmdstr = f'show ip route {ipaddr_str} json'
    try:
        output = json.loads(router.vtysh_cmd(cmdstr))
    except:
        output = {}

    ret = topotest.json_cmp(output, expected)
    if ret is None:
        num_nh_expected = len(expected[ipaddr_str][0]["nexthops"])
        num_nh_observed = len(output[ipaddr_str][0]["nexthops"])
        if num_nh_expected == num_nh_observed:
            return ret
        return "{}, prefix {} does not have the correct number of nexthops : observed {}, expected {}".format(
            router.name, ipaddr_str, num_nh_observed, num_nh_expected
        )
    return ret


def iproute2_check_path_selection(router, ipaddr_str, expected, vrf_name=None):
    if not topotest.iproute2_is_json_capable():
        return None

    if vrf_name:
        cmdstr = f'ip -json route show vrf {vrf_name} {ipaddr_str}'
    else:
        cmdstr = f'ip -json route show {ipaddr_str}'
    try:
        output = json.loads(cmdstr)
    except:
        output = []

    return topotest.json_cmp(output, expected)