#!/usr/bin/env python3
"""
Docker Factory Reset Script for Ubuntu
Completely resets Docker to a fresh installation state.
"""
import subprocess
import sys
import os
import time
from typing import List


class ProgressBar:
    """Simple progress indicator for operations."""

    def __init__(self, total: int, description: str = ""):
        self.total = total
        self.current = 0
        self.description = description

    def update(self, increment: int = 1):
        self.current += increment
        self._display()

    def _display(self):
        if self.total == 0:
            percentage = 100
        else:
            percentage = int((self.current / self.total) * 100)

        bar_length = 40
        filled = int((percentage / 100) * bar_length)
        bar = '█' * filled + '░' * (bar_length - filled)

        print(f"\r  [{bar}] {percentage}% ({self.current}/{self.total}) {self.description}", end='', flush=True)

    def complete(self):
        self.current = self.total
        self._display()
        print()  # New line after completion


def run_cmd_realtime(cmd: List[str], check: bool = False, show_output: bool = True) -> int:
    """
    Run a command and stream stdout/stderr in real-time.
    Returns the exit code.
    If check=True, exits the script on non-zero return code.
    """
    if show_output:
        print(f"\n▶ RUN: {' '.join(cmd)}")
        print("─" * 60)

    try:
        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,  # Merge stderr into stdout
            text=True,
            bufsize=1,  # Line buffered
        )

        # Stream output line by line
        if show_output:
            for line in process.stdout:
                print(line, end='')  # Print without adding extra newline
        else:
            # Consume output even if not showing
            for _line in process.stdout:
                pass

        process.wait()
        code = process.returncode

        if show_output:
            if code != 0:
                print(f"✗ Command exited with code {code}")
                if check:
                    sys.exit(code)
            else:
                print(f"✓ Command completed successfully")

        return code

    except FileNotFoundError:
        print(f"✗ ERROR: Command not found: {' '.join(cmd)}")
        sys.exit(1)


def run_cmd_quiet(cmd: List[str]) -> tuple:
    """
    Run a command quietly and return (exit_code, stdout, stderr).
    Used for listing/checking operations.
    """
    try:
        result = subprocess.run(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )
        return result.returncode, result.stdout.strip(), result.stderr.strip()
    except FileNotFoundError:
        print(f"✗ ERROR: Command not found: {' '.join(cmd)}")
        sys.exit(1)


def require_root():
    if os.geteuid() != 0:
        print("✗ This script must be run as root (use sudo).")
        sys.exit(1)


def docker_installed() -> bool:
    code, out, err = run_cmd_quiet(["docker", "--version"])
    if code != 0:
        print("✗ Docker does not appear to be installed or accessible.")
        if err:
            print(err)
        return False
    print(f"✓ Docker found: {out}")
    return True


def list_ids(cmd: List[str]) -> List[str]:
    code, out, _err = run_cmd_quiet(cmd)
    if code != 0:
        return []
    return [line.strip() for line in out.splitlines() if line.strip()]


def stop_and_remove_containers():
    print("\n" + "=" * 60)
    print("STEP 1/8: CONTAINERS")
    print("=" * 60)

    container_ids = list_ids(["docker", "ps", "-aq"])
    if not container_ids:
        print("ℹ No containers found.")
        return

    print(f"ℹ Found {len(container_ids)} container(s)")

    running_ids = list_ids(["docker", "ps", "-q"])
    if running_ids:
        print(f"\nℹ Stopping {len(running_ids)} running container(s)...")
        progress = ProgressBar(len(running_ids), "stopping containers")

        for container_id in running_ids:
            run_cmd_realtime(["docker", "stop", container_id], show_output=False)
            progress.update()
        progress.complete()
    else:
        print("ℹ No running containers to stop.")

    print(f"\nℹ Removing all {len(container_ids)} container(s)...")
    progress = ProgressBar(len(container_ids), "removing containers")

    for container_id in container_ids:
        run_cmd_realtime(["docker", "rm", "-f", container_id], show_output=False)
        progress.update()
    progress.complete()

    print("✓ All containers removed")


def remove_images():
    print("\n" + "=" * 60)
    print("STEP 2/8: IMAGES")
    print("=" * 60)

    image_ids = list_ids(["docker", "images", "-aq"])
    if not image_ids:
        print("ℹ No images found.")
        return

    print(f"ℹ Found {len(image_ids)} image(s)")
    print("\nℹ Removing all images...")
    progress = ProgressBar(len(image_ids), "removing images")

    for image_id in image_ids:
        run_cmd_realtime(["docker", "rmi", "-f", image_id], show_output=False)
        progress.update()
    progress.complete()

    print("✓ All images removed")


def remove_networks():
    print("\n" + "=" * 60)
    print("STEP 3/8: NETWORKS")
    print("=" * 60)

    code, out, err = run_cmd_quiet(["docker", "network", "ls", "--format", "{{.Name}}"])
    if code != 0:
        print(f"✗ Error listing networks: {err}")
        return

    default_networks = {"bridge", "host", "none"}
    networks = [n.strip() for n in out.splitlines() if n.strip()]
    custom_networks = [n for n in networks if n not in default_networks]

    if not custom_networks:
        print("ℹ No custom networks found.")
        return

    print(f"ℹ Found {len(custom_networks)} custom network(s): {', '.join(custom_networks)}")
    print("\nℹ Removing custom networks...")
    progress = ProgressBar(len(custom_networks), "removing networks")

    for network in custom_networks:
        run_cmd_realtime(["docker", "network", "rm", network], show_output=False)
        progress.update()
    progress.complete()

    print("✓ All custom networks removed")


