#!/usr/bin/env python3
"""
wisp-backup.py — Pull configs and logs from CCR towers and upload to S3.

Usage:
    ./wisp-backup.py                    # Uses config.yaml in same directory
    ./wisp-backup.py --config /path/to/config.yaml
    ./wisp-backup.py --vpn-only         # Just connect VPN, don't backup
    ./wisp-backup.py --dry-run          # Show what would be done, don't upload
"""

import argparse
import datetime
import gzip
import io
import os
import paramiko
import shutil
import subprocess
import sys
import time
import yaml
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed


def load_config(path: str) -> dict:
    with open(path, 'r') as f:
        cfg = yaml.safe_load(f)
    return cfg


def run_vpn_script(action: str, script_dir: str) -> bool:
    """Call home-router-vpn.sh up/down/status."""
    vpn_script = os.path.join(script_dir, "home-router-vpn.sh")
    if not os.path.exists(vpn_script):
        print(f"[!] VPN script not found: {vpn_script}")
        return False

    result = subprocess.run(
        ["bash", vpn_script, action],
        capture_output=True, text=True, timeout=30
    )
    for line in result.stdout.splitlines():
        print(f"  {line}")
    if result.stderr.strip():
        for line in result.stderr.splitlines():
            print(f"  [stderr] {line}")

    return result.returncode == 0


def ssh_connect(host: str, port: int, username: str, key_path: str = None,
                password: str = None, timeout: int = 30):
    """Open an SSH connection to a CCR router."""
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

    if key_path and os.path.exists(key_path):
        # Try Ed25519 first, fall back to RSA
        try:
            key = paramiko.Ed25519Key.from_private_key_file(key_path)
        except paramiko.ssh_exception.SSHException:
            key = paramiko.RSAKey.from_private_key_file(key_path)
        client.connect(host, port=port, username=username, pkey=key,
                       timeout=timeout, allow_agent=False, look_for_keys=False)
    elif password:
        client.connect(host, port=port, username=username, password=password,
                       timeout=timeout, allow_agent=False, look_for_keys=False)
    else:
        # Try default key locations
        client.connect(host, port=port, username=username,
                       timeout=timeout, allow_agent=True)
    return client


def run_ssh_command(client, command: str, timeout: int = 30) -> str:
    """Run a command over SSH and return stdout."""
    _, stdout, stderr = client.exec_command(command, timeout=timeout)
    exit_status = stdout.channel.recv_exit_status()
    output = stdout.read().decode('utf-8', errors='replace')

    if exit_status != 0:
        err = stderr.read().decode('utf-8', errors='replace')
        print(f"    [warn] Command exited {exit_status}: {err.strip()}")
    return output


def fetch_ccr_config(client) -> str:
    """Fetch full CCR config via /export terse."""
    print("    Fetching config ...")
    config = run_ssh_command(client, "/export terse")
    return config


def fetch_ccr_logs(client, lines: int = 500) -> str:
    """Fetch recent CCR log entries."""
    print(f"    Fetching last {lines} log lines ...")
    logs = run_ssh_command(client, f"/log print without-paging count-only={lines}")
    return logs


def compress_text(text: str, filename: str) -> str:
    """Write text to a .gz file and return the path."""
    gz_path = filename + ".gz"
    with gzip.open(gz_path, 'wt', encoding='utf-8') as f:
        f.write(text)
    return gz_path


def upload_to_s3(local_path: str, s3_path: str, bucket: str, region: str,
                 endpoint: str = None):
    """Upload file to S3 (or Wasabi/S3-compatible) using awscli."""
    dest = f"s3://{bucket}/{s3_path}"
    cmd = ["aws", "s3", "cp", local_path, dest, "--region", region]
    if endpoint:
        cmd += ["--endpoint-url", endpoint]
    result = subprocess.run(
        cmd, capture_output=True, text=True, timeout=60
    )
    if result.returncode == 0:
        print(f"    [s3] Uploaded → {dest}")
    else:
        print(f"    [s3] Upload failed: {result.stderr.strip()}")
    return result.returncode == 0


