#!/usr/bin/env python3

"""
mass_update.py - Mass update uTun servers via SSH
Usage: python3 mass_update.py [fast|full] [-c config] [-j jobs]

Config format (JSON):
{
    "servers": [
        {"host": "10.0.0.1", "workdir": "/opt/utun"},
        {"host": "10.0.0.2", "port": 2222, "user": "admin", "workdir": "/root/utun"}
    ]
}

Default config: ~/.utun_servers.json
"""

import argparse
import json
import subprocess
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

DEFAULT_CONFIG = Path.home() / ".utun_servers.json"
SSH_TIMEOUT = 300  # 5 minutes


def load_servers(config_path):
    """Load server list from JSON config."""
    with open(config_path, 'r') as f:
        data = json.load(f)
    return data.get("servers", [])


def update_server(server, command):
    """Run update on a single server via SSH."""
    host = server["host"]
    port = server.get("port", 22)
    workdir = server.get("workdir", "~/utun")
    user = server.get("user", "root")

    ssh_cmd = [
        "ssh",
        "-p", str(port),
        "-o", "StrictHostKeyChecking=no",
        "-o", "ConnectTimeout=10",
        "-o", "BatchMode=yes",
        f"{user}@{host}",
        f"cd {workdir} && ./tools/utun_update.sh {command}"
    ]

    try:
        result = subprocess.run(
            ssh_cmd,
            capture_output=True,
            text=True,
            timeout=SSH_TIMEOUT
        )
        return {
            "host": host,
            "returncode": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
    except subprocess.TimeoutExpired:
        return {
            "host": host,
            "returncode": -1,
            "stdout": "",
            "stderr": f"Timeout after {SSH_TIMEOUT}s"
        }
    except Exception as e:
        return {
            "host": host,
            "returncode": -1,
            "stdout": "",
            "stderr": str(e)
        }


def main():
    parser = argparse.ArgumentParser(
        description="Mass update uTun servers via SSH"
    )
    parser.add_argument(
        "command",
        choices=["fast", "full"],
        help="Update mode: fast (make) or full (autotools rebuild)"
    )
    parser.add_argument(
        "-c", "--config",
        default=str(DEFAULT_CONFIG),
        help=f"Servers config file (default: {DEFAULT_CONFIG})"
    )
    parser.add_argument(
        "-j", "--jobs",
        type=int,
        default=10,
        help="Parallel SSH jobs (default: 10)"
    )

    args = parser.parse_args()

    config_path = Path(args.config)
    if not config_path.exists():
        print(f"ERROR: Config not found: {config_path}", file=sys.stderr)
        print("\nCreate it with format:", file=sys.stderr)
        example = {
            "servers": [
                {"host": "10.0.0.1", "workdir": "/opt/utun"},
                {"host": "10.0.0.2", "port": 2222, "user": "admin"}
            ]
        }
        print(json.dumps(example, indent=2), file=sys.stderr)
        sys.exit(1)

    servers = load_servers(config_path)
    if not servers:
        print("ERROR: No servers found in config", file=sys.stderr)
        sys.exit(1)

    print(f"Updating {len(servers)} servers in '{args.command}' mode...")
    print(f"Parallel jobs: {args.jobs}")
    print(f"Config: {config_path}")
    print()

    success = 0
    failed = 0

    with ThreadPoolExecutor(max_workers=args.jobs) as executor:
        futures = {
            executor.submit(update_server, s, args.command): s
            for s in servers
        }

        for future in as_completed(futures):
            result = future.result()
            host = result["host"]
            rc = result["returncode"]

            print(f"{'='*60}")
            print(f"Host: {host}  |  Exit code: {rc}")
            print(f"{'='*60}")

            if result["stdout"]:
                print(result["stdout"])
            if result["stderr"]:
                print(result["stderr"], file=sys.stderr, end="")

            if rc == 0:
                success += 1
            else:
                failed += 1

    print(f"{'='*60}")
    print(f"Done: {success} succeeded, {failed} failed")

    sys.exit(0 if failed == 0 else 1)


if __name__ == "__main__":
    main()