#!/usr/bin/env python3
"""
SSI vs WPAI Benchmark Runner v2 — preserved copy of the script used
for runs 1-4 (1k SSI fresh, 1k WPAI fresh, 50k SSI fresh+images, 50k SSI update).

Originally lived at /tmp/benchmark-runner.py but /tmp/ got cleared.
Runs 5-7 (50k WPAI fresh, 50k SSI fresh re-run, 50k WPAI update) were
launched manually via direct `ssh nohup ... wp ssi 2 ...` and `ssh nohup
... wp all-import run 2 ...` commands because SSH timeouts in this runner
killed the 50k WPAI run at the 2-hour mark.

Telemetry collected per run:
- sar -u -r -d 5: CPU, RAM, disk I/O at 5-second intervals
- curl probes every N seconds: homepage, /shop/, /product-category/men/,
  /product-category/women/, /cart/, /wp-admin/edit.php?post_type=product,
  /wp-admin/ — response times + category product counts
- Pre/post product count and disk usage
- SSI stage timings (parsed from CLI output)
- Final export CSV for field-by-field comparison
"""
import subprocess, json, time, os, threading, re, sys
from datetime import datetime

SNAPSHOT_ID = 390805640  # benchmark-ready-empty
HETZNER_TOKEN = "ATgw36XpTsBbQqcGOHvgRgsF38nyBD85ZLi9ZI1vWztzg8qOmIpFTMfJhBRUDNoN"
SERVER_ID = 132884991
WP_PATH = "/var/www/tests.superspeedyplugins.com/htdocs/"
SSH_HOST = "benchmarking"
DATA_DIR = "/var/www/superspeedy/htdocs/wp-content/plugins/super-speedy-imports/.benchmark-data"
SITE_URL = "https://tests.superspeedyplugins.com"
COOKIE_JAR = "/tmp/benchmark-cookies.txt"

SSI_UPLOADS = "/var/www/tests.superspeedyplugins.com/htdocs/wp-content/uploads/super-speedy-imports"
WPAI_FILES = "/var/www/tests.superspeedyplugins.com/htdocs/wp-content/uploads/wpallimport/files"

PROBE_URLS = [
    ("homepage", "/"),
    ("shop", "/shop/"),
    ("category_men", "/product-category/men/"),
    ("category_women", "/product-category/women/"),
    ("cart", "/cart/"),
    ("admin_products", "/wp-admin/edit.php?post_type=product"),
    ("admin_dashboard", "/wp-admin/"),
]

def log(msg):
    ts = datetime.now().strftime("%H:%M:%S")
    print(f"[{ts}] {msg}", flush=True)

def ssh(cmd, timeout=7200):
    full = f'ssh -o ConnectTimeout=10 -o ServerAliveInterval=30 {SSH_HOST} "{cmd}"'
    r = subprocess.run(full, shell=True, capture_output=True, text=True, timeout=timeout)
    return r.stdout + r.stderr

def wp(cmd, timeout=7200):
    return ssh(f"sudo -u tests46427 wp {cmd} --path={WP_PATH}", timeout)