def backup_tower(tower: dict, cfg: dict, date_str: str, temp_dir: str) -> dict:
    """Backup a single tower. Returns dict of results."""
    name = tower['name']
    ip = tower['ip']
    site = tower['site']
    port = tower.get('ssh_port', 22)
    user = tower.get('ssh_user', 'admin')
    key_path = tower.get('ssh_key')
    password = tower.get('ssh_password')
    log_lines = cfg.get('backup', {}).get('log_lines', 500)
    compress = cfg.get('backup', {}).get('compress', True)
    timeout = cfg.get('backup', {}).get('timeout_seconds', 60)

    s3_cfg = cfg.get('s3', {})
    bucket = s3_cfg['bucket']
    region = s3_cfg.get('region', 'us-east-1')
    prefix = s3_cfg.get('prefix', 'wisp-backups')
    endpoint = s3_cfg.get('endpoint')  # Wasabi or S3-compatible endpoint

    results = {
        'tower': name,
        'site': site,
        'ip': ip,
        'config': False,
        'logs': False,
        'error': None
    }

    # Create local dirs
    tower_temp = Path(temp_dir) / site / name / date_str
    tower_temp.mkdir(parents=True, exist_ok=True)

    s3_config_base = f"{prefix}/configs/{site}/{date_str}"
    s3_logs_base = f"{prefix}/logs/{site}/{date_str}"

    print(f"\n  [{name}] ({ip}) — backing up ...")

    try:
        client = ssh_connect(ip, port, user, key_path, password, timeout=timeout)
    except Exception as e:
        results['error'] = f"SSH connection failed: {e}"
        print(f"    [FAIL] {results['error']}")
        return results

    try:
        # --- Config ---
        config_text = fetch_ccr_config(client)
        config_filename = f"{tower_temp}/{name}-{date_str}.rsc"
        with open(config_filename, 'w') as f:
            f.write(config_text)

        if compress:
            config_filename = compress_text(config_text,
                                            str(tower_temp / f"{name}-{date_str}.rsc"))
            remote_config = f"{s3_config_base}/{name}-{date_str}.rsc.gz"
        else:
            remote_config = f"{s3_config_base}/{name}-{date_str}.rsc"

        upload_to_s3(config_filename, remote_config, bucket, region, endpoint)
        results['config'] = True

        # --- Logs ---
        logs_text = fetch_ccr_logs(client, log_lines)
        log_filename = f"{tower_temp}/{name}-{date_str}.log"
        with open(log_filename, 'w') as f:
            f.write(logs_text)

        if compress:
            log_filename = compress_text(logs_text,
                                         str(tower_temp / f"{name}-{date_str}.log"))
            remote_log = f"{s3_logs_base}/{name}-{date_str}.log.gz"
        else:
            remote_log = f"{s3_logs_base}/{name}-{date_str}.log"

        upload_to_s3(log_filename, remote_log, bucket, region, endpoint)
        results['logs'] = True

    except Exception as e:
        results['error'] = str(e)
        print(f"    [FAIL] {e}")
    finally:
        client.close()

    return results


def cleanup_old_local(temp_dir: str, retention_days: int):
    """Remove local backup copies older than retention_days."""
    cutoff = time.time() - (retention_days * 86400)
    removed = 0
    for root, dirs, files in os.walk(temp_dir):
        for f in files:
            path = os.path.join(root, f)
            if os.path.getmtime(path) < cutoff:
                os.remove(path)
                removed += 1
        # Clean empty dirs
        for d in dirs:
            dpath = os.path.join(root, d)
            try:
                os.rmdir(dpath)
            except OSError:
                pass
    if removed:
        print(f"[*] Cleaned up {removed} old local files")


