#!/usr/bin/python3

import pycol
from clicmd import *
import sys
import traceback
import textwrap
import xml.etree.ElementTree as ET
import xml.dom.minidom
import pydoc

default_conf = """
<configuration name="test_cli">
  <interfaces>
    <interface>
      <name>gre1</name>
      <description>GRE Tunnel</description>
      <encapsulation>
        <gre-tunnel>
          <gre-port>gre1</gre-port>
          <local-endpoint>1.1.1.1</local-endpoint>
          <remote-endpoint>1.1.1.2</remote-endpoint>
        </gre-tunnel>
      </encapsulation>
    </interface>
    <interface>
      <name>gre2</name>
      <description>GRE Tunnel</description>
      <encapsulation>
        <gre-tunnel>
          <gre-port>gre2</gre-port>
          <local-endpoint>2.2.2.1</local-endpoint>
          <remote-endpoint>2.2.2.2</remote-endpoint>
        </gre-tunnel>
      </encapsulation>
    </interface>
  </interfaces>
</configuration>
"""

xmlconf = ET.XML(default_conf)
cur_xmlconf = xmlconf
ip_list = []

# from http://stackoverflow.com/questions/749796/pretty-printing-xml-in-python
def xml_indent(elem, level=0):
    i = "\n" + level*"  "
    if len(elem):
        if not elem.text or not elem.text.strip():
            elem.text = i + "  "
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
        for elem in elem:
            xml_indent(elem, level+1)
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
    else:
        if level and (not elem.tail or not elem.tail.strip()):
            elem.tail = i

def text_indent(s, indent = 0, end = "\n"):
    return " " * indent + s + end

def help_callback(cmdline, kvargs):
    cmd_wr = textwrap.TextWrapper(initial_indent="  ",
                                   subsequent_indent="  ")
    token_wr = textwrap.TextWrapper(initial_indent="  ",
                                   subsequent_indent="      ")
    s = ""
    for c in cmdline.context:
        s += "%s\n"%(c.to_expr())
        lines = c.get_help().split("\n")
        # first line is the command + its help
        s += ("\n".join(cmd_wr.wrap(lines[0]))) + "\n"
        for l in lines[1:]:
            s += ("\n".join(token_wr.wrap(l))) + "\n"
    pydoc.pager(s)

# used in main and gre context
help_cmd = CmdBuilder(
    expr = "help",
    token_desc = {
        "help": TextCliCmd("help", help_str = "Show the help."),
    },
    cb = help_callback,
    help_str = "help for command toto"
)

def get_interface_text_conf(conf, indent = 0):
    s = ""
    local = conf.find("encapsulation/gre-tunnel/local-endpoint")
    remote = conf.find("encapsulation/gre-tunnel/remote-endpoint")
    if local != None and remote != None:
        bind_str = "bind %s %s"%(local.text, remote.text)
        vrfid = conf.find("encapsulation/gre-tunnel/link-vrf-id")
        if vrfid != None:
            bind_str += " link-vrfid %s"%(vrfid.text)
        s += text_indent(bind_str, indent)
    key = conf.find("encapsulation/gre-tunnel/key")
    if key != None:
        if key.text == None:
            s += text_indent("key enable", indent)
        else:
            s += text_indent("key %s"%(key.text), indent)
    checksum = conf.find("encapsulation/gre-tunnel/checksum")
    if checksum != None:
        s += text_indent("checksum enable", indent)
    return s

def get_main_text_conf(conf, indent = 0):
    s = text_indent("# GRE", indent)
    for node in xmlconf.findall("interfaces/interface"):
        s += text_indent(node.find("name").text, indent)
        s += get_interface_text_conf(node, indent=(indent + 4))
    return s

def display_cb(cmdline, kvargs):
    if "xml" in kvargs:
        xml_indent(cur_xmlconf)
        s = ET.tostring(cur_xmlconf, method="xml", encoding="unicode")
    else:
        if cmdline.context == main_ctx:
            s = get_main_text_conf(cur_xmlconf)
        else:
            s = get_interface_text_conf(cur_xmlconf)
    pydoc.pager(s)

# used in main and gre context
display_cmd = CmdBuilder(
    expr = "display [xml]",
    token_desc = {
        "display": TextCliCmd("display", help_str = "Display the " +
                              "current configuration."),
        "xml": TextCliCmd("xml", help_str = "If present, dump in XML format."),
    },
    cb = display_cb,
    help_str = "Display the current configuration."
)

def bind_cb(cmdline, kvargs):
    tunnel = cur_xmlconf.find("encapsulation/gre-tunnel")
    local = tunnel.find("local-endpoint")
    if local == None:
        local = ET.XML("<local-endpoint/>")
        tunnel.append(local)
    local.text = kvargs["<local-ip>"]
    remote = tunnel.find("remote-endpoint")
    if remote == None:
        remote = ET.XML("<remote-endpoint/>")
        tunnel.append(remote)
    remote.text = kvargs["<remote-ip>"]
    vrfid = tunnel.find("link-vrf-id")
    if vrfid != None:
        tunnel.remove(vrfid)
    if "<vrfid>" in kvargs:
        vrfid = ET.XML("<link-vrf-id/>")
        tunnel.append(vrfid)
        vrfid.text = str(kvargs["<vrfid>"])

def key_cb(cmdline, kvargs):
    tunnel = cur_xmlconf.find("encapsulation/gre-tunnel")
    key = tunnel.find("key")
    if key != None:
        tunnel.remove(key)
    if "disable" in kvargs:
        return
    key = ET.XML("<key/>")
    tunnel.append(key)
    if "<key>" in kvargs:
        key.text = str(kvargs["<key>"])

