#!/usr/bin/env python3
"""
Weave Template (Claude) — Experiment lifecycle manager.

Default experiment: deploy a single Ubuntu 22.04 VM on a verified-available
random site, record CPU/memory/network utilization for DURATION seconds,
then download the results to the local data/ directory.

Commands:
    start    SLICE_NAME  — Deploy VM and run utilization capture
    stop     SLICE_NAME  — Delete slice
    monitor  SLICE_NAME  — Health check (exit 1 on failure)
    collect  SLICE_NAME  — Download data from VM to data/
    analyze  SLICE_NAME  — Parse results and write summary.json
"""
import os
import sys
import json
import time
from datetime import datetime, timezone

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
DATA_DIR = os.path.join(SCRIPT_DIR, os.environ.get("DATA_DIR", "data"))
DURATION = int(os.environ.get("DURATION", "60"))

# Metadata file tracks run info for notebooks
METADATA_FILE = os.path.join(DATA_DIR, "metadata.json")


def _save_metadata(metadata):
    """Save run metadata to data/metadata.json."""
    os.makedirs(DATA_DIR, exist_ok=True)
    with open(METADATA_FILE, "w") as f:
        json.dump(metadata, f, indent=2)


def _load_metadata():
    """Load run metadata if it exists."""
    if os.path.exists(METADATA_FILE):
        with open(METADATA_FILE) as f:
            return json.load(f)
    return {}


def start(slice_name):
    """Build, submit, and run utilization capture on a single-VM slice."""
    from fabrictestbed_extensions.fablib.fablib import FablibManager
    fablib = FablibManager()

    start_time = datetime.now(timezone.utc).isoformat()

    # Step 1: Build and submit the slice using FABlib
    print(f"### PROGRESS: Running — creating slice '{slice_name}'...")
    my_slice = fablib.new_slice(name=slice_name)

    # Pick a random available site
    site = fablib.get_random_site()
    print(f"### PROGRESS: Running — selected site: {site}")

    node = my_slice.add_node(name="node1", site=site,
                             cores=2, ram=8, disk=10,
                             image="default_ubuntu_22")

    print(f"### PROGRESS: Running — submitting slice...")
    my_slice.submit()

    print("### PROGRESS: Running — waiting for slice to be ready...")
    my_slice.wait_ssh(progress=True)

    node = my_slice.get_node("node1")
    mgmt_ip = node.get_management_ip()
    print(f"### PROGRESS: Running — VM ready at {mgmt_ip}")

    # Step 2: Install monitoring tools (previously in boot_config)
    print("### PROGRESS: Running — installing monitoring tools...")
    node.execute("sudo apt-get update -qq && sudo apt-get install -y -qq sysstat jq", quiet=True)

    # Step 3: Run utilization capture
    node.execute("mkdir -p /tmp/experiment_data", quiet=True)
    print(f"### PROGRESS: Running — recording utilization for {DURATION}s...")

    # CPU and memory via vmstat (1-second intervals)
    capture_script = f"""#!/bin/bash
set -e
cd /tmp/experiment_data

# Capture vmstat: procs, memory, swap, io, system, cpu
vmstat -t 1 {DURATION} | awk '
NR==1 {{ next }}
NR==2 {{ for(i=1;i<=NF;i++) h[i]=$i; next }}
{{
  printf "{{\\\"timestamp\\\":\\\"%s %s\\\",\\\"procs_r\\\":%s,\\\"procs_b\\\":%s,\\\"mem_swpd\\\":%s,\\\"mem_free\\\":%s,\\\"mem_buff\\\":%s,\\\"mem_cache\\\":%s,\\\"cpu_us\\\":%s,\\\"cpu_sy\\\":%s,\\\"cpu_id\\\":%s,\\\"cpu_wa\\\":%s}}\\n", $(NF-1), $NF, $1, $2, $3, $4, $5, $6, $(NF-4), $(NF-3), $(NF-2), $(NF-5)
}}' > vmstat.json

# Capture memory snapshots via free
for i in $(seq 1 {DURATION}); do
  free -b | awk -v ts="$(date -Iseconds)" '
    /^Mem:/ {{ printf "{{\\\"timestamp\\\":\\\"%s\\\",\\\"total\\\":%s,\\\"used\\\":%s,\\\"free\\\":%s,\\\"shared\\\":%s,\\\"buff_cache\\\":%s,\\\"available\\\":%s}}\\n", ts, $2, $3, $4, $5, $6, $7 }}'
  sleep 1
done > memory.json

# Capture network via /proc/net/dev
IFACE=$(ip -o link show | awk -F': ' '!/lo/{{print $2; exit}}')
for i in $(seq 1 {DURATION}); do
  RX1=$(cat /sys/class/net/$IFACE/statistics/rx_bytes)
  TX1=$(cat /sys/class/net/$IFACE/statistics/tx_bytes)
  sleep 1
  RX2=$(cat /sys/class/net/$IFACE/statistics/rx_bytes)
  TX2=$(cat /sys/class/net/$IFACE/statistics/tx_bytes)
  echo "{{\\\"timestamp\\\":\\\"$(date -Iseconds)\\\",\\\"interface\\\":\\\"$IFACE\\\",\\\"rx_bytes_per_sec\\\":$((RX2-RX1)),\\\"tx_bytes_per_sec\\\":$((TX2-TX1))}}"
done > network.json

echo "DONE"
"""
    # Upload and run the capture script
    node.execute("cat > /tmp/experiment_data/capture.sh << 'SCRIPT_EOF'\n" + capture_script + "\nSCRIPT_EOF", quiet=True)
    node.execute("chmod +x /tmp/experiment_data/capture.sh", quiet=True)

    # Run the capture (this takes DURATION seconds)
    stdout, stderr = node.execute(f"cd /tmp/experiment_data && bash capture.sh 2>&1", quiet=True)
    if "DONE" not in stdout:
        print(f"WARNING: Capture may not have completed cleanly")
        print(f"  stdout: {stdout[-200:]}")

    print(f"### PROGRESS: Running — utilization capture complete.")

    # Save metadata
    end_time = datetime.now(timezone.utc).isoformat()
    metadata = {
        "slice_name": slice_name,
        "site": site,
        "node_ip": mgmt_ip,
        "image": "default_ubuntu_22",
        "cores": 2,
        "ram_gb": 8,
        "disk_gb": 10,
        "duration_seconds": DURATION,
        "start_time": start_time,
        "end_time": end_time,
        "status": "data_captured"
    }
    _save_metadata(metadata)

    print(f"### PROGRESS: Running — measurement phase complete.")


