#!/usr/bin/env python3
"""
....................................

     JET PROPULSION LABORATORY
------------------------------------
         ___  _______  ___
        |   ||       ||   |
        |   ||    _  ||   |
        |   ||   |_| ||   |
     ___|   ||    ___||   |___
    |       ||   |    |       |
    |_______||___|    |_______|

------------------------------------
 CALIFORNIA INSTITUTE OF TECHNOLOGY
------------------------------------

*****************************************************************************
 Title: Net model generation script for ION configuration
 Author: Nate Richard
 Modified: 01/06/2024
 Company: JPL
 Date:   11/17/2020

 File: ion_config
 Description:
           Script that takes a simple JSON model file and converts it to an
           JSON net model file of use in generating ION config files.
           Python 3.8.16

*****************************************************************************
"""

import argparse
import json
from pathlib import Path
from typing import Dict, List, Union
from typing_extensions import Literal, TypedDict, Never


class NetModelParmDict(TypedDict):
    """Dictionary for parameters to construct Net Model."""

    IP: Union[List[str], str]
    NODE: str
    DEST: List[str]
    PROTOCOL: List[str]
    RATE: List[int]
    SERVICES: Union[List[str], List[Never]]


class NetHostsDict(TypedDict):
    """Dictionary for host in Net Model."""

    hostName: str
    hostDesc: str
    ipAddrs: List[str]


class NetNodesDict(TypedDict):
    """Dictionary for nodes in Net Model."""

    nodeName: str
    nodeDesc: str
    nodeHost: str
    nodeType: Literal["ion", "dtn"]
    endpointID: str
    services: List[str]


class NetHopsDict(TypedDict):
    """Dictionary for net hops in Net Model."""

    hopName: str
    hopDesc: str
    fromNode: str
    fromIP: str
    toNode: str
    toIP: str
    bpLayer: Literal["ltp", "tcp", "udp", "stcp", "dccp"]
    ltpLayer: Literal["", "udp", "dccp"]
    portnum: int
    maxRate: int
    symmetric: bool


def create_netmodel(src: Dict, nm_name: str) -> str:
    """Generate a net model based of provided JSON string."""
    # Create net_hosts dictionary
    net_hosts = gen_net_hosts(src)

    # Create net_nodes dictionary
    net_nodes = gen_net_nodes(src)

    # Create net_hops dictionary
    net_hops = gen_net_hops(src)

    return {
        "netModelName": nm_name,
        "netModelDesc": "Python generated net model",
        "netHosts": net_hosts,
        "netNodes": net_nodes,
        "netHops": net_hops,
    }


def gen_net_hosts(indict: NetModelParmDict) -> NetHostsDict:
    """Combine host and node lists into a net model compatible dictionary."""
    net_hosts: NetHostsDict = {}
    for host in indict.keys():
        net_hosts[host] = {
            "hostName": host,
            "hostDesc": f"{host} Host",
            "ipAddrs": [indict[host]["IP"]],
        }

    return net_hosts


def gen_net_nodes(indict: NetModelParmDict) -> NetNodesDict:
    """Combine node parameters into a single dictionary for net model."""
    net_nodes: NetNodesDict = {}
    for node, parms in indict.items():
        net_nodes[node] = {
            "nodeName": node,
            "nodeDesc": f"{node} ION node on {node}",
            "nodeHost": node,
            "nodeType": "ion",
            "endpointID": parms["NODE"],
            "services": parms["SERVICES"],
        }

    return net_nodes


def gen_net_hops(indict: NetModelParmDict) -> NetHopsDict:
    """Combine node and link parameters to define hops between nodes for net model."""
    net_hops: NetHopsDict = {}
    for node, parms in indict.items():
        dest = parms["DEST"]
        proto = parms["PROTOCOL"]
        rate = parms["RATE"]
        # Assume we have multiple destinations but only 1 protocol
        # Assume data rates are defined for each destination even if they're
        # the same
        if len(dest) > 1 and len(proto) == 1:
            for dest_node, speed in zip(dest, rate):
                hop = f"{node}2{dest_node}"
                if proto[0] == "ltp":
                    layer = "udp"
                    sym = False
                    portnum = 1113
                else:
                    layer = ""
                    sym = True
                    portnum = 4556
                net_hops[hop] = {
                    "hopName": hop,
                    "hopDesc": f"Link from {node} to {dest_node}",
                    "fromNode": node,
                    "fromIP": parms["IP"],
                    "toNode": dest_node,
                    "toIP": indict[dest_node]["IP"],
                    "bpLayer": proto[0],
                    "ltpLayer": layer,
                    "portNum": portnum,
                    "maxRate": int(speed),
                    "symmetric": sym,
                }

        # Assume we have multiple destinations and multiple protocols
        # Assumes each destination has a unique data rate
        elif len(dest) > 1 and len(proto) > 1:
            for dest_node, bpproto, speed in zip(dest, proto, rate):
                hop = f"{node}2{dest_node}"
                if bpproto == "ltp":
                    layer = "udp"
                    sym = False
                    portnum = 1113
                else:
                    layer = ""
                    sym = True
                    portnum = 4556
                net_hops[hop] = {
                    "hopName": hop,
                    "hopDesc": f"Link from {node} to {dest_node}",
                    "fromNode": node,
                    "fromIP": parms["IP"],
                    "toNode": dest_node,
                    "toIP": indict[dest_node]["IP"],
                    "bpLayer": bpproto,
                    "ltpLayer": layer,
                    "portNum": portnum,
                    "maxRate": int(speed),
                    "symmetric": sym,
                }

        # Assume we have 1 destination, 1 protocol, and it hasn't been
        # defined yet
        elif len(dest) == 1 and dest[0] != "":
            hop = f"{node}2{dest[0]}"
            if proto[0] == "ltp":
                layer = "udp"
                sym = False
                portnum = 1113
            else:
                layer = ""
                sym = True
                portnum = 4556
            net_hops[hop] = {
                "hopName": hop,
                "hopDesc": f"Link from {node} to {dest[0]}",
                "fromNode": node,
                "fromIP": parms["IP"],
                "toNode": dest[0],
                "toIP": indict[dest[0]]["IP"],
                "bpLayer": proto[0],
                "ltpLayer": layer,
                "portNum": portnum,
                "maxRate": int(rate[0]),
                "symmetric": sym,
            }

    return net_hops


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Convert simple node model into a net model."
    )
    parser.add_argument(
        "model_file", help="JSON file defining parameters to create net model"
    )
    args = parser.parse_args()

    model_name = Path(args.model_file).stem
    net_model_name = f"{model_name}-net_model.json"
    with open(args.model_file, "r", encoding="utf-8") as mfile:
        with open(net_model_name, "w", encoding="utf-8") as nfile:
            json.dump(create_netmodel(json.load(mfile), model_name), nfile, indent=4)
