#!/usr/bin/env python3
"""
SeQUeNCe Experiment Lab weave lifecycle script.

This weave provisions one FABRIC VM, installs the SeQUeNCe simulator, and runs
one selected experiment against one selected topology.
"""

import os
import sys
from pathlib import Path

from fabrictestbed_extensions.fablib.fablib import FablibManager


SCRIPT_DIR = Path(__file__).resolve().parent
TOOLS_DIR = SCRIPT_DIR / "tools"

DEFAULTS = {
    "EXPERIMENT": "simple_request",
    "TOPOLOGY": "two_node",
    "STOP_TIME_PS": 3_000_000_000_000,
    "REQUEST_START_PS": 1_000_000_000_000,
    "REQUEST_END_PS": 2_500_000_000_000,
    "REQUEST_MEMORY": 1,
    "REQUEST_FIDELITY": 0.8,
    "RANDOM_REQUEST_COUNT": 3,
}

SUPPORTED = {
    "simple_request": {"two_node", "star_network"},
    "random_requests": {"two_node", "star_network"},
    "teleport_demo": {"teleport_2node"},
}


def _env_int(name: str, default: int) -> int:
    value = os.environ.get(name)
    return int(value) if value not in (None, "") else default


def _env_float(name: str, default: float) -> float:
    value = os.environ.get(name)
    return float(value) if value not in (None, "") else default


def _read_settings() -> dict:
    settings = {
        "experiment": os.environ.get("EXPERIMENT", DEFAULTS["EXPERIMENT"]).strip(),
        "topology": os.environ.get("TOPOLOGY", DEFAULTS["TOPOLOGY"]).strip(),
        "stop_time_ps": _env_int("STOP_TIME_PS", DEFAULTS["STOP_TIME_PS"]),
        "request_start_ps": _env_int("REQUEST_START_PS", DEFAULTS["REQUEST_START_PS"]),
        "request_end_ps": _env_int("REQUEST_END_PS", DEFAULTS["REQUEST_END_PS"]),
        "request_memory": _env_int("REQUEST_MEMORY", DEFAULTS["REQUEST_MEMORY"]),
        "request_fidelity": _env_float("REQUEST_FIDELITY", DEFAULTS["REQUEST_FIDELITY"]),
        "random_request_count": _env_int("RANDOM_REQUEST_COUNT", DEFAULTS["RANDOM_REQUEST_COUNT"]),
    }
    valid_topologies = SUPPORTED.get(settings["experiment"])
    if valid_topologies is None:
        raise ValueError(
            f"Unsupported EXPERIMENT '{settings['experiment']}'. Supported: {', '.join(SUPPORTED)}"
        )
    if settings["topology"] not in valid_topologies:
        allowed = ", ".join(sorted(valid_topologies))
        raise ValueError(
            f"Topology '{settings['topology']}' is not valid for experiment '{settings['experiment']}'. "
            f"Allowed: {allowed}"
        )
    return settings


def _run_or_fail(node, command: str, step_msg: str) -> str:
    print(step_msg)
    stdout, stderr = node.execute(command, quiet=True)
    if stderr.strip():
        print(stderr.strip())
    return stdout


def start(slice_name: str) -> None:
    settings = _read_settings()

    print("### PROGRESS: Step 1/5 - Creating FABRIC slice...")
    print(f"  Experiment: {settings['experiment']}")
    print(f"  Topology:   {settings['topology']}")

    fablib = FablibManager()
    my_slice = fablib.new_slice(name=slice_name)
    node = my_slice.add_node(
        name="simulator",
        cores=4,
        ram=16,
        disk=50,
        image="default_ubuntu_24",
    )

    print("### PROGRESS: Step 2/5 - Submitting slice and waiting for SSH...")
    my_slice.submit()
    my_slice.wait_ssh(progress=True)

    print("### PROGRESS: Step 3/5 - Uploading experiment files to VM...")
    node = my_slice.get_node(name="simulator")
    #node.execute("mkdir -p ~/tools/topologies", quiet=True)
    node.upload_directory(str(TOOLS_DIR), ".")
    node.execute("ls -R tools")

    print("### PROGRESS: Step 4/5 - Installing SeQUeNCe and dependencies...")
    bootstrap_cmd = (
        "sudo apt-get update -qq && "
        "sudo DEBIAN_FRONTEND=noninteractive apt-get install -y -qq python3-venv python3-pip && "
        "python3 -m venv ~/sequence-venv && "
        "~/sequence-venv/bin/python -m pip install -q --upgrade pip && "
        "~/sequence-venv/bin/python -m pip install -q sequence"
    )
    _run_or_fail(node, bootstrap_cmd, "  Installing Python runtime and SeQUeNCe package...")

    print("### PROGRESS: Step 5/5 - Running selected SeQUeNCe experiment...")
    run_cmd = (
        "~/sequence-venv/bin/python ~/tools/run_sequence_experiment.py "
        f"--experiment {settings['experiment']} "
        f"--topology {settings['topology']} "
        f"--stop-time {settings['stop_time_ps']} "
        f"--request-start {settings['request_start_ps']} "
        f"--request-end {settings['request_end_ps']} "
        f"--request-memory {settings['request_memory']} "
        f"--request-fidelity {settings['request_fidelity']} "
        f"--random-request-count {settings['random_request_count']}"
    )
    output = _run_or_fail(node, run_cmd, "  Executing simulation now...")
    print(output.strip())

    print("### PROGRESS: READY! Simulation completed.")
    print("  The VM stays available for inspection.")
    print("  Click Stop to tear down the slice.")


def stop(slice_name: str) -> None:
    fablib = FablibManager()
    try:
        my_slice = fablib.get_slice(name=slice_name)
        print(f"### PROGRESS: Deleting slice '{slice_name}'...")
        my_slice.delete()
        print(f"### PROGRESS: Slice '{slice_name}' deleted.")
    except Exception as exc:
        print(f"### PROGRESS: Slice not found or already deleted: {exc}")


def monitor(slice_name: str) -> None:
    fablib = FablibManager()
    my_slice = fablib.get_slice(name=slice_name)
    state = str(my_slice.get_state())
    if "StableOK" not in state:
        print(f"ERROR: Slice state is {state} (expected StableOK)")
        sys.exit(1)

    for node in my_slice.get_nodes():
        stdout, stderr = node.execute("echo ok", quiet=True)
        if stderr.strip() or "ok" not in stdout:
            print(f"ERROR: Node {node.get_name()} health check failed")
            sys.exit(1)

    print(f"### PROGRESS: Slice healthy (state: {state})")


if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: sequence_experiment_lab.py {start|stop|monitor} SLICE_NAME")
        sys.exit(1)

    action = sys.argv[1]
    slice_name_arg = sys.argv[2]

    if action == "start":
        start(slice_name_arg)
    elif action == "stop":
        stop(slice_name_arg)
    elif action == "monitor":
        monitor(slice_name_arg)
    else:
        print(f"Unknown action: {action}")
        sys.exit(1)