def stop(slice_name):
    """Delete the experiment slice."""
    from fabrictestbed_extensions.fablib.fablib import FablibManager
    fablib = FablibManager()
    try:
        my_slice = fablib.get_slice(name=slice_name)
        print(f"### PROGRESS: Cleaning up — deleting slice '{slice_name}'...")
        my_slice.delete()
        print(f"### PROGRESS: Slice '{slice_name}' deleted.")
        # Update metadata
        metadata = _load_metadata()
        if metadata:
            metadata["status"] = "deleted"
            metadata["deleted_at"] = datetime.now(timezone.utc).isoformat()
            _save_metadata(metadata)
    except Exception as e:
        print(f"### PROGRESS: Slice not found or already deleted: {e}")


def monitor(slice_name):
    """Health check — exit 1 on failure."""
    from fabrictestbed_extensions.fablib.fablib import FablibManager
    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}")
        sys.exit(1)
    node = my_slice.get_node("node1")
    try:
        stdout, _ = node.execute("echo ok", quiet=True)
        if "ok" not in stdout:
            raise Exception("unexpected output")
    except Exception as e:
        print(f"ERROR: Node health check failed: {e}")
        sys.exit(1)
    print(f"### PROGRESS: Node healthy (state: {state})")


def collect(slice_name):
    """Download experiment data from the VM to local data/ directory."""
    from fabrictestbed_extensions.fablib.fablib import FablibManager
    fablib = FablibManager()

    os.makedirs(DATA_DIR, exist_ok=True)
    print(f"### PROGRESS: Running — collecting data from '{slice_name}'...")

    my_slice = fablib.get_slice(name=slice_name)
    node = my_slice.get_node("node1")

    files_to_collect = ["vmstat.json", "memory.json", "network.json"]
    collected = []

    for filename in files_to_collect:
        remote_path = f"/tmp/experiment_data/{filename}"
        local_path = os.path.join(DATA_DIR, filename)
        try:
            stdout, _ = node.execute(f"cat {remote_path}", quiet=True)
            if stdout.strip():
                # Convert newline-delimited JSON to a proper JSON array
                lines = [line.strip() for line in stdout.strip().split('\n') if line.strip()]
                records = []
                for line in lines:
                    try:
                        records.append(json.loads(line))
                    except json.JSONDecodeError:
                        continue
                with open(local_path, 'w') as f:
                    json.dump(records, f, indent=2)
                collected.append(filename)
                print(f"  Collected {filename}: {len(records)} records")
            else:
                print(f"  WARNING: {filename} is empty")
        except Exception as e:
            print(f"  WARNING: Could not collect {filename}: {e}")

    # Update metadata
    metadata = _load_metadata()
    metadata["status"] = "data_collected"
    metadata["collected_files"] = collected
    metadata["collected_at"] = datetime.now(timezone.utc).isoformat()
    _save_metadata(metadata)

    print(f"### PROGRESS: Running — collected {len(collected)}/{len(files_to_collect)} files to {DATA_DIR}/")


