#!/usr/bin/env python3
"""
PTP over Wi‑Fi (two nodes)

- Exactly two nodes are used and defined at the top of the file.
- You can (and must) set the Wi‑Fi interface and PHC (/dev/ptpX) per node here.
- ptp4l configuration files/paths remain the same as in your Ethernet script.
- Keeps your "no CLOCK_REALTIME" option as a toggle (default: False). When disabled,
  we do not touch CLOCK_REALTIME at all; when enabled, we sync CLOCK_REALTIME from PHC.

Usage:
    ./run_ptp_wifi_two_nodes.py <duration_seconds>

Example:
    ./run_ptp_wifi_two_nodes.py 120

Notes:
- Assumes passwordless SSH to both nodes.
- Logs are pulled after the run into ./logs/<timestamped-folder>/<host>/
- ptp4l binary path can be switched below.
"""

import os
import sys
import time
import datetime
import subprocess
from typing import Tuple

# ==============================
# --- USER SETTINGS (edit)  ---
# ==============================

GM_HOST = "apu00"
SL_HOST = "apu01"

WIFI_IF_GM = "mesh0"
WIFI_IF_SL = "mesh0"

PHC_DEV_GM = "/dev/ptp3"
PHC_DEV_SL = "/dev/ptp3"

USE_PHC2SYS = True

# PTP4L_INSTANCE = "ptp4l"
PTP4L_INSTANCE = "/home/apu/wifi-ptp/ptp/ptp4l"
PTP_REMOTE_CONFIG_PATH = "/opt/ptp_conf/"
GM_FILE = "ptp4l_grandmaster.conf"
SV_FILE = "ptp4l_slave.conf"

GRANDMASTER_LOG = "/tmp/ptp4l_grandmaster.log"
SLAVE_LOG      = "/tmp/ptp4l_slave.log"
PHC2SYS_LOG    = "/tmp/phc2sys_wifi.log"

LOCAL_LOG_BASE = "logs"
PHC2SYS_TUNING = "-m -w --step_threshold=1"

# RBM
RECEIVER_PATH    = "/opt/timestamping/receiver_hwts_logger"
BROADCASTER_PATH = "/home/apu/testbed_files/experiments/reference_broadcast/broadcaster"
RECEIVER_IFACES_GM = ("eth0",)
RECEIVER_IFACES_SL = ("eth0",)
# BROADCAST_REMOTE_HOST = GM_HOST
BROADCAST_IF          = WIFI_IF_GM
BROADCAST_RATE_PPS    = 1000

# ==============================
# --- Helpers
# ==============================

def sh(host: str, cmd: str, check: bool = True) -> subprocess.CompletedProcess:
    return subprocess.run(["ssh", host, cmd], check=check, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)


def write_ptp_config(host: str, role: str, conf_json_path: str) -> None:
    import json
    with open(conf_json_path, "r") as f:
        data = json.load(f)
    lines = []
    for section, options in data.items():
        lines.append(f"[{section}]")
        for k, v in options.items():
            lines.append(f"{k} {v}")
    conf_str = "\n".join(lines)
    fname = GM_FILE if role == "grandmaster" else SV_FILE
    sh(host, f"echo '{conf_str}' | sudo tee '{PTP_REMOTE_CONFIG_PATH}/{fname}' > /dev/null")
    print(f"[{host}] wrote {fname}")


def delete_remote_logs(host: str) -> None:
    cmd = f"sudo rm -f {GRANDMASTER_LOG} {SLAVE_LOG} {PHC2SYS_LOG}"
    try:
        sh(host, cmd)
        print(f"[{host}] cleared old logs")
    except subprocess.CalledProcessError:
        pass


def stop_services(host: str) -> None:
    for proc in ("ptp4l", "phc2sys", "receiver_hwts_logger", "broadcaster"):
        try:
            sh(host, f"pkill -f {proc} || true")
        except subprocess.CalledProcessError:
            pass
    try:
        sh(host, "sync")
    except subprocess.CalledProcessError:
        pass
    print(f"[{host}] stopped all services and flushed")


