#!/usr/bin/env python3

import time
import subprocess
import os
from datetime import datetime

#  GLOBAL CONFIG 
SSH_KEY = os.path.expanduser("~/.ssh/vnfkey")
SSH_OPTS = [
    "-o", "BatchMode=yes",
    "-o", "StrictHostKeyChecking=accept-new",
    "-o", "ConnectTimeout=5",
    "-n",  
    "-f",  
]

INTERVAL = 5            
LOG_DIR = "~/scripts"   

# Publisher (runs once, never killed)
PUBLISHER = {
    "name": "publisher",
    "ip":   "10.20.20.26",  # python-vnf
    "cmd":  "python3 -u ~/publisher.py",
    "port": None,
}

AGGREGATORS = [
    {
        "name": "cpp-vnf",
        "ip":   "10.20.20.14",
        "cmd":  "~/pointcloud_concatenator",
        "port": 5562,
    },
    {
        "name": "python-vnf",
        "ip":   "10.20.20.26",
        "cmd":  "python3 -u ~/diversification.py",
        "port": 5563,
    },
    {
        "name": "java-vnf",
        "ip":   "10.20.20.113",
        "cmd":  "java -cp ~/jeromq-0.5.2.jar:. PointCloudConcatenator",
        "port": 5560,
    }
]

#  SSH UTILITY 
def ssh_run(ip: str, remote_cmd: str, *, capture=False):
    cmd = ["ssh", *SSH_OPTS, "-i", SSH_KEY, f"ubuntu@{ip}", remote_cmd]
    return subprocess.run(cmd, capture_output=capture, text=True)

def free_port(ip: str, port: int | None):
    if port:
        ssh_run(ip, f"fuser -k {port}/tcp || true && sleep 1")

#  VNF STARTER 
def start_vnf(vnf: dict, kill_old: bool) -> bool:
    ip, cmd = vnf["ip"], vnf["cmd"]
    logfile = f"{LOG_DIR}/{vnf['name']}.log"

    if kill_old:
        ssh_run(ip, f"pkill -f '{cmd}' || true")
        free_port(ip, vnf["port"])

    remote = f"mkdir -p {LOG_DIR} && nohup {cmd} > {logfile} 2>&1 < /dev/null &"
    print(f"\n[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] {vnf['name']} → {ip}")
    ssh_run(ip, remote)
    print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] started.")

    # Wait 2s for PID and 3s for log content (total 5s)
    time.sleep(2)
    pid = (
        ssh_run(ip, f"pgrep -f '{cmd}' | head -n1", capture=True)
        .stdout.strip() or "n/a"
    )
    time.sleep(10)  # Give process time to generate output

    tail = (
        ssh_run(ip, f"tail -n 5 {logfile} || true", capture=True)
        .stdout.strip() or "(empty)"
    )
    print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] pid: {pid}")
    print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] last log lines:\n{tail}")
    
    alive = False
    if pid != "n/a":
        alive = ssh_run(ip, f"ps -p {pid}", capture=True).returncode == 0
    crashed = (not alive) or ("Segmentation fault" in tail)

    if crashed:
        print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}]  {vnf['name']} may have crashed.")
    return crashed

# MAIN LOOP 
print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] Controller running — switch every {INTERVAL}s")

start_vnf(PUBLISHER, kill_old=False)

while True:
    for v in AGGREGATORS:
        start_vnf(v, kill_old=True)
        print(f"[{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] sleep {INTERVAL}s\n")
        time.sleep(INTERVAL)