def analyze(slice_name):
    """Parse collected data and generate summary statistics."""
    print(f"### PROGRESS: Running — analyzing data for '{slice_name}'...")

    summary = {"slice_name": slice_name, "metrics": {}}

    # Analyze CPU data (from vmstat.json)
    vmstat_path = os.path.join(DATA_DIR, "vmstat.json")
    if os.path.exists(vmstat_path):
        with open(vmstat_path) as f:
            records = json.load(f)
        if records:
            cpu_user = [r.get("cpu_us", 0) for r in records]
            cpu_sys = [r.get("cpu_sy", 0) for r in records]
            cpu_idle = [r.get("cpu_id", 0) for r in records]
            summary["metrics"]["cpu"] = {
                "user_avg": round(sum(cpu_user) / len(cpu_user), 2),
                "system_avg": round(sum(cpu_sys) / len(cpu_sys), 2),
                "idle_avg": round(sum(cpu_idle) / len(cpu_idle), 2),
                "user_max": max(cpu_user),
                "system_max": max(cpu_sys),
                "samples": len(records)
            }
            print(f"  CPU: avg user={summary['metrics']['cpu']['user_avg']}%, "
                  f"sys={summary['metrics']['cpu']['system_avg']}%, "
                  f"idle={summary['metrics']['cpu']['idle_avg']}%")

    # Analyze memory data
    memory_path = os.path.join(DATA_DIR, "memory.json")
    if os.path.exists(memory_path):
        with open(memory_path) as f:
            records = json.load(f)
        if records:
            used = [r.get("used", 0) for r in records]
            available = [r.get("available", 0) for r in records]
            total = records[0].get("total", 1)
            summary["metrics"]["memory"] = {
                "total_bytes": total,
                "used_avg_bytes": round(sum(used) / len(used)),
                "used_avg_pct": round(100 * sum(used) / len(used) / total, 2) if total else 0,
                "available_avg_bytes": round(sum(available) / len(available)),
                "samples": len(records)
            }
            print(f"  Memory: avg used={summary['metrics']['memory']['used_avg_pct']}% "
                  f"of {round(total / 1e9, 1)}GB")

    # Analyze network data
    network_path = os.path.join(DATA_DIR, "network.json")
    if os.path.exists(network_path):
        with open(network_path) as f:
            records = json.load(f)
        if records:
            rx = [r.get("rx_bytes_per_sec", 0) for r in records]
            tx = [r.get("tx_bytes_per_sec", 0) for r in records]
            summary["metrics"]["network"] = {
                "interface": records[0].get("interface", "unknown"),
                "rx_avg_bytes_per_sec": round(sum(rx) / len(rx)),
                "tx_avg_bytes_per_sec": round(sum(tx) / len(tx)),
                "rx_max_bytes_per_sec": max(rx),
                "tx_max_bytes_per_sec": max(tx),
                "samples": len(records)
            }
            print(f"  Network: avg rx={round(sum(rx)/len(rx)/1024, 1)}KB/s, "
                  f"tx={round(sum(tx)/len(tx)/1024, 1)}KB/s")

    # Load and merge metadata
    metadata = _load_metadata()
    summary["metadata"] = metadata
    summary["analyzed_at"] = datetime.now(timezone.utc).isoformat()

    # Write summary
    summary_path = os.path.join(DATA_DIR, "summary.json")
    with open(summary_path, "w") as f:
        json.dump(summary, f, indent=2)

    print(f"### PROGRESS: Running — summary written to {summary_path}")


if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: experiment.py {start|stop|monitor|collect|analyze} SLICE_NAME")
        sys.exit(1)
    action, name = sys.argv[1], sys.argv[2]
    actions = {"start": start, "stop": stop, "monitor": monitor,
               "collect": collect, "analyze": analyze}
    if action not in actions:
        print(f"Unknown action: {action}")
        sys.exit(1)
    actions[action](name)