def hetzner_api(method, path, data=None):
    import urllib.request
    url = f"https://api.hetzner.cloud/v1{path}"
    body = json.dumps(data).encode() if data else None
    req = urllib.request.Request(url, data=body, method=method,
        headers={"Authorization": f"Bearer {HETZNER_TOKEN}", "Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=30) as resp:
        return json.loads(resp.read())

def restore_snapshot():
    log("Restoring snapshot...")
    r = hetzner_api("POST", f"/servers/{SERVER_ID}/actions/rebuild", {"image": SNAPSHOT_ID})
    action_id = r["action"]["id"]
    while True:
        time.sleep(15)
        try:
            status = hetzner_api("GET", f"/actions/{action_id}")["action"]["status"]
        except:
            status = "unknown"
        if status == "success":
            break
        if status == "error":
            raise RuntimeError("Snapshot restore failed")
    for _ in range(60):
        try:
            result = subprocess.run(f"ssh -o ConnectTimeout=5 {SSH_HOST} echo OK",
                shell=True, capture_output=True, text=True, timeout=10)
            if "OK" in result.stdout:
                break
        except:
            pass
        time.sleep(5)
    # Upload helper PHP scripts (SCP avoids shell quote-escaping issues)
    subprocess.run(f"scp /tmp/count-products.php /tmp/count-products-all.php {SSH_HOST}:/tmp/",
                   shell=True, capture_output=True, timeout=15)
    log("Server restored and SSH ready")

def login_wp_admin():
    """Get wp-admin auth cookie on the SERVER (avoids Cloudflare blocks)"""
    ssh("curl -s -c /tmp/admin-cookies.txt -d 'log=benchmark&pwd=BenchmarkAdmin2026!&wp-submit=Log+In&redirect_to=%2Fwp-admin%2F&testcookie=1' "
        "-b /tmp/admin-cookies.txt -L -k 'https://tests.superspeedyplugins.com/wp-login.php' -o /dev/null")
    log("wp-admin cookie acquired (server-side)")

def start_sar(run_name):
    sar_file = f"/tmp/sar-{run_name}.log"
    ssh(f"nohup sar -u -r -d 5 > {sar_file} 2>&1 & echo sar_started")
    return sar_file

def stop_sar_and_collect(sar_file, run_name):
    ssh("pkill -f 'sar -u' || true")
    time.sleep(2)
    local_path = os.path.join(DATA_DIR, f"sar-{run_name}.log")
    subprocess.run(f"scp {SSH_HOST}:{sar_file} {local_path}", shell=True, capture_output=True)
    return local_path

def curl_probe(url, via_server=False):
    """Probe a URL. via_server=True runs curl on the test server (for wp-admin behind Cloudflare)"""
    try:
        if via_server:
            out = ssh(f"curl -s -o /dev/null -w '%{{http_code}} %{{time_total}}' -L -k --max-time 30 -b /tmp/admin-cookies.txt '{url}'", timeout=35)
        else:
            r = subprocess.run(
                f"curl -s -o /dev/null -w '%{{http_code}} %{{time_total}}' -L --max-time 30 '{url}'",
                shell=True, capture_output=True, text=True, timeout=35
            )
            out = r.stdout
        parts = out.strip().split()
        for i in range(len(parts)):
            if parts[i].isdigit() and len(parts[i]) == 3:
                status = int(parts[i])
                time_s = float(parts[i+1]) if i+1 < len(parts) else 30.0
                return {"status": status, "time_s": round(time_s, 3)}
        return {"status": 0, "time_s": 30.0}
    except:
        return {"status": 0, "time_s": 30.0}

def probe_urls_once():
    results = {}
    for name, path in PROBE_URLS:
        url = f"{SITE_URL}{path}"
        via_server = "admin" in name
        results[name] = curl_probe(url, via_server)

    for cat_name in ["category_men", "category_women"]:
        cat_path = [p for n, p in PROBE_URLS if n == cat_name][0]
        try:
            r = subprocess.run(
                f"curl -s -L --max-time 15 '{SITE_URL}{cat_path}'",
                shell=True, capture_output=True, text=True, timeout=20
            )
            m = re.search(r'(\d+)\s+results?', r.stdout)
            if not m:
                m = re.search(r'of\s+(\d+)\s+results', r.stdout)
            results[f"{cat_name}_product_count"] = int(m.group(1)) if m else 0
        except:
            results[f"{cat_name}_product_count"] = 0
    return results

class ProbeCollector(threading.Thread):
    def __init__(self, interval=30):
        super().__init__(daemon=True)
        self.interval = interval
        self.samples = []
        self.running = True
        self.start_time = time.time()

    def run(self):
        while self.running:
            t = time.time()
            try:
                sample = probe_urls_once()
                sample["timestamp"] = datetime.utcnow().isoformat()
                sample["elapsed_s"] = round(t - self.start_time, 1)
                self.samples.append(sample)
            except Exception as e:
                log(f"  Probe error: {e}")
            deadline = t + self.interval
            while self.running and time.time() < deadline:
                time.sleep(1)

    def start_collecting(self):
        self.start_time = time.time()
        self.start()

    def stop_collecting(self):
        self.running = False
        self.join(timeout=35)
        return self.samples

def get_disk_usage():
    out = ssh("du -sm /var/www/tests.superspeedyplugins.com/htdocs/wp-content/uploads/ 2>/dev/null")
    for line in out.strip().split('\n'):
        m = re.match(r'(\d+)', line.strip())
        if m:
            return int(m.group(1))
    return 0

def get_product_count():
    out = ssh("sudo -u tests46427 php /tmp/count-products.php 2>/dev/null")
    try:
        return int(out.strip())
    except ValueError:
        return 0

def get_product_count_all():
    """Including trashed"""
    out = ssh("sudo -u tests46427 php /tmp/count-products-all.php 2>/dev/null")
    counts = {}
    for line in out.strip().split('\n'):
        parts = line.strip().split('\t')
        if len(parts) == 2:
            try:
                counts[parts[0]] = int(parts[1])
            except ValueError:
                pass
    return counts

def run_ssi(import_id, workers=8, file_override=None):
    cmd = f"ssi {import_id} --workers={workers}"
    if file_override:
        cmd += f" --file={file_override}"
    return wp(cmd, timeout=43200)  # 12 hours

def run_wpai(import_id):
    return wp(f"all-import run {import_id}", timeout=86400)  # 24 hours

def export_products(run_name):
    remote_path = f"/tmp/export-{run_name}.csv"
    wp(f"ssi export {remote_path} product --exclude-dates")
    local = os.path.join(DATA_DIR, f"export-{run_name}.csv")
    subprocess.run(f"scp {SSH_HOST}:{remote_path} {local}", shell=True, capture_output=True)
    try:
        with open(local) as f:
            count = sum(1 for _ in f) - 1
    except:
        count = 0
    return local, count

def parse_ssi_output(output):
    metrics = {}
    lines = output.split('\n')
    for line in lines:
        line = line.strip()
        m = re.match(r'Total Duration \(s\)\s+([\d.]+)', line)
        if m:
            metrics['total_duration_s'] = float(m.group(1))
        m = re.match(r'Peak Memory \(MB\)\s+([\d.]+)', line)
        if m:
            metrics['peak_memory_mb'] = float(m.group(1))
        m = re.match(r'([\w-]+)\s+([\w\s_-]+?)\s{2,}([\d,]+)\s', line)
        if m:
            stage = m.group(1).strip()
            metric = m.group(2).strip().lower().replace(' ', '_')
            val = m.group(3).replace(',', '')
            try:
                metrics[f"{stage}__{metric}"] = int(val)
            except:
                pass
    for line in lines:
        m = re.search(r"Updating meta_key='([^']+)'.*?(\d+) row\(s\), ([\d.]+)s", line)
        if m:
            metrics[f"postmeta_update_{m.group(1)}_rows"] = int(m.group(2))
            metrics[f"postmeta_update_{m.group(1)}_seconds"] = float(m.group(3))
    return metrics

def parse_wpai_output(output):
    metrics = {}
    m = re.search(r'Created\s+Updated\s+Skipped\s+Deleted.*?\n(\d+)\s+(\d+)\s+(\d+)\s+(\d+)', output)
    if m:
        metrics['created'] = int(m.group(1))
        metrics['updated'] = int(m.group(2))
        metrics['skipped'] = int(m.group(3))
        metrics['deleted'] = int(m.group(4))
    m = re.search(r'time:\s+(\d+)\s+(minute|second|hour)', output)
    if m:
        val = int(m.group(1))
        unit = m.group(2)
        if 'minute' in unit: val *= 60
        elif 'hour' in unit: val *= 3600
        metrics['reported_time_s'] = val
    return metrics

def parse_sar_file(sar_path):
    """Extract CPU and memory time series from sar log"""
    cpu_series = []
    mem_series = []
    try:
        with open(sar_path) as f:
            lines = f.readlines()
        for i, line in enumerate(lines):
            parts = line.split()
            if 'all' in parts and len(parts) >= 8:
                idx = parts.index('all')
                if idx >= 1 and parts[0][0].isdigit():
                    try:
                        cpu_series.append({
                            "time": parts[0],
                            "user_pct": float(parts[idx+1]),
                            "system_pct": float(parts[idx+3]),
                            "iowait_pct": float(parts[idx+4]),
                            "idle_pct": float(parts[idx+6])
                        })
                    except (ValueError, IndexError):
                        pass
            if 'kbmemfree' in line and i + 1 < len(lines):
                next_parts = lines[i+1].split()
                if len(next_parts) >= 5 and next_parts[0][0].isdigit():
                    try:
                        for j, p in enumerate(next_parts):
                            if '.' in p and float(p) < 100 and j > 2:
                                mem_series.append({
                                    "time": next_parts[0],
                                    "mem_used_pct": float(p)
                                })
                                break
                    except (ValueError, IndexError):
                        pass
    except:
        pass
    return {"cpu": cpu_series, "memory": mem_series}

def run_benchmark(run_name, plugin, import_id, restore=True, file_override=None,
                  probe_interval=30, wpai_file_overwrite=None):
    log(f"═══ RUN: {run_name} ({plugin}) ═══")

    result = {
        "run_name": run_name,
        "plugin": plugin,
        "import_id": import_id,
        "file_override": file_override,
        "started_at": datetime.utcnow().isoformat(),
    }

    if restore:
        restore_snapshot()

    login_wp_admin()

    if wpai_file_overwrite:
        ssh(f"cp {wpai_file_overwrite} {WPAI_FILES}/shirt-test-50k-with-images.csv")
        log(f"  Overwrote WPAI file with {wpai_file_overwrite}")

    result["disk_before_mb"] = get_disk_usage()
    result["products_before"] = get_product_count()
    log(f"  Pre-run: {result['products_before']} products, {result['disk_before_mb']}MB disk")

    sar_file = start_sar(run_name)
    prober = ProbeCollector(interval=probe_interval)
    prober.start_collecting()

    wall_start = time.time()
    if plugin == "ssi":
        raw_output = run_ssi(import_id, file_override=file_override)
        result["parsed_metrics"] = parse_ssi_output(raw_output)
    else:
        raw_output = run_wpai(import_id)
        result["parsed_metrics"] = parse_wpai_output(raw_output)
    wall_end = time.time()

    result["wall_clock_s"] = round(wall_end - wall_start, 2)
    result["raw_output_tail"] = raw_output[-3000:] if len(raw_output) > 3000 else raw_output

    probe_samples = prober.stop_collecting()
    result["probe_samples"] = probe_samples
    sar_local = stop_sar_and_collect(sar_file, run_name)
    result["sar_file"] = sar_local
    result["sar_parsed"] = parse_sar_file(sar_local)

    result["disk_after_mb"] = get_disk_usage()
    result["products_after"] = get_product_count()
    result["products_by_status"] = get_product_count_all()
    result["disk_delta_mb"] = result["disk_after_mb"] - result["disk_before_mb"]
    result["finished_at"] = datetime.utcnow().isoformat()

    products_created = result["products_after"]
    if result["wall_clock_s"] > 0 and products_created > 0:
        result["items_per_second"] = round(products_created / result["wall_clock_s"], 2)
    else:
        result["items_per_second"] = 0

    log(f"  Done: {result['wall_clock_s']}s, {result['products_after']} products (published), "
        f"disk +{result['disk_delta_mb']}MB, {result['items_per_second']} items/s")
    log(f"  Probes collected: {len(probe_samples)}")

    try:
        export_path, export_count = export_products(run_name)
        result["export_file"] = export_path
        result["export_count"] = export_count
        log(f"  Export: {export_count} products")
    except Exception as e:
        result["export_file"] = None
        result["export_count"] = 0
        log(f"  Export failed: {e}")

    result_file = os.path.join(DATA_DIR, f"{run_name}.json")
    with open(result_file, 'w') as f:
        json.dump(result, f, indent=2, default=str)
    log(f"  Saved: {result_file}")

    return result

def run_all():
    """The original orchestration — runs 1-4 worked, run 5 (WPAI 50k) hit SSH timeout
    at 2h despite import continuing for ~66h on the server. After that, I switched to
    direct SSH `nohup ... wp ssi 2 ...` and manual snapshot management to handle the
    multi-day WPAI run plus the comparison re-runs and update tests."""
    os.makedirs(DATA_DIR, exist_ok=True)
    all_results = []

    # Run 1: SSI 1k fresh
    all_results.append(run_benchmark("ssi-1k-fresh", "ssi", 1, restore=True, probe_interval=5))
    # Run 2: WPAI 1k fresh
    all_results.append(run_benchmark("wpai-1k-fresh", "wpai", 1, restore=True, probe_interval=10))
    # Run 3: SSI 50k fresh+images
    all_results.append(run_benchmark("ssi-50k-fresh-images", "ssi", 2, restore=True, probe_interval=30))
    # Run 4: SSI 50k update (NO restore — runs against Run 3's data)
    update_csv = f"{SSI_UPLOADS}/shirt-test-50k-update.csv"
    all_results.append(run_benchmark("ssi-50k-update", "ssi", 2, restore=False,
                                     file_override=update_csv, probe_interval=30))
    # Run 5: WPAI 50k fresh+images — HIT SSH TIMEOUT, finished manually
    all_results.append(run_benchmark("wpai-50k-fresh-images", "wpai", 2, restore=True, probe_interval=60))
    # Run 6: WPAI 50k update — done manually after SSI re-run + snapshot restore
    all_results.append(run_benchmark("wpai-50k-update", "wpai", 2, restore=False,
                                     wpai_file_overwrite=f"{SSI_UPLOADS}/shirt-test-50k-update.csv",
                                     probe_interval=60))

    combined_file = os.path.join(DATA_DIR, "all-results.json")
    with open(combined_file, 'w') as f:
        json.dump(all_results, f, indent=2, default=str)

    log(f"\n{'='*70}\nALL BENCHMARKS COMPLETE\n{'='*70}")
    log(f"{'Run':<30} {'Time':>10} {'Products':>10} {'Items/s':>10} {'Disk':>8}")
    log("-" * 72)
    for r in all_results:
        log(f"{r['run_name']:<30} {r['wall_clock_s']:>8.1f}s {r['products_after']:>10} "
            f"{r['items_per_second']:>10.1f} {r['disk_delta_mb']:>6}MB")

if __name__ == "__main__":
    run_all()
