#!/usr/bin/env python3
"""
iperf-benchmark-weave: Two-node bandwidth test using iperf3 over L2Bridge
"""

import sys
import time
from fabric_fablib import fablib

def start(slice_name: str, site: str = "RENC"):
    """Create, submit, configure, and run iperf3 benchmark."""
    print(f"Starting iperf benchmark weave: {slice_name}")
    
    # Create slice
    slice_obj = fablib.new_slice(name=slice_name)
    
    # Add two nodes at the same site
    server_node = slice_obj.add_node(name="iperf-server", site=site, cores=2, ram=8, disk=20)
    client_node = slice_obj.add_node(name="iperf-client", site=site, cores=2, ram=8, disk=20)
    
    # Add NICs
    server_nic = server_node.add_component(name="nic1", model="NIC_Basic")
    client_nic = client_node.add_component(name="nic1", model="NIC_Basic")
    
    # Create L2Bridge network
    network = slice_obj.add_network(name="l2-net", type="L2Bridge", 
                                    interfaces=[f"iperf-server-nic1-p1", f"iperf-client-nic1-p1"])
    
    # Submit slice
    print("Submitting slice for provisioning...")
    slice_obj.submit()
    
    # Wait for StableOK
    print("Waiting for slice to reach StableOK...")
    slice_obj.wait(wait_for="StableOK", timeout=600)
    
    # Configure networking with static IPs
    print("Configuring L2 network with static IPs...")
    server_node.execute("sudo ip addr add 192.168.100.1/24 dev enp7s0")
    server_node.execute("sudo ip link set enp7s0 up")
    client_node.execute("sudo ip addr add 192.168.100.2/24 dev enp7s0")
    client_node.execute("sudo ip link set enp7s0 up")
    
    # Verify connectivity
    print("Verifying connectivity...")
    rc, out = client_node.execute("ping -c 3 192.168.100.1")
    print(f"Ping result: {out}")
    
    # Install iperf3
    print("Installing iperf3 on both nodes...")
    server_node.execute("sudo apt-get update -y -qq && sudo apt-get install -y -qq iperf3", timeout=300)
    client_node.execute("sudo apt-get update -y -qq && sudo apt-get install -y -qq iperf3", timeout=300)
    
    # Start iperf3 server
    print("Starting iperf3 server...")
    server_node.execute("iperf3 -s -D")
    time.sleep(2)
    
    # Run iperf3 client test
    print("Running iperf3 benchmark (10 seconds)...")
    rc, output = client_node.execute("iperf3 -c 192.168.100.1 -t 10")
    print(f"\n=== IPERF3 RESULTS ===\n{output}")
    
    # Save results
    with open("iperf_results.txt", "w") as f:
        f.write(output)
    
    print("Results saved to iperf_results.txt")
    return {"slice_name": slice_name, "status": "completed", "results": output}

def stop(slice_name: str):
    """Delete the slice."""
    print(f"Stopping weave: {slice_name}")
    try:
        slice_obj = fablib.get_slice(slice_name)
        slice_obj.delete()
        print(f"Slice {slice_name} deleted")
    except Exception as e:
        print(f"Error deleting slice: {e}")
    return {"slice_name": slice_name, "status": "stopped"}

def monitor(slice_name: str):
    """Check slice status."""
    try:
        slice_obj = fablib.get_slice(slice_name)
        state = slice_obj.get_state()
        nodes = slice_obj.get_nodes()
        return {"slice_name": slice_name, "state": state, "node_count": len(nodes)}
    except Exception as e:
        return {"slice_name": slice_name, "error": str(e)}

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: iperf_benchmark.py <start|stop|monitor> [slice_name] [site]")
        sys.exit(1)
    
    action = sys.argv[1]
    slice_name = sys.argv[2] if len(sys.argv) > 2 else "iperf-benchmark-slice"
    site = sys.argv[3] if len(sys.argv) > 3 else "RENC"
    
    if action == "start":
        result = start(slice_name, site)
    elif action == "stop":
        result = stop(slice_name)
    elif action == "monitor":
        result = monitor(slice_name)
    else:
        print(f"Unknown action: {action}")
        sys.exit(1)
    
    print(result)