#!/usr/bin/env python3

import os
import sys
import time
import datetime
import subprocess
import json

# Paths and configuration
PTP_REMOTE_CONFIG_PATH = "/opt/ptp_conf/"
PTP_CONFIG_DIR = "ptp_config/adhoc"

GM_FILE = "ptp4l_grandmaster.conf"
SV_FILE = "ptp4l_slave.conf"
MA_FILE = "ptp4l_master.conf"

PTP4L_INSTANCE = "/home/apu/wifi-ptp/ptp/ptp4l"

GRANDMASTER_LOG = "/tmp/ptp4l_ibss_grandmaster.log"
SLAVE_LOG = "/tmp/ptp4l_ibss_slave.log"
MASTER_LOG = "/tmp/ptp4l_ibss_master.log"
PHC2SYS_LOG = "/tmp/phc2sys_ibss.log"

INTERFACE = "adhoc0"

def load_ptp_config(filepath):
    with open(filepath, "r") as f:
        data = json.load(f)

    lines = []
    for section, options in data.items():
        lines.append(f"[{section}]")
        for key, value in options.items():
            lines.append(f"{key} {value}")
    return "\n".join(lines)

def get_hostname(node_id):
    return f"apu{str(node_id).zfill(2)}"

def write_ptp_config_to_node(hostname, role, conf_str):
    filename = {
        "grandmaster": GM_FILE,
        "slave": SV_FILE,
        "master": MA_FILE
    }.get(role)

    if not filename:
        raise ValueError("Invalid role for config write")

    subprocess.run(
        ["ssh", hostname, f"echo '{conf_str}' | sudo tee '{PTP_REMOTE_CONFIG_PATH}/{filename}' > /dev/null"],
        shell=False,
        check=True
    )

def start_ptp_on_node_ibss(node_id, nodes):
    hostname = get_hostname(node_id)
    min_node = min(nodes)
    max_node = max(nodes)

    if node_id == min_node:
        config = load_ptp_config(os.path.join(PTP_CONFIG_DIR, "ptp4l_grandmaster.json"))
        write_ptp_config_to_node(hostname, "grandmaster", config)
        cmd = f"{PTP4L_INSTANCE} -i {INTERFACE} -m -H -f {PTP_REMOTE_CONFIG_PATH}{GM_FILE} > {GRANDMASTER_LOG} 2>&1 &"

    elif node_id == max_node:
        config = load_ptp_config(os.path.join(PTP_CONFIG_DIR, "ptp4l_bc_slave.json"))
        write_ptp_config_to_node(hostname, "slave", config)
        cmd = f"{PTP4L_INSTANCE} -i {INTERFACE} -m -H -f {PTP_REMOTE_CONFIG_PATH}{SV_FILE} > {SLAVE_LOG} 2>&1 &"

    else:
        config_slave = load_ptp_config(os.path.join(PTP_CONFIG_DIR, "ptp4l_bc_slave.json"))
        config_master = load_ptp_config(os.path.join(PTP_CONFIG_DIR, "ptp4l_bc_master.json"))
        write_ptp_config_to_node(hostname, "slave", config_slave)
        write_ptp_config_to_node(hostname, "master", config_master)

        cmd_slave = f"{PTP4L_INSTANCE} -i {INTERFACE} -m -H -f {PTP_REMOTE_CONFIG_PATH}{SV_FILE} > {SLAVE_LOG} 2>&1 &"
        cmd_phc2sys = f"phc2sys -s {INTERFACE} -c CLOCK_REALTIME -O 0 -m > {PHC2SYS_LOG} 2>&1 &"
        cmd_master = f"{PTP4L_INSTANCE} -i {INTERFACE} -m -H -f {PTP_REMOTE_CONFIG_PATH}{MA_FILE} > {MASTER_LOG} 2>&1 &"
        cmd = f"{cmd_slave} {cmd_phc2sys} {cmd_master}"

    print(f"[{hostname}] Starting PTP processes")
    subprocess.run(["ssh", hostname, cmd], check=True)

def stop_ptp_on_nodes(nodes):
    for node_id in nodes:
        hostname = get_hostname(node_id)
        subprocess.run(["ssh", hostname, "pkill -x ptp4l || true"], check=False)
        subprocess.run(["ssh", hostname, "pkill -x phc2sys || true"], check=False)

def collect_logs(nodes, output_path):
    with open(output_path, "w") as outfile:
        for node_id in nodes:
            hostname = get_hostname(node_id)
            for log in [GRANDMASTER_LOG, SLAVE_LOG, MASTER_LOG, PHC2SYS_LOG]:
                try:
                    result = subprocess.run([
                        "ssh", hostname, f"cat {log}"
                    ], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
                    outfile.write(f"===== BEGIN LOG FROM {hostname} [{log}] =====\n")
                    outfile.write(result.stdout)
                    outfile.write(f"\n===== END LOG FROM {hostname} =====\n\n")
                except subprocess.CalledProcessError:
                    continue

def main():
    if len(sys.argv) < 3:
        print("Usage: run_ptp_ibss_experiment.py <duration_seconds> <node_numbers...>")
        sys.exit(1)

    duration = int(sys.argv[1])
    nodes = [int(n) for n in sys.argv[2:]]

    print(f"Running PTP IBSS experiment for {duration}s on nodes: {nodes}")
    stop_ptp_on_nodes(nodes)

    for node_id in nodes:
        start_ptp_on_node_ibss(node_id, nodes)

    print(f"Sleeping {duration}s for experiment run...")
    time.sleep(duration)

    print("Stopping PTP processes and collecting logs...")
    stop_ptp_on_nodes(nodes)

    log_filename = datetime.datetime.now().strftime("ptp_ibss_%Y%m%d_%H%M%S.log")
    collect_logs(nodes, log_filename)
    print(f"Logs saved to {log_filename}")

if __name__ == "__main__":
    main()