#!/usr/bin/env python3
"""
Docker Backup/Restore Tool
Backup and restore containers (config + volumes) across multiple Docker hosts.
"""

import argparse
import json
import os
import sys
import tarfile
import tempfile
import time
from datetime import datetime
from io import BytesIO
from pathlib import Path

import docker
from docker.errors import DockerException, NotFound, APIError
from rich.console import Console
from rich.table import Table
from rich import print as rprint

console = Console()

# ── Host connection helpers ──────────────────────────────────────────────────

def connect(host: str) -> docker.DockerClient:
    """Connect to a Docker host. 'local' uses the default socket."""
    try:
        if host in ("local", "unix:///var/run/docker.sock", ""):
            client = docker.from_env()
        elif host.startswith("ssh://"):
            client = docker.DockerClient(base_url=host, use_ssh_client=True)
        else:
            # tcp:// or host:port shorthand
            if not host.startswith("tcp://") and "://" not in host:
                host = f"tcp://{host}"
            client = docker.DockerClient(base_url=host)
        client.ping()
        return client
    except DockerException as e:
        console.print(f"[red]Cannot connect to {host}: {e}[/red]")
        sys.exit(1)


# ── Backup ───────────────────────────────────────────────────────────────────

def backup_container(client: docker.DockerClient, container_name: str,
                     output_dir: Path, save_image: bool = False):
    """
    Backup a single container to output_dir/<container_name>_<timestamp>.tar
    Contents:
      config.json   — full docker inspect output
      volumes/      — one .tar per named/anonymous volume
      image.tar     — saved image (optional, can be large)
    """
    try:
        container = client.containers.get(container_name)
    except NotFound:
        console.print(f"[red]Container not found: {container_name}[/red]")
        sys.exit(1)

    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    safe_name = container.name.lstrip("/").replace("/", "_")
    out_file = output_dir / f"{safe_name}_{ts}.tar"

    console.print(f"Backing up [cyan]{container.name}[/cyan] → [green]{out_file}[/green]")

    inspect = client.api.inspect_container(container.id)

    with tarfile.open(out_file, "w") as tar:

        # 1. config.json
        config_bytes = json.dumps(inspect, indent=2).encode()
        info = tarfile.TarInfo(name="config.json")
        info.size = len(config_bytes)
        info.mtime = int(time.time())
        tar.addfile(info, BytesIO(config_bytes))
        console.print("  [dim]✓ config.json[/dim]")

        # 2. Volumes
        mounts = inspect.get("Mounts", [])
        for mount in mounts:
            mtype = mount.get("Type", "")
            if mtype not in ("volume", "bind"):
                continue

            if mtype == "volume":
                vol_name = mount.get("Name", "")
                dest = mount.get("Destination", "")
                label = vol_name or dest.replace("/", "_")
            else:
                src = mount.get("Source", "")
                label = src.replace("/", "_").lstrip("_")
                dest = src

            archive_name = f"volumes/{label}.tar"
            console.print(f"  [dim]Archiving volume: {label} ({dest})[/dim]")

            try:
                # Use docker cp stream to get volume data
                stream, _ = client.api.get_archive(container.id, dest)
                vol_data = b"".join(stream)
                vol_info = tarfile.TarInfo(name=archive_name)
                vol_info.size = len(vol_data)
                vol_info.mtime = int(time.time())
                tar.addfile(vol_info, BytesIO(vol_data))
                console.print(f"  [dim]✓ {archive_name} ({len(vol_data)//1024} KB)[/dim]")
            except (APIError, Exception) as e:
                console.print(f"  [yellow]⚠ Could not archive {label}: {e}[/yellow]")

        # 3. Image (optional)
        if save_image:
            image_tag = inspect["Config"].get("Image", "")
            console.print(f"  [dim]Saving image: {image_tag} (may take a while)…[/dim]")
            try:
                image = client.images.get(image_tag)
                img_data = b"".join(image.save())
                img_info = tarfile.TarInfo(name="image.tar")
                img_info.size = len(img_data)
                img_info.mtime = int(time.time())
                tar.addfile(img_info, BytesIO(img_data))
                console.print(f"  [dim]✓ image.tar ({len(img_data)//1024//1024} MB)[/dim]")
            except Exception as e:
                console.print(f"  [yellow]⚠ Could not save image: {e}[/yellow]")

    console.print(f"[green]✓ Backup complete:[/green] {out_file}")
    return out_file


