from fabrictestbed_extensions.fablib.attestable_switch import Attestable_Switch
from fabrictestbed_extensions.fablib.fablib import FablibManager as fablib_manager
from ipaddress import ip_address, IPv4Address, IPv4Network
from concurrent import futures
import time
from datetime import datetime
from . import custom_setup

start_time = None
end_time = None
cached_answer = None

def start_exercise(vslice, macs, ips, subnets, hosts, routers):
    global start_time
    global cached_answer
    start_time = time.time()

    tunnel_ips = {
        "h1": ["15.1.1.11"],
        "h2": ["15.1.1.12"]
    }
    file_r1 = "scripts/rules_r1.txt"
    file_r2 = "scripts/rules_r2.txt"
    file_r3 = "scripts/rules_r3.txt"
    file_r4 = "scripts/rules_r4.txt"
    file_r5 = "scripts/rules_r5.txt"
    file_r6 = "scripts/rules_r6.txt"
    file_r7 = "scripts/rules_r7.txt"
    file_r8 = "scripts/rules_r8.txt"
    file_r9 = "scripts/rules_r9.txt"
    files = {
        "r1": file_r1,
        "r2": file_r2,
        "r3": file_r3,
        "r4": file_r4,
        "r5": file_r5,
        "r6": file_r6,
        "r7": file_r7,
        "r8": file_r8,
        "r9": file_r9
    }

    teids = {}
    ctr = 1
    for hostname in hosts:
        teids[hostname] = {}
        for destname in hosts:
            if hostname == destname:
                continue
            teids[hostname][destname] = ctr
            ctr = ctr + 1

    table_host = "table_host"
    table_internet = "table_host"
    table_tunnel = "table_internet"
    forward_host = "forward_remove_tunnel"
    a_forward_router = "forward_add_tunnel"
    forward_router = "forward"
        
    with open(file_r1, 'w') as outfile:
        outfile.write("load_new_config_file router-2.json\n")
        outfile.write("swap_configs\n")
        for i in range(3):
            outfile.write(f"table_add {table_host} {forward_host} {tunnel_ips['h1'][0]} 0 {i} => {macs['h1'][0]} {macs['r1'][2]} 2\n")
            outfile.write(f"table_add {table_host} {forward_host} {tunnel_ips['h1'][0]} 1 {i} => {macs['h1'][0]} {macs['r1'][2]} 2\n")
            if i == 1:
                outfile.write(f"table_add {table_tunnel} {a_forward_router} {subnets['h2'][0]} 2 {i} => {macs['r2'][1]} {macs['r1'][0]} 0 {tunnel_ips['h1'][0]} {tunnel_ips['h2'][0]} {teids['h1']['h2']}\n")
            else:
                outfile.write(f"table_add {table_tunnel} {a_forward_router} {subnets['h2'][0]} 2 {i} => {macs['r8'][0]} {macs['r1'][1]} 1 {tunnel_ips['h1'][0]} {tunnel_ips['h2'][0]} {teids['h1']['h2']}\n")
        
    with open(file_r2, 'w') as outfile:
        outfile.write("load_new_config_file router-2.json\n")
        outfile.write("swap_configs\n")
        for i in range(3):
            outfile.write(f"table_add {table_internet} {forward_router} {tunnel_ips['h1'][0]} 0 {i} => {macs['r1'][0]} {macs['r2'][1]} 1\n")
            outfile.write(f"table_add {table_internet} {forward_router} {tunnel_ips['h1'][0]} 2 {i} => {macs['r1'][0]} {macs['r2'][1]} 1\n")
            outfile.write(f"table_add {table_internet} {forward_router} {tunnel_ips['h2'][0]} 2 {i} => {macs['r3'][1]} {macs['r2'][0]} 0\n")
            if i == 1:
                outfile.write(f"table_add {table_internet} {forward_router} {tunnel_ips['h2'][0]} 1 {i} => {macs['r9'][0]} {macs['r2'][2]} 2\n")
            else:
                outfile.write(f"table_add {table_internet} {forward_router} {tunnel_ips['h2'][0]} 1 {i} => {macs['r3'][1]} {macs['r2'][0]} 0\n")
        
    with open(file_r3, 'w') as outfile:
        outfile.write("load_new_config_file router-2.json\n")
        outfile.write("swap_configs\n")
        for i in range(3):
            outfile.write(f"table_add {table_internet} {forward_router} {tunnel_ips['h2'][0]} 1 {i} => {macs['r4'][1]} {macs['r3'][0]} 0\n")
            outfile.write(f"table_add {table_internet} {forward_router} {tunnel_ips['h1'][0]} 0 {i} => {macs['r2'][0]} {macs['r3'][1]} 1\n")
        
    with open(file_r4, 'w') as outfile:
        outfile.write("load_new_config_file router-2.json\n")
        outfile.write("swap_configs\n")
        for i in range(3):
            outfile.write(f"table_add {table_internet} {forward_router} {tunnel_ips['h2'][0]} 1 {i} => {macs['r5'][1]} {macs['r4'][0]} 0\n")
            outfile.write(f"table_add {table_internet} {forward_router} {tunnel_ips['h2'][0]} 2 {i} => {macs['r5'][1]} {macs['r4'][0]} 0\n")
            outfile.write(f"table_add {table_internet} {forward_router} {tunnel_ips['h1'][0]} 2 {i} => {macs['r3'][0]} {macs['r4'][1]} 1\n")
            if i == 1:
                outfile.write(f"table_add {table_internet} {forward_router} {tunnel_ips['h1'][0]} 0 {i} => {macs['r9'][1]} {macs['r4'][2]} 2\n")
            else:
                outfile.write(f"table_add {table_internet} {forward_router} {tunnel_ips['h1'][0]} 0 {i} => {macs['r3'][0]} {macs['r4'][1]} 1\n")
        
    with open(file_r5, 'w') as outfile:
        outfile.write("load_new_config_file router-2.json\n")
        outfile.write("swap_configs\n")
        for i in range(3):
            outfile.write(f"table_add {table_host} {forward_host} {tunnel_ips['h2'][0]} 0 {i} => {macs['h2'][0]} {macs['r1'][2]} 2\n")
            outfile.write(f"table_add {table_host} {forward_host} {tunnel_ips['h2'][0]} 1 {i} => {macs['h2'][0]} {macs['r1'][2]} 2\n")
            if i == 1:
                outfile.write(f"table_add {table_tunnel} {a_forward_router} {subnets['h1'][0]} 2 {i} => {macs['r6'][1]} {macs['r1'][0]} 0 {tunnel_ips['h2'][0]} {tunnel_ips['h1'][0]} {teids['h2']['h1']}\n")
            else:
                outfile.write(f"table_add {table_tunnel} {a_forward_router} {subnets['h1'][0]} 2 {i} => {macs['r4'][0]} {macs['r1'][1]} 1 {tunnel_ips['h2'][0]} {tunnel_ips['h1'][0]} {teids['h2']['h1']}\n")
        
    with open(file_r6, 'w') as outfile:
        outfile.write("load_new_config_file router-2.json\n")
        outfile.write("swap_configs\n")
        for i in range(3):
            outfile.write(f"table_add {table_internet} {forward_router} {tunnel_ips['h2'][0]} 0 {i} => {macs['r5'][0]} {macs['r6'][1]} 1\n")
            outfile.write(f"table_add {table_internet} {forward_router} {tunnel_ips['h2'][0]} 2 {i} => {macs['r5'][0]} {macs['r6'][1]} 1\n")
            outfile.write(f"table_add {table_internet} {forward_router} {tunnel_ips['h1'][0]} 2 {i} => {macs['r7'][1]} {macs['r6'][0]} 0\n")
            if i == 1:
                outfile.write(f"table_add {table_internet} {forward_router} {tunnel_ips['h1'][0]} 1 {i} => {macs['r9'][2]} {macs['r6'][2]} 2\n")
            else:
                outfile.write(f"table_add {table_internet} {forward_router} {tunnel_ips['h1'][0]} 1 {i} => {macs['r7'][1]} {macs['r6'][0]} 0\n")
        
    with open(file_r7, 'w') as outfile:
        outfile.write("load_new_config_file router-2.json\n")
        outfile.write("swap_configs\n")
        for i in range(3):
            outfile.write(f"table_add {table_internet} {forward_router} {tunnel_ips['h2'][0]} 0 {i} => {macs['r6'][0]} {macs['r7'][1]} 1\n")
            outfile.write(f"table_add {table_internet} {forward_router} {tunnel_ips['h1'][0]} 1 {i} => {macs['r8'][1]} {macs['r7'][0]} 0\n")
        
    with open(file_r8, 'w') as outfile:
        outfile.write("load_new_config_file router-2.json\n")
        outfile.write("swap_configs\n")
        for i in range(3):
            outfile.write(f"table_add {table_internet} {forward_router} {tunnel_ips['h1'][0]} 1 {i} => {macs['r1'][1]} {macs['r8'][0]} 0\n")
            outfile.write(f"table_add {table_internet} {forward_router} {tunnel_ips['h1'][0]} 2 {i} => {macs['r1'][1]} {macs['r8'][0]} 0\n")
            outfile.write(f"table_add {table_internet} {forward_router} {tunnel_ips['h2'][0]} 2 {i} => {macs['r7'][0]} {macs['r8'][1]} 1\n")
            if i == 1:
                outfile.write(f"table_add {table_internet} {forward_router} {tunnel_ips['h2'][0]} 0 {i} => {macs['r9'][3]} {macs['r8'][2]} 2\n")
            else:
                outfile.write(f"table_add {table_internet} {forward_router} {tunnel_ips['h2'][0]} 0 {i} => {macs['r7'][0]} {macs['r8'][1]} 1\n")
        
    with open(file_r9, 'w') as outfile:
        outfile.write("load_new_config_file router-1.json\n")
        outfile.write("swap_configs\n")
        for i in range(3):
            for hostname in ['h1', 'h2']:
                if i == 1:
                    outfile.write(f"table_add {table_internet} {forward_router} {subnets[hostname][0]} 0 {i} => {macs['r4'][2]} {macs['r9'][1]} 1\n")
                    outfile.write(f"table_add {table_internet} {forward_router} {subnets[hostname][0]} 1 {i} => {macs['r6'][2]} {macs['r9'][2]} 2\n")
                    outfile.write(f"table_add {table_internet} {forward_router} {subnets[hostname][0]} 2 {i} => {macs['r4'][2]} {macs['r9'][1]} 1\n")
                    outfile.write(f"table_add {table_internet} {forward_router} {subnets[hostname][0]} 3 {i} => {macs['r4'][2]} {macs['r9'][1]} 1\n")
                else:
                    outfile.write(f"table_add {table_internet} {forward_router} {subnets[hostname][0]} 0 {i} => {macs['r6'][2]} {macs['r9'][2]} 2\n")
                    outfile.write(f"table_add {table_internet} {forward_router} {subnets[hostname][0]} 1 {i} => {macs['r6'][2]} {macs['r9'][2]} 2\n")
                    outfile.write(f"table_add {table_internet} {forward_router} {subnets[hostname][0]} 2 {i} => {macs['r4'][2]} {macs['r9'][1]} 1\n")
                    outfile.write(f"table_add {table_internet} {forward_router} {subnets[hostname][0]} 3 {i} => {macs['r6'][2]} {macs['r9'][2]} 2\n")

    jobs = []
    for routername, router in routers.items():
        jobs.append(router.upload_file_thread(files[routername], "/home/ubuntu/rules.txt"))
    futures.wait(jobs)
    jobs = []
    for routername, router in routers.items():
        jobs.append(router.execute_thread(f"cat rules.txt | simple_switch_CLI"))
    futures.wait(jobs)
    
    overall = True
    h1 = hosts['h1']
    print('h1 pinging h2 10 times')
    for _ in range(10):
        print(f"[{datetime.now()}] hosts['h1'].ping_test(ips['h2'][0])...", end="")
        pingtest = h1.ping_test(ips['h2'][0])
        if pingtest:
            print("True")
        else:
            print("False")

    cached_answer = "r9"
    return


def finish_exercise(vslice, misconfigured_switch):
    global start_time
    global end_time
    global cached_answer
    end_time = time.time()

    correct = misconfigured_switch == cached_answer

    if correct:
        print("Answer is correct!")
    else:
        print("Answer is incorrect -- you can retry it if you wish. To retry this attempt, update the 'misconfigured_switch' value and then re-run this cell.")