def remove_volumes():
    print("\n" + "=" * 60)
    print("STEP 4/8: VOLUMES")
    print("=" * 60)

    volume_names = list_ids(["docker", "volume", "ls", "-q"])
    if not volume_names:
        print("ℹ No volumes found.")
        return

    print(f"ℹ Found {len(volume_names)} volume(s)")
    print("\nℹ Removing all volumes...")
    progress = ProgressBar(len(volume_names), "removing volumes")

    for volume in volume_names:
        run_cmd_realtime(["docker", "volume", "rm", "-f", volume], show_output=False)
        progress.update()
    progress.complete()

    print("✓ All volumes removed")


def prune_system():
    print("\n" + "=" * 60)
    print("STEP 5/8: SYSTEM PRUNE & BUILD CACHE")
    print("=" * 60)

    print("ℹ Pruning builder cache...")
    run_cmd_realtime(["docker", "builder", "prune", "-af"])

    print("\nℹ Running system prune (extra safety)...")
    run_cmd_realtime(["docker", "system", "prune", "-af", "--volumes"])


def stop_docker_service():
    print("\n" + "=" * 60)
    print("STEP 6/8: STOP DOCKER DAEMON")
    print("=" * 60)

    print("ℹ Stopping Docker service...")
    run_cmd_realtime(["systemctl", "stop", "docker"])

    print("\nℹ Stopping containerd service...")
    run_cmd_realtime(["systemctl", "stop", "containerd"])

    print("✓ Docker daemon stopped")


def wipe_docker_data_dirs():
    print("\n" + "=" * 60)
    print("STEP 7/8: DELETE DOCKER DATA DIRECTORIES")
    print("=" * 60)

    paths = ["/var/lib/docker", "/var/lib/containerd"]
    existing_paths = [p for p in paths if os.path.exists(p)]

    if not existing_paths:
        print("ℹ No Docker data directories found.")
        return

    print(f"ℹ Found {len(existing_paths)} data director(ies) to delete")
    progress = ProgressBar(len(existing_paths), "deleting directories")

    for p in paths:
        if os.path.exists(p):
            print(f"\nℹ Deleting {p} ...")
            code, out, _err = run_cmd_quiet(["du", "-sh", p])
            if code == 0 and out:
                print(f"  Size: {out.split()[0]}")

            run_cmd_realtime(["rm", "-rf", p], show_output=False)
            progress.update()
        else:
            print(f"ℹ {p} does not exist; skipping.")

    progress.complete()
    print("✓ All Docker data directories deleted")


def start_docker_service():
    print("\n" + "=" * 60)
    print("STEP 8/8: START DOCKER DAEMON")
    print("=" * 60)

    print("ℹ Starting containerd service...")
    run_cmd_realtime(["systemctl", "start", "containerd"])

    print("\nℹ Starting Docker service...")
    run_cmd_realtime(["systemctl", "start", "docker"])

    print("\nℹ Waiting for Docker to initialize...")
    for i in range(5):
        time.sleep(1)
        code, out, _err = run_cmd_quiet(["systemctl", "is-active", "docker"])
        if code == 0 and out.strip() == "active":
            print(f"✓ Docker service is active (took {i+1}s)")
            break
    else:
        print("⚠ Warning: Docker service may not be active yet.")

    print("\nℹ Running docker info to verify...")
    run_cmd_realtime(["docker", "info"])


def show_final_state():
    print("\n" + "=" * 60)
    print("FINAL DOCKER STATE")
    print("=" * 60)

    cmds = [
        (["docker", "ps", "-a"], "CONTAINERS"),
        (["docker", "images"], "IMAGES"),
        (["docker", "network", "ls"], "NETWORKS"),
        (["docker", "volume", "ls"], "VOLUMES"),
    ]

    for cmd, title in cmds:
        print(f"\n▶ {title}")
        print("─" * 60)
        code, out, err = run_cmd_quiet(cmd)
        if code == 0:
            print(out if out else "(empty)")
        else:
            print("✗ Command failed")
            if err:
                print(err)


def show_banner():
    banner = """
╔════════════════════════════════════════════════════════════╗
║                                                            ║
║        DOCKER FACTORY RESET FOR UBUNTU                     ║
║                                                            ║
║        Completely resets Docker to fresh state             ║
║                                                            ║
╚════════════════════════════════════════════════════════════╝
"""
    print(banner)


def main():
    show_banner()
    require_root()

    if not docker_installed():
        sys.exit(1)

    print("\n⚠ WARNING: This script will COMPLETELY RESET Docker on this Ubuntu host.")
    print("\nIt will:")
    print("  • Stop and remove ALL containers (running and stopped)")
    print("  • Remove ALL images")
    print("  • Remove ALL volumes")
    print("  • Remove ALL custom networks")
    print("  • Prune ALL build cache and system data")
    print("  • STOP Docker daemon")
    print("  • DELETE /var/lib/docker and /var/lib/containerd")
    print("  • START Docker daemon again")
    print("\n⚠ You will lose ALL Docker data on this machine.")
    print("\n" + "─" * 60)

    confirm = input("\nType 'RESET_DOCKER_UBUNTU' to continue: ")
    if confirm != "RESET_DOCKER_UBUNTU":
        print("\n✗ Aborted.")
        sys.exit(0)

    start_time = time.time()

    stop_and_remove_containers()
    remove_images()
    remove_networks()
    remove_volumes()
    prune_system()
    stop_docker_service()
    wipe_docker_data_dirs()
    start_docker_service()
    show_final_state()

    elapsed_time = time.time() - start_time

    print("\n" + "=" * 60)
    print("✓ DOCKER FACTORY RESET COMPLETE")
    print(f"  Total time: {elapsed_time:.2f} seconds")
    print("=" * 60)
    print("\nDocker has been reset to a fresh installation state.\n")


if __name__ == "__main__":
    main()