def checksum_cb(cmdline, kvargs):
    # XXX handle input/output
    tunnel = cur_xmlconf.find("encapsulation/gre-tunnel")
    checksum = tunnel.find("checksum")
    if checksum != None:
        tunnel.remove(checksum)
    if "disable" in kvargs:
        return
    checksum = ET.XML("<checksum/>")
    tunnel.append(checksum)

def exit_cb(cmdline, kvargs):
    global cur_xmlconf, xmlconf
    cmdline.set_context("main")
    cmdline.set_prompt("router{} ")
    cur_xmlconf = xmlconf

gre_ctx = PycolContext([
    CmdBuilder(
        expr = "bind <local-ip> <remote-ip> [link-vrfid <vrfid>]",
        token_desc = {
            "bind": TextCliCmd("bind", help_str = "Bind the tunnel."),
            "<local-ip>": IPv4CliCmd(help_str = "The local IPv4 address (A.B.C.D)."),
            "<remote-ip>": IPv4CliCmd(help_str = "The remote IPv4 address (A.B.C.D)."),
            "link-vrfid": TextCliCmd("link-vrfid", help_str = ""),
            "<vrfid>": IntCliCmd(help_str = "Link vrfid."),
        },
        cb = bind_cb,
        help_str = "Bind the GRE tunnel on local and remote addresses."
    ),
    CmdBuilder(
        expr = "checksum|checksum-input|checksum-output enable|disable",
        token_desc = {
            "checksum": TextCliCmd("checksum", help_str = ""),
            "checksum-input": TextCliCmd("checksum-input", help_str = ""),
            "checksum-output": TextCliCmd("checksum-output", help_str = ""),
            "enable": TextCliCmd("enable", help_str = ""),
            "disable": TextCliCmd("disable", help_str = ""),
        },
        cb = checksum_cb,
        help_str = "Enable or disable GRE checksum."
    ),
    CmdBuilder(
        expr = "key <key>|enable|disable",
        token_desc = {
            "key": TextCliCmd("key", help_str = ""),
            "<key>": IntCliCmd(val_min = 0, val_max = (1<<32) - 1,
                               help_str = "The GRE key."),
            "enable": TextCliCmd("enable", help_str = ""),
            "disable": TextCliCmd("disable", help_str = ""),
        },
        cb = key_cb,
        help_str = "Enable or disable GRE checksum."
    ),
    CmdBuilder(
        expr = "exit",
        token_desc = {
            "exit": TextCliCmd("exit", help_str = "Exit from GRE context."),
        },
        cb = exit_cb,
        help_str = "Exit from GRE context."
    ),
    display_cmd,
    help_cmd,
])

def list_gre_ctx():
    return [ node.text for node in \
             xmlconf.findall("interfaces/interface/encapsulation/gre-tunnel/gre-port") ]

def enter_gre(ifname):
    global xmlconf, cur_xmlconf
    for node in xmlconf.findall("interfaces/interface"):
        if node.find("name").text == ifname:
            cur_xmlconf = node
            break
    cmdline.set_prompt("router{%s} "%(ifname))
    cmdline.set_context("gre")

def enter_gre_cb(cmdline, kvargs):
    ifname = kvargs["<gre-ctx>"]
    return enter_gre(ifname)

def create_gre_cb(cmdline, kvargs):
    global xmlconf, cur_xmlconf

    # create the context first
    ifname = kvargs["<new-gre-ctx>"]
    xml_str = """
          <interface>
            <name>{name}</name>
            <description>GRE Tunnel</description>
            <encapsulation>
              <gre-tunnel>
                <gre-port>{name}</gre-port>
              </gre-tunnel>
            </encapsulation>
          </interface>
        """.format(name = ifname)
    node = ET.XML(xml_str)
    xmlconf.find("interfaces").append(node)
    return enter_gre(ifname)

def delete_gre(cmdline, kvargs):
    interfaces = xmlconf.find("interfaces")
    for interface in interfaces.findall("interface"):
        if interface.find("name").text == kvargs["<gre-ctx>"]:
            interfaces.remove(interface)
            break

main_ctx = PycolContext([
    CmdBuilder(
        expr = "<gre-ctx>",
        token_desc = {
            "<gre-ctx>": ChoiceCliCmd(choice = list_gre_ctx,
                                      help_str = "An existing gre context."),
        },
        cb = enter_gre_cb,
        help_str = "Enter an existing GRE context."
    ),
    CmdBuilder(
        expr = "<new-gre-ctx>",
        token_desc = {
            "<new-gre-ctx>": RegexpCliCmd("gre[0-9]+",
                                          help_str = "A new gre context (format " +
                                          "is gre[0-9]+)."),
        },
        cb = create_gre_cb,
        help_str = "Create and enter a new GRE context."
    ),
    CmdBuilder(
        expr = "delete <gre-ctx>",
        token_desc = {
            "delete": TextCliCmd("delete", help_str = ""),
            "<gre-ctx>": ChoiceCliCmd(choice = list_gre_ctx,
                                      help_str = "An existing gre context."),
        },
        cb = delete_gre,
        help_str = "Delete a GRE context."
    ),
    display_cmd,
    help_cmd,
])


cmdline = pycol.Cmdline()
cmdline.add_context("main", main_ctx)
cmdline.add_context("gre", gre_ctx)
cmdline.set_context("main")
cmdline.set_prompt("router{} ")
cmdline.input_loop()