#!/usr/bin/env python3
import docker
import os
import datetime
import argparse
import subprocess
import sys

BACKUP_DIR = "/backup/docker_volumes"

import requests

PORTAINER_URL = "https://port.sectorq.eu/api"
API_KEY = "ptr_/5RkMCT/j3BTaL32vMSDtXFi76yOXRKVFOrUtzMsl5Y="
HEADERS = {"X-API-Key": f"{API_KEY}"}
ENDPOINT_ID = 4
def get_stack_by_name(stack_name):
    url = f"{PORTAINER_URL}/stacks"
    r = requests.get(url, headers=HEADERS)
    r.raise_for_status()
    stacks = r.json()
    for s in stacks:
        if s['Name'] == stack_name:
            return s
    return None

def stop_stack(stack_name):
    stack = get_stack_by_name(stack_name)
    if not stack:
        print(f"[INFO] Stack {stack_name} not found.")
        return

    print(f"Stopping stack {stack_name} via Portainer API...")
    url = f"{PORTAINER_URL}/stacks/{stack['Id']}/stop?endpointId={ENDPOINT_ID}"
    print("URL:", url)
    r = requests.post(url, headers=HEADERS)
    if r.status_code == 200:
        print(f"[OK] Stack {stack_name} stopped.")
    else:
        print(f"[ERROR] Failed to stop stack: {r.status_code} {r.text}")

def start_stack(stack_name):
    stack = get_stack_by_name(stack_name)
    if not stack:
        print(f"[INFO] Stack {stack_name} not found.")
        return

    print(f"Starting stack {stack_name} via Portainer API...")
    url = f"{PORTAINER_URL}/stacks/{stack['Id']}/start?endpointId={ENDPOINT_ID}"
    r = requests.post(url, headers=HEADERS)
    if r.status_code == 200:
        print(f"[OK] Stack {stack_name} started.")
    else:
        print(f"[ERROR] Failed to start stack: {r.status_code} {r.text}")

def to_timestamp():
    return datetime.datetime.now().strftime("%Y%m%d-%H%M%S")

def get_containers_using_volume(client, vol_name):
    containers = []
    for c in client.containers.list():
        for m in c.attrs['Mounts']:
            if m.get('Name') == vol_name:
                containers.append(c)
    return containers

def backup_volumes(client):
    os.makedirs(BACKUP_DIR, exist_ok=True)
    timestamp = to_timestamp()
    volumes = client.volumes.list()
    if not volumes:
        print("No Docker volumes found.")
        return

    for vol in volumes:
        vol_name = vol.name
        backup_file = os.path.join(BACKUP_DIR, f"{vol_name}-{timestamp}.tar.gz")
        containers = get_containers_using_volume(client, vol_name)

        # Track stacks that need to be stopped
        stacks_to_restart = {}

        for c in containers:
            stack = c.labels.get('com.docker.stack.namespace')
            if stack:
                # Stop stack instead of individual container
                if stack not in stacks_to_restart:
                    stacks_to_restart[stack] = True
                    stop_stack(stack)
            else:
                # Stop individual container
                print(f"Stopping container {c.name} for backup...")
                c.stop()

        # Backup using busybox
        try:
            print(f"Backing up volume {vol_name} → {backup_file}")
            client.containers.run(
                image="busybox",
                command=f"tar czf /backup/{vol_name}-{timestamp}.tar.gz -C /volume .",
                remove=True,
                volumes={
                    vol_name: {'bind': '/volume', 'mode': 'ro'},
                    BACKUP_DIR: {'bind': '/backup', 'mode': 'rw'}
                },
            )
            print(f"[OK] Backup completed: {backup_file}")
        except Exception as e:
            print(f"[ERROR] Failed to backup {vol_name}: {e}")

        # Restart individual containers
        for c in containers:
            stack = c.labels.get('com.docker.stack.namespace')
            if not stack:
                print(f"Starting container {c.name} after backup...")
                c.start()

        # Restart stacks
        for stack in stacks_to_restart.keys():
            start_stack(stack)


def restore_volume(client, vol_name, backup_file):
    containers = get_containers_using_volume(client, vol_name)
    stacks_to_restart = {}

    for c in containers:
        stack = c.labels.get('com.docker.stack.namespace')
        if stack:
            if stack not in stacks_to_restart:
                stacks_to_restart[stack] = True
                stop_stack(stack)
        else:
            print(f"Stopping container {c.name} for restore...")
            c.stop()

    try:
        print(f"Restoring volume {vol_name} from {backup_file}")
        client.containers.run(
            image="busybox",
            command=f"tar xzf /backup/{os.path.basename(backup_file)} -C /volume",
            remove=True,
            volumes={
                vol_name: {'bind': '/volume', 'mode': 'rw'},
                os.path.dirname(backup_file): {'bind': '/backup', 'mode': 'rw'}
            },
        )
        print(f"[OK] Restore completed: {vol_name}")
    except Exception as e:
        print(f"[ERROR] Failed to restore {vol_name}: {e}")

    # Restart containers
    for c in containers:
        stack = c.labels.get('com.docker.stack.namespace')
        if not stack:
            print(f"Starting container {c.name} after restore...")
            c.start()

    # Restart stacks
    for stack in stacks_to_restart.keys():
        stack_file = f"/deployments/{stack}.yml"
        if os.path.exists(stack_file):
            start_stack(stack_file)
        else:
            print(f"[WARNING] Stack file {stack_file} not found. Start manually.")

def main():
    parser = argparse.ArgumentParser(description="Backup or restore Docker volumes.")
    parser.add_argument("action", choices=["backup", "restore"], help="Action to perform")
    parser.add_argument("--volume", help="Volume name (required for restore)")
    parser.add_argument("--file", help="Backup file (required for restore)")
    args = parser.parse_args()

    client = docker.from_env()

    if args.action == "backup":
        backup_volumes(client)
    elif args.action == "restore":
        if not args.volume or not args.file:
            print("[ERROR] --volume and --file are required for restore")
            sys.exit(1)
        restore_volume(client, args.volume, args.file)

if __name__ == "__main__":
    main()