# ── Restore ──────────────────────────────────────────────────────────────────

def restore_container(client: docker.DockerClient, backup_file: Path,
                      new_name: str | None = None, start: bool = False,
                      load_image: bool = False):
    """
    Restore a container from a backup .tar produced by backup_container().
    """
    console.print(f"Restoring from [cyan]{backup_file}[/cyan]")

    with tarfile.open(backup_file, "r") as tar:
        # 1. Read config
        config_member = tar.getmember("config.json")
        config = json.loads(tar.extractfile(config_member).read())

    orig_name = config["Name"].lstrip("/")
    container_name = new_name or orig_name
    image_name = config["Config"]["Image"]

    console.print(f"  Original name : [cyan]{orig_name}[/cyan]")
    console.print(f"  Restore as    : [cyan]{container_name}[/cyan]")
    console.print(f"  Image         : [cyan]{image_name}[/cyan]")

    # 2. Optionally load image from backup
    if load_image:
        with tarfile.open(backup_file, "r") as tar:
            try:
                img_member = tar.getmember("image.tar")
                console.print("  [dim]Loading image from backup…[/dim]")
                img_data = tar.extractfile(img_member).read()
                client.images.load(img_data)
                console.print("  [dim]✓ Image loaded[/dim]")
            except KeyError:
                console.print("  [yellow]⚠ No image.tar in backup, skipping load[/yellow]")

    # 3. Pull image if not present
    try:
        client.images.get(image_name)
        console.print(f"  [dim]✓ Image {image_name} already present[/dim]")
    except NotFound:
        console.print(f"  [dim]Pulling {image_name}…[/dim]")
        try:
            client.images.pull(image_name)
        except Exception as e:
            console.print(f"  [red]Cannot pull image: {e}[/red]")
            sys.exit(1)

    # 4. Check for name collision
    try:
        existing = client.containers.get(container_name)
        console.print(f"  [yellow]Container {container_name} already exists (id={existing.short_id})[/yellow]")
        answer = input("  Remove existing container? [y/N] ").strip().lower()
        if answer == "y":
            existing.remove(force=True)
        else:
            console.print("[red]Restore aborted.[/red]")
            sys.exit(1)
    except NotFound:
        pass

    # 5. Recreate volumes and restore data
    with tarfile.open(backup_file, "r") as tar:
        volume_members = [m for m in tar.getmembers()
                          if m.name.startswith("volumes/") and m.name.endswith(".tar")]

    mounts = config.get("Mounts", [])
    vol_map: dict[str, str] = {}  # label -> actual volume name or bind path

    for mount in mounts:
        mtype = mount.get("Type", "")
        dest = mount.get("Destination", "")
        if mtype == "volume":
            vol_name = mount.get("Name", "")
            label = vol_name or dest.replace("/", "_")
        else:
            src = mount.get("Source", "")
            label = src.replace("/", "_").lstrip("_")
            vol_map[label] = src
            continue  # bind mounts: path must exist on target host

        # Create volume if missing
        try:
            client.volumes.get(vol_name)
            console.print(f"  [dim]Volume {vol_name} already exists[/dim]")
        except NotFound:
            client.volumes.create(name=vol_name)
            console.print(f"  [dim]✓ Created volume {vol_name}[/dim]")
        vol_map[label] = vol_name

    # Restore volume data via a temporary Alpine container
    with tarfile.open(backup_file, "r") as tar:
        for vmember in volume_members:
            label = vmember.name[len("volumes/"):-len(".tar")]
            if label not in vol_map:
                console.print(f"  [yellow]⚠ No matching mount for volume backup: {label}[/yellow]")
                continue

            vol_name = vol_map[label]
            dest_path = next(
                (m["Destination"] for m in mounts
                 if (m.get("Name") == vol_name or
                     m.get("Source", "").replace("/", "_").lstrip("_") == label)),
                None
            )
            if not dest_path:
                continue

            console.print(f"  [dim]Restoring volume data: {vol_name} → {dest_path}[/dim]")
            vol_data = tar.extractfile(vmember).read()

            # Spin up temp container with the volume mounted, restore via tar stream
            if vol_map[label] == vol_name:
                volumes_arg = {vol_name: {"bind": "/restore_target", "mode": "rw"}}
            else:
                volumes_arg = {vol_name: {"bind": "/restore_target", "mode": "rw"}}

            try:
                restore_ctr = client.containers.run(
                    "alpine",
                    command="sh -c 'cd /restore_target && tar xf /tmp/vol.tar --strip-components=1'",
                    volumes=volumes_arg,
                    detach=True,
                    remove=False,
                )
                # Upload the tar to the temp container
                client.api.put_archive(restore_ctr.id, "/tmp", vol_data)
                restore_ctr.wait()
                restore_ctr.remove()
                console.print(f"  [dim]✓ Volume data restored: {vol_name}[/dim]")
            except Exception as e:
                console.print(f"  [yellow]⚠ Could not restore volume data for {vol_name}: {e}[/yellow]")

    # 6. Rebuild container create kwargs from inspect
    cfg = config["Config"]
    host_cfg = config["HostConfig"]

    create_kwargs: dict = {
        "image": image_name,
        "name": container_name,
        "command": cfg.get("Cmd"),
        "entrypoint": cfg.get("Entrypoint"),
        "environment": cfg.get("Env") or [],
        "working_dir": cfg.get("WorkingDir") or None,
        "user": cfg.get("User") or None,
        "hostname": cfg.get("Hostname") or None,
        "domainname": cfg.get("Domainname") or None,
        "labels": cfg.get("Labels") or {},
        "tty": cfg.get("Tty", False),
        "stdin_open": cfg.get("OpenStdin", False),
    }

    # Ports
    exposed = cfg.get("ExposedPorts") or {}
    port_bindings = host_cfg.get("PortBindings") or {}
    if exposed or port_bindings:
        ports = {}
        for port_proto in exposed:
            bindings = port_bindings.get(port_proto)
            if bindings:
                host_ports = [b.get("HostPort") for b in bindings if b.get("HostPort")]
                ports[port_proto] = host_ports or None
            else:
                ports[port_proto] = None
        create_kwargs["ports"] = ports

    # Restart policy
    restart = host_cfg.get("RestartPolicy", {})
    if restart.get("Name"):
        create_kwargs["restart_policy"] = {
            "Name": restart["Name"],
            "MaximumRetryCount": restart.get("MaximumRetryCount", 0),
        }

    # Network mode
    net_mode = host_cfg.get("NetworkMode", "default")
    if net_mode not in ("default", "bridge"):
        create_kwargs["network_mode"] = net_mode

    # Volumes mounts
    volume_binds = []
    for mount in mounts:
        mtype = mount.get("Type", "")
        dest = mount.get("Destination", "")
        mode = "rw" if not mount.get("RW") is False else "ro"
        if mtype == "volume":
            volume_binds.append(f"{mount['Name']}:{dest}:{mode}")
        elif mtype == "bind":
            volume_binds.append(f"{mount['Source']}:{dest}:{mode}")
    if volume_binds:
        create_kwargs["volumes"] = volume_binds

    # Extra host config fields
    if host_cfg.get("Privileged"):
        create_kwargs["privileged"] = True
    if host_cfg.get("NetworkMode") == "host":
        create_kwargs["network_mode"] = "host"
    cap_add = host_cfg.get("CapAdd") or []
    if cap_add:
        create_kwargs["cap_add"] = cap_add
    devices = host_cfg.get("Devices") or []
    if devices:
        create_kwargs["devices"] = [
            f"{d['PathOnHost']}:{d['PathInContainer']}:{d['CgroupPermissions']}"
            for d in devices
        ]

    # 7. Create container
    try:
        ctr = client.containers.create(**{k: v for k, v in create_kwargs.items() if v is not None})
        console.print(f"[green]✓ Container created:[/green] {ctr.name} ({ctr.short_id})")
    except Exception as e:
        console.print(f"[red]Failed to create container: {e}[/red]")
        sys.exit(1)

    if start:
        ctr.start()
        console.print(f"[green]✓ Container started[/green]")

    return ctr