def start_gm(host: str, wifi_if: str) -> None:
    cmd = f"{PTP4L_INSTANCE} -i {wifi_if} -p {PHC_DEV_GM} -m -H -f {PTP_REMOTE_CONFIG_PATH}{GM_FILE} > {GRANDMASTER_LOG} 2>&1 &"
    sh(host, cmd)
    print(f"[{host}] ptp4l grandmaster on {wifi_if}")


def start_slave(host: str, wifi_if: str) -> None:
    cmd = f"{PTP4L_INSTANCE} -i {wifi_if} -p {PHC_DEV_SL} -m -H -f {PTP_REMOTE_CONFIG_PATH}{SV_FILE} > {SLAVE_LOG} 2>&1 &"
    sh(host, cmd)
    print(f"[{host}] ptp4l slave with command: {cmd}")
    print(f"[{host}] ptp4l slave on {wifi_if}")


def start_phc2sys(host: str, phc_dev: str) -> None:
    if not USE_PHC2SYS:
        print(f"[{host}] phc2sys disabled")
        return
    phc = phc_dev if phc_dev.startswith("/dev/") else f"/dev/{phc_dev}"
    cmd = f"phc2sys -s {phc} -c eth0 -O 0 {PHC2SYS_TUNING} > {PHC2SYS_LOG} 2>&1 &"
    sh(host, cmd)
    print(f"[{host}] phc2sys: {phc} -> eth0")


def is_ptp_running(host: str) -> bool:
    try:
        subprocess.run(["ssh", host, "pgrep -x ptp4l > /dev/null"], check=True)
        return True
    except subprocess.CalledProcessError:
        return False


def wait_until_running(hosts: Tuple[str, str], retries: int = 10, delay: float = 1.5) -> bool:
    for _ in range(retries):
        if all(is_ptp_running(h) for h in hosts):
            print("All ptp4l processes are running.")
            return True
        time.sleep(delay)
    return False


def collect_logs(host: str, local_dir: str, gm: bool) -> None:
    os.makedirs(local_dir, exist_ok=True)
    src = GRANDMASTER_LOG if gm else SLAVE_LOG
    label = "ptp4l_grandmaster" if gm else "ptp4l_slave"
    try:
        out = sh(host, f'bash -lc "sync; sleep 1; [ -f \"{src}\" ] && cat \"{src}\""')
        with open(os.path.join(local_dir, f"{label}.csv"), "w") as f:
            f.write(out.stdout.replace("\x00", ""))
    except subprocess.CalledProcessError as e:
        with open(os.path.join(local_dir, f"{label}_ERROR.txt"), "w") as f:
            f.write(e.stderr or "")
    try:
        out = sh(host, f'bash -lc "[ -f \"{PHC2SYS_LOG}\" ] && cat \"{PHC2SYS_LOG}\""')
        if out.stdout:
            with open(os.path.join(local_dir, "phc2sys.csv"), "w") as f:
                f.write(out.stdout)
    except subprocess.CalledProcessError:
        pass
    tslog_dir = os.path.join(local_dir, "tslogs")
    os.makedirs(tslog_dir, exist_ok=True)
    res = subprocess.run(["ssh", host, r'sh -lc "ls -td /tmp/tslog_* 2>/dev/null | head -n2"'], capture_output=True, text=True)
    folders = [line.strip() for line in res.stdout.splitlines() if line.strip()]
    copied = set()
    for fld in folders:
        lsres = subprocess.run(["ssh", host, f'sh -lc \'ls -1 "{fld}"/*.csv 2>/dev/null\''], capture_output=True, text=True)
        for remote_file in lsres.stdout.splitlines():
            base = os.path.basename(remote_file)
            if base in copied:
                continue
            subprocess.run(["scp", "-q", f"{host}:{remote_file}", os.path.join(tslog_dir, base)], check=False)
            copied.add(base)