def main():
    parser = argparse.ArgumentParser(description="WISP CCR Backup to S3")
    parser.add_argument('--config', default=None,
                        help='Path to config.yaml (default: ./config.yaml)')
    parser.add_argument('--vpn-only', action='store_true',
                        help='Just connect VPN, do not backup')
    parser.add_argument('--dry-run', action='store_true',
                        help='Show what would be done without executing')
    parser.add_argument('--tower', default=None,
                        help='Backup only a specific tower name')
    parser.add_argument('--parallel', type=int, default=5,
                        help='Max parallel tower backups (default: 5)')
    args = parser.parse_args()

    # Resolve config path
    if args.config:
        config_path = args.config
    else:
        config_path = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                   "config.yaml")

    if not os.path.exists(config_path):
        print(f"[!] Config not found: {config_path}")
        print("    Create it from the template or pass --config")
        sys.exit(1)

    cfg = load_config(config_path)
    script_dir = os.path.dirname(os.path.abspath(__file__))
    temp_dir = cfg.get('backup', {}).get('temp_dir', '/tmp/wisp-backups')
    retention_days = cfg.get('backup', {}).get('retention_days', 7)

    date_str = datetime.datetime.now().strftime('%Y-%m-%d')

    print(f"╔══════════════════════════════════════════╗")
    print(f"║  WISP Backup — {date_str}")
    print(f"╚══════════════════════════════════════════╝")

    # Step 1: Connect VPN (only if router not already reachable via WireGuard)
    vpn_was_connected = False
    print("\n[*] Step 1: Checking connectivity via WireGuard ...")
    router_ip = cfg.get('towers', [{}])[0].get('ip', '10.77.0.2') if cfg.get('towers') else '10.77.0.2'
    ping_check = subprocess.run(
        ["ping", "-c", "1", "-W", "2", router_ip],
        capture_output=True, timeout=5
    )
    if ping_check.returncode == 0:
        print(f"  [+] Router reachable at {router_ip} via WireGuard — skipping VPN")
    else:
        print(f"  [-] Router not reachable via WireGuard, trying L2TP/IPsec VPN ...")
        if not run_vpn_script("up", script_dir):
            print("[-] VPN connection failed. Aborting.")
            sys.exit(1)
        vpn_was_connected = True

    # Safety: confirm internet still works before proceeding
    print("[*] Verifying internet connectivity ...")
    check = subprocess.run(
        ["ping", "-c", "1", "-W", "3", "8.8.8.8"],
        capture_output=True, timeout=10
    )
    if check.returncode != 0:
        print("[-] Internet connectivity issue after VPN/backup prep")
        if vpn_was_connected:
            run_vpn_script("down", script_dir)
        sys.exit(1)
    print("[+] Internet OK — proceeding with backup")

    if args.vpn_only:
        print("\n[*] VPN connected (--vpn-only mode). Exiting.")
        print("    Run 'home-router-vpn.sh down' to disconnect.")
        return

    # Step 2: Backup towers
    print("\n[*] Step 2: Backing up towers ...")
    towers = cfg.get('towers', [])
    if args.tower:
        towers = [t for t in towers if t['name'] == args.tower]
        if not towers:
            print(f"[-] Tower '{args.tower}' not found in config")
            if vpn_was_connected:
                run_vpn_script("down", script_dir)
            sys.exit(1)

    if args.dry_run:
        print("\n  [DRY RUN] Would backup these towers:")
        for t in towers:
            print(f"    - {t['name']} ({t['ip']}) @ site: {t['site']}")
        print(f"\n  [DRY RUN] Would upload to s3://{cfg['s3']['bucket']}/")
        if vpn_was_connected:
            run_vpn_script("down", script_dir)
        return

    results = []
    with ThreadPoolExecutor(max_workers=args.parallel) as executor:
        futures = {
            executor.submit(backup_tower, t, cfg, date_str, temp_dir): t
            for t in towers
        }
        for future in as_completed(futures):
            results.append(future.result())

    # Step 3: Disconnect VPN if we connected it
    print("\n[*] Step 3: Disconnecting VPN ...")
    if vpn_was_connected:
        run_vpn_script("down", script_dir)
    else:
        print("  (WireGuard already active — no VPN teardown needed)")

    # Step 4: Cleanup
    print("\n[*] Step 4: Cleaning up local temp files ...")
    cleanup_old_local(temp_dir, retention_days)

    # Summary
    success = [r for r in results if r['error'] is None]
    failed = [r for r in results if r['error'] is not None]

    print(f"\n╔══════════════════════════════════════════╗")
    print(f"║  Summary: {len(success)} OK, {len(failed)} Failed")
    print(f"╚══════════════════════════════════════════╝")
    for r in success:
        ck = "✓" if r['config'] else "✗"
        lk = "✓" if r['logs'] else "✗"
        print(f"  ✓ {r['tower']} ({r['site']}) — config {ck}  logs {lk}")
    for r in failed:
        print(f"  ✗ {r['tower']} ({r['site']}) — {r['error']}")

    if failed:
        sys.exit(1)


if __name__ == '__main__':
    main()