# ── List ─────────────────────────────────────────────────────────────────────

def list_containers(hosts: list[str], all_containers: bool = False):
    table = Table(title="Docker Containers")
    table.add_column("Host", style="cyan")
    table.add_column("Name", style="white")
    table.add_column("Image", style="dim")
    table.add_column("Status", style="green")
    table.add_column("Ports", style="dim")

    for host in hosts:
        client = connect(host)
        label = host
        try:
            containers = client.containers.list(all=all_containers)
            for c in containers:
                ports = ", ".join(
                    f"{v[0]['HostPort']}→{k}" if v else k
                    for k, v in (c.ports or {}).items()
                ) or "-"
                status = c.status
                color = "green" if status == "running" else "yellow"
                table.add_row(label, c.name, c.image.tags[0] if c.image.tags else c.image.short_id,
                              f"[{color}]{status}[/{color}]", ports)
                label = ""  # only show host label on first row
        except Exception as e:
            table.add_row(host, f"[red]ERROR: {e}[/red]", "", "", "")

    console.print(table)


# ── CLI ──────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(
        description="Backup and restore Docker containers across multiple hosts.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # List containers on multiple hosts
  docker_backup.py list --host local --host tcp://192.168.1.10:2375

  # Backup a container on a remote host
  docker_backup.py backup myapp --host tcp://192.168.1.10:2375 --output /backups

  # Backup with image saved
  docker_backup.py backup myapp --host local --save-image

  # Restore to a different host (pull image from registry)
  docker_backup.py restore /backups/myapp_20240101_120000.tar --host tcp://192.168.1.20:2375

  # Restore with a new name and auto-start
  docker_backup.py restore /backups/myapp_20240101_120000.tar --name myapp2 --start

  # Restore using image from backup file (no registry needed)
  docker_backup.py restore /backups/myapp_20240101_120000.tar --load-image --start
        """
    )

    sub = parser.add_subparsers(dest="command", required=True)

    # list
    p_list = sub.add_parser("list", help="List containers across hosts")
    p_list.add_argument("--host", action="append", default=[], dest="hosts",
                        metavar="HOST", help="Docker host (repeatable). Use 'local' for local socket.")
    p_list.add_argument("-a", "--all", action="store_true", help="Show stopped containers too")

    # backup
    p_back = sub.add_parser("backup", help="Backup a container")
    p_back.add_argument("container", help="Container name or ID")
    p_back.add_argument("--host", default="local", metavar="HOST",
                        help="Docker host (default: local)")
    p_back.add_argument("--output", "-o", default=".", metavar="DIR",
                        help="Output directory (default: current dir)")
    p_back.add_argument("--save-image", action="store_true",
                        help="Also save the container image into the backup (large!)")

    # restore
    p_rest = sub.add_parser("restore", help="Restore a container from backup")
    p_rest.add_argument("backup", help="Path to the backup .tar file")
    p_rest.add_argument("--host", default="local", metavar="HOST",
                        help="Docker host to restore to (default: local)")
    p_rest.add_argument("--name", metavar="NAME",
                        help="Override container name (default: original name)")
    p_rest.add_argument("--start", action="store_true",
                        help="Start the container after restoring")
    p_rest.add_argument("--load-image", action="store_true",
                        help="Load image from backup instead of pulling from registry")

    args = parser.parse_args()

    if args.command == "list":
        hosts = args.hosts or ["local"]
        list_containers(hosts, all_containers=args.all)

    elif args.command == "backup":
        client = connect(args.host)
        output_dir = Path(args.output)
        output_dir.mkdir(parents=True, exist_ok=True)
        backup_container(client, args.container, output_dir, save_image=args.save_image)

    elif args.command == "restore":
        client = connect(args.host)
        restore_container(client, Path(args.backup),
                          new_name=args.name,
                          start=args.start,
                          load_image=args.load_image)


if __name__ == "__main__":
    main()