# RBM helpers
def start_receivers(host: str, ifaces: tuple[str, ...]) -> None:
    for iface in ifaces:
        cmd = f"sudo taskset -c 1 chrt -f 90 {RECEIVER_PATH} {iface} > /dev/null 2>&1 &"
        sh(host, cmd)

def wait_for_receivers(host: str, expected: int, timeout: int = 10) -> None:
    for _ in range(timeout):
        res = subprocess.run(["ssh", host, "pgrep -fa receiver_hwts_logger | wc -l"], capture_output=True, text=True)
        if int(res.stdout.strip() or 0) >= expected:
            return
        time.sleep(1)
    print(f"ERROR: receivers not running on {host}")
    sys.exit(1)

def start_broadcaster_remote(host: str, iface: str, rate_pps: int) -> None:
    cmd = f"sudo taskset -c 1 chrt -f 90 {BROADCASTER_PATH} {iface} {rate_pps} > /dev/null 2>&1 &"
    sh(host, cmd)
    
def start_broadcaster():
    cmd = f"sudo taskset -c 1 chrt -f 90 {BROADCASTER_PATH} eth0 1000 > /dev/null 2>&1 & echo $!"
    proc = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, text=True)
    return int(proc.stdout.strip())

def stop_broadcaster_remote(host: str) -> None:
    sh(host, "pkill -f broadcaster || true")

# ==============================
# --- Main
# ==============================

def main():
    if len(sys.argv) != 2:
        print("Usage: run_ptp_wifi_two_nodes.py <duration_seconds>")
        sys.exit(1)
    duration = int(sys.argv[1])
    for host in (GM_HOST, SL_HOST):
        stop_services(host)
        delete_remote_logs(host)
    start_receivers(GM_HOST, RECEIVER_IFACES_GM)
    start_receivers(SL_HOST, RECEIVER_IFACES_SL)
    wait_for_receivers(GM_HOST, len(RECEIVER_IFACES_GM))
    wait_for_receivers(SL_HOST, len(RECEIVER_IFACES_SL))
    
    # start_broadcaster_remote(BROADCAST_REMOTE_HOST, BROADCAST_IF, BROADCAST_RATE_PPS)
    broadcaster_pid = start_broadcaster()
    
    gm_json = os.path.join("ptp_config", "eth", "ptp4l_grandmaster.json")
    sl_json = os.path.join("ptp_config", "eth", "ptp4l_bc_slave.json")
    write_ptp_config(GM_HOST, "grandmaster", gm_json)
    write_ptp_config(SL_HOST, "slave", sl_json)
    
    start_gm(GM_HOST, WIFI_IF_GM)
    start_slave(SL_HOST, WIFI_IF_SL)
    
    start_phc2sys(GM_HOST, PHC_DEV_GM)
    start_phc2sys(SL_HOST, PHC_DEV_SL)
    
    if not wait_until_running((GM_HOST, SL_HOST)):
        sys.exit(2)
    for remaining in range(duration, 0, -1):
        sys.stdout.write(f"\rTime left: {remaining:4d} s")
        sys.stdout.flush()
        time.sleep(1)
    print("\rTime left:    0 s\n")
    for host in (GM_HOST, SL_HOST):
        stop_services(host)
    
    print("Stopping local reference broadcaster")
    subprocess.run(["sudo", "pkill", "broadcaster"])
    
    # stop_broadcaster_remote(BROADCAST_REMOTE_HOST)
    
    timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    folder = os.path.join(LOCAL_LOG_BASE, f"{timestamp}_wifi_ptp4l_{duration}s_{GM_HOST}-{SL_HOST}")
    collect_logs(GM_HOST, os.path.join(folder, GM_HOST), gm=True)
    collect_logs(SL_HOST, os.path.join(folder, SL_HOST), gm=False)
    print(f"Done. Logs saved under: {folder}")

if __name__ == "__main__":
    main()