#!/usr/bin/env python3
import subprocess
import time
import paramiko
import os
from datetime import datetime


def ssh_execute_command(host, username, command):
    """Execute a command on a remote host via SSH and capture its output."""
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    ssh.connect(hostname=host, username=username)
    stdin, stdout, stderr = ssh.exec_command(command)
    output = stdout.read().decode()
    error = stderr.read().decode()
    ssh.close()
    if error:
        raise RuntimeError(f"Error executing command '{command}': {error}")
    return output


def ssh_execute_command_no_wait(host, username, command):
    """Execute a command on a remote host via SSH without waiting for completion."""
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    ssh.connect(hostname=host, username=username)
    ssh.exec_command(command)
    ssh.close()


def fetch_file_from_dut(host, username, remote_path, local_path):
    """Fetch a file from the DUT to the local machine."""
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    ssh.connect(hostname=host, username=username)
    sftp = ssh.open_sftp()
    sftp.get(remote_path, local_path)
    sftp.close()
    ssh.close()


def run_measurement(sample_number, supply, test_case, output_dir):
    """Run the measure.py script."""
    subprocess.run(
        ["python3", "measure.py", str(sample_number), supply, test_case],
        check=True
    )


def run_iperf3_server(output_dir):
    """Run iperf3 in server mode on the host and save its log."""
    print("Starting iperf3 server on the host...")
    log_file = os.path.join(output_dir, "iperf3_server.log")
    server_process = subprocess.Popen(
        ["iperf3", "-s"],
        stdout=open(log_file, "w"),
        stderr=subprocess.STDOUT
    )
    time.sleep(2)  # Allow the server to start
    return server_process


def stop_iperf3_server(server_process):
    """Stop the iperf3 server."""
    print("Stopping iperf3 server on the host...")
    server_process.terminate()
    server_process.wait()


def start_background_command_and_measure(host, username, command, remote_log_file, local_log_file, sample_number, supply, test_case, output_dir):
    """
    Start a background command on the DUT, perform the measurement while it's running,
    and fetch the log file after the command finishes.
    """
    # Start the command in the background and redirect output to a log file on the DUT
    nohup_command = f"nohup {command} > {remote_log_file} 2>&1 &"
    ssh_execute_command_no_wait(host, username, nohup_command)
    print(f"Command '{command}' started on DUT. Log will be saved to {remote_log_file}")

    # Wait briefly to ensure the command is running
    time.sleep(5)

    # Perform the measurement while the command is running
    run_measurement(sample_number, supply, test_case, output_dir)

    # Wait for the command to finish (optional: adjust this based on the command duration)
    time.sleep(5)

    # Fetch the log file from the DUT
    fetch_file_from_dut(host, username, remote_log_file, local_log_file)
    print(f"Log file fetched from DUT: {local_log_file}")


def main():
    sample_number = 8
    dut_ip = "10.11.0.45"  # Replace with the actual DUT IP
    host_ip = "10.11.0.91"  # Replace with the actual host IP
    username = "caster"

    # Ethernet interface IPs for EthPrim and EthSec
    eth_prim_ip = "10.11.0.45"
    eth_sec_ip = "10.11.0.70"

    # Create a timestamped folder for logs and results
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_dir = os.path.join("results", f"test_{timestamp}")
    os.makedirs(output_dir, exist_ok=True)

    print("Measuring baseline...")
    run_measurement(sample_number, "external", "baseline", output_dir)

    print("Running stress test...")
    # Run the stress test in the background and measure while it's running
    remote_stress_log = "/tmp/stress.log"
    local_stress_log = os.path.join(output_dir, "stress.log")
    start_background_command_and_measure(
        dut_ip,
        username,
        "stress -c 2 --io 2 --vm 2 --timeout 30",
        remote_stress_log,
        local_stress_log,
        sample_number,
        "external",
        "CPU",
        output_dir
    )

    print("Running iperf3 test for EthPrim...")
    iperf3_server = run_iperf3_server(output_dir)  # Start iperf3 server on the host
    try:
        remote_iperf3_client_log = "/tmp/iperf3_client_prim.log"
        local_iperf3_client_log = os.path.join(output_dir, "iperf3_client_prim.log")
        start_background_command_and_measure(
            dut_ip,
            username,
            f"iperf3 -c {host_ip} -B {eth_prim_ip}",
            remote_iperf3_client_log,
            local_iperf3_client_log,
            sample_number,
            "external",
            "EthPrim",
            output_dir
        )
    finally:
        stop_iperf3_server(iperf3_server)  # Ensure the server is stopped

    print("Running iperf3 test for EthSec...")
    iperf3_server = run_iperf3_server(output_dir)  # Start iperf3 server on the host
    try:
        remote_iperf3_client_log = "/tmp/iperf3_client_sec.log"
        local_iperf3_client_log = os.path.join(output_dir, "iperf3_client_sec.log")
        start_background_command_and_measure(
            dut_ip,
            username,
            f"iperf3 -c {host_ip} -B {eth_sec_ip}",
            remote_iperf3_client_log,
            local_iperf3_client_log,
            sample_number,
            "external",
            "EthSec",
            output_dir
        )
    finally:
        stop_iperf3_server(iperf3_server)  # Ensure the server is stopped

    input("Switch power to PoE and press Enter to continue...")

    print("Measuring baseline with PoE...")
    run_measurement(sample_number, "PoE", "baseline", output_dir)

    print("Running stress test with PoE...")
    # Run the stress test in the background and measure while it's running
    remote_stress_log_poe = "/tmp/stress_poe.log"
    local_stress_log_poe = os.path.join(output_dir, "stress_poe.log")
    start_background_command_and_measure(
        dut_ip,
        username,
        "stress -c 2 --io 2 --vm 2 --timeout 30",
        remote_stress_log_poe,
        local_stress_log_poe,
        sample_number,
        "PoE",
        "CPU",
        output_dir
    )

    print("Running iperf3 test for EthPrim with PoE...")
    iperf3_server = run_iperf3_server(output_dir)  # Start iperf3 server on the host
    try:
        remote_iperf3_client_log_poe = "/tmp/iperf3_client_prim_poe.log"
        local_iperf3_client_log_poe = os.path.join(output_dir, "iperf3_client_prim_poe.log")
        start_background_command_and_measure(
            dut_ip,
            username,
            f"iperf3 -c {host_ip} -B {eth_prim_ip}",
            remote_iperf3_client_log_poe,
            local_iperf3_client_log_poe,
            sample_number,
            "PoE",
            "EthPrim",
            output_dir
        )
    finally:
        stop_iperf3_server(iperf3_server)  # Ensure the server is stopped

    print("Running iperf3 test for EthSec with PoE...")
    iperf3_server = run_iperf3_server(output_dir)  # Start iperf3 server on the host
    try:
        remote_iperf3_client_log_poe = "/tmp/iperf3_client_sec_poe.log"
        local_iperf3_client_log_poe = os.path.join(output_dir, "iperf3_client_sec_poe.log")
        start_background_command_and_measure(
            dut_ip,
            username,
            f"iperf3 -c {host_ip} -B {eth_sec_ip}",
            remote_iperf3_client_log_poe,
            local_iperf3_client_log_poe,
            sample_number,
            "PoE",
            "EthSec",
            output_dir
        )
    finally:
        stop_iperf3_server(iperf3_server)  # Ensure the server is stopped

    print("All tests completed.")


if __name__ == "__main__":
    main()