from pprint import pprint
import re
from collections import defaultdict

import yaml
from scrapli import Scrapli
from scrapli.exceptions import ScrapliException
from rich.console import Console
from rich.table import Table


def connect_ssh(params, command, verbose=True):
    if verbose:
        print("Connect...", params["host"])
    try:
        with Scrapli(**params) as ssh:
            prompt = ssh.get_prompt()
            host = re.search("(\S+)[>#]", prompt).group(1)
            reply = ssh.send_command(command)
            output = reply.result
            return host, output
    except ScrapliException as error:
        print(error)


def parse_sh_ip_int_br_output(output):
    regex = (
        r"(?P<port>\S+) +\S+ +\w+ \w+ +"
        r"(?P<admin_status>administratively down|up|down) +(?P<status>up|down)"
    )
    interface_type_regex = (
        r"(?P<loopback>Loopback\d+)"
        r"|(?P<tunnel>Tunnel\d+)"
        r"|(?P<subinterface>\S+\d\.\d+)"
        r"|(?P<vlan>Vlan\d+)"
        r"|(?P<physical>\S+)"
    )
    keys = ("physical", "subinterface", "loopback", "tunnel", "vlan")
    status = ["administratively down", "down", "up"]

    devices_stats = {key: {k: [] for k in status} for key in keys}
    for m in re.finditer(regex, output):
        admin_status, status, port = m.group("admin_status", "status", "port")
        intf_type_m = re.search(interface_type_regex, port)
        intf_type = intf_type_m.lastgroup
        if admin_status == "administratively down":
            status = admin_status
        devices_stats[intf_type][status].append(port)
    return devices_stats


def collect_ports_stats(devices, command="sh ip int br"):
    stats = {}
    for device in devices:
        result = connect_ssh(device, command)
        if result:
            host, output = result
            stats[host] = parse_sh_ip_int_br_output(output)
    return stats


def get_port_stats(data):
    data_stats = {}
    for device, ports in data.items():
        data_stats[device] = {}
        for port_type, details in ports.items():
            data_stats[device][port_type] = {}
            for status, port_numbers in details.items():
                data_stats[device][port_type][status] = [str(len(port_numbers))]
    return data_stats


def filter_only_physical_ports(data):
    data_new = {}
    for device, ports in data.items():
        data_new[device] = {}
        for port_type, details in ports.items():
            if port_type == "physical":
                data_new[device][port_type] = details
    return data_new


def be_pretty(data, physical_only=True, stats_only=False):
    columns = ["Device", "Port type", "Admin Down", "Down", "Up"]

    if physical_only:
        data = filter_only_physical_ports(data)

    data_stats = get_port_stats(data)
    draw_table(data_stats, columns)
    if stats_only:
        return

    draw_table(data, columns)


def draw_table(data, columns):
    table = Table()
    for name in columns:
        table.add_column(name, justify="left")

    for device, ports in data.items():
        for index, (port_type, details) in enumerate(ports.items()):
            row = (
                device,
                port_type,
                *["\n".join(port_numbers) for port_numbers in details.values()],
            )
            if index == 0:
                table.add_row(*row)
            else:
                table.add_row("", *row[1:])

    console = Console()
    console.print(table)


if __name__ == "__main__":
    with open("devices.yaml") as f:
        devices = yaml.safe_load(f)
    stats = collect_ports_stats(devices)
    # be_pretty(stats, stats_only=True)
    be_pretty(stats, physical_only=False)
