deu/deu.py

#!/usr/bin/env python3

import sys
import logging
import subprocess
import time
from pathlib import Path
from typing import Dict, Any, Optional, List
import yaml
import toml
import argparse
from dataclasses import dataclass
from enum import Enum, auto
import random
import string

# Configure logging
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


class ShellType(Enum):
    BASH = auto()
    SH = auto()
    UNKNOWN = auto()


@dataclass
class ContainerConfig:
    path: Path
    service: str
    image: str
    shell: ShellType = ShellType.UNKNOWN


def validate_path(path: Path) -> None:
    """Validate the provided path."""
    if not path:
        raise ValueError("Path cannot be empty")

    if not path.is_absolute():
        path = path.absolute()

    if not path.parent.exists():
        raise FileNotFoundError(f"Parent directory {path.parent} does not exist")


def validate_image(image: str) -> None:
    """Validate the Docker image name."""
    if not image:
        raise ValueError("Docker image cannot be empty")

    # Basic validation for image format
    if not any(char in image for char in [":", "@"]):
        logger.warning("No tag or digest specified for the image. Using 'latest' tag")


def check_image_exists(image: str) -> bool:
    """Check if the Docker image exists locally."""
    try:
        subprocess.run(
            ["docker", "image", "inspect", image], check=True, capture_output=True
        )
        return True
    except subprocess.CalledProcessError:
        return False


def pull_image(image: str) -> None:
    """Pull the Docker image if it doesn't exist locally."""
    if not check_image_exists(image):
        logger.info(f"Pulling image {image}...")
        try:
            subprocess.run(["docker", "pull", image], check=True)
        except subprocess.CalledProcessError as e:
            logger.error(f"Failed to pull image {image}: {e}")
            raise


def detect_shell(image: str) -> ShellType:
    """Detect available shell in the image."""
    try:
        # Try bash first
        subprocess.run(
            ["docker", "run", "--rm", image, "which", "bash"],
            check=True,
            capture_output=True,
        )
        return ShellType.BASH
    except subprocess.CalledProcessError:
        try:
            # Fallback to sh
            subprocess.run(
                ["docker", "run", "--rm", image, "which", "sh"],
                check=True,
                capture_output=True,
            )
            return ShellType.SH
        except subprocess.CalledProcessError:
            return ShellType.UNKNOWN


def create_docker_compose_config(
    workspace_path: Path, image: str, service_name: str
) -> Dict[str, Any]:
    """Create the Docker Compose configuration."""
    return {
        "services": {
            service_name: {
                "image": image,
                "command": "tail -f /dev/null",  # Keep container running with a shell
                "volumes": [f"{workspace_path.absolute()}:/workspace"],
                "working_dir": "/workspace",
                "stdin_open": True,
                "tty": True,
                "environment": {
                    "TERM": "xterm-256color",
                },
                "security_opt": ["seccomp=unconfined"],
            }
        }
    }


def create_toml_config(
    container_path: Path, service_name: str, image: str, shell: ShellType
) -> Dict[str, Any]:
    """Create the TOML configuration."""
    return {
        "container": {
            "path": str(container_path.absolute()),
            "service": service_name,
            "image": image,
            "shell": shell.name,
        }
    }


def write_config(path: Path, config: Dict[str, Any], config_type: str) -> None:
    """Write configuration to a file."""
    try:
        if config_type == "yaml":
            path.write_text(yaml.dump(config, sort_keys=False))
        elif config_type == "toml":
            path.write_text(toml.dumps(config))
        logger.info(f"Created {config_type.upper()} file at: {path}")
    except Exception as e:
        logger.error(f"Failed to write {config_type.upper()} file: {e}")
        raise


def run_docker_compose_command(
    container_path: Path, command: str, detach: bool = False
) -> None:
    """Run a Docker Compose command with retries."""
    max_retries = 3
    retry_delay = 1

    for attempt in range(max_retries):
        try:
            cmd = [
                "docker",
                "compose",
                "-f",
                str(container_path / "docker-compose.yml"),
            ]
            cmd.extend(command.split())
            if detach:
                cmd.append("-d")

            subprocess.run(cmd, check=True)
            return
        except subprocess.CalledProcessError as e:
            if attempt == max_retries - 1:
                logger.error(
                    f"Failed to execute Docker Compose command after {max_retries} attempts: {e}"
                )
                raise
            logger.warning(f"Command failed, retrying in {retry_delay} seconds...")
            time.sleep(retry_delay)
            retry_delay *= 2


def get_container_config() -> ContainerConfig:
    """Get the container configuration from deu.toml."""
    try:
        config = toml.loads((Path.cwd() / "deu.toml").read_text())
        container_config = config["container"]
        return ContainerConfig(
            path=Path(container_config["path"]),
            service=container_config["service"],
            image=container_config["image"],
            shell=ShellType[container_config["shell"]],
        )
    except Exception as e:
        logger.error(f"Failed to read container configuration: {e}")
        sys.exit(1)


def get_global_config_dir() -> Path:
    """Get the global configuration directory."""
    return Path.home() / ".config" / "deu"


def find_container_config(
    container_name: Optional[str] = None,
) -> Optional[ContainerConfig]:
    """Find container configuration, checking local then global."""
    # If container name is provided, search for it
    if container_name:
        # First check local deu.toml
        local_toml = Path.cwd() / "deu.toml"
        if local_toml.exists():
            try:
                config = toml.loads(local_toml.read_text())
                if config["container"]["service"] == container_name:
                    return get_container_config()
            except Exception:
                pass

        # Then check global configs
        global_dir = get_global_config_dir()
        if global_dir.exists():
            for toml_file in global_dir.glob("*.toml"):
                try:
                    config = toml.loads(toml_file.read_text())
                    if config["container"]["service"] == container_name:
                        return ContainerConfig(
                            path=Path(config["container"]["path"]),
                            service=config["container"]["service"],
                            image=config["container"]["image"],
                            shell=ShellType[config["container"]["shell"]],
                        )
                except Exception:
                    continue
        return None

    # If no container name, just check local deu.toml
    local_toml = Path.cwd() / "deu.toml"
    if local_toml.exists():
        return get_container_config()
    return None


def create_dev_container(
    path: str, image: str, service_name: str, is_global: bool = False
) -> None:
    """Create a development container configuration."""
    try:
        if is_global:
            # Create global config directory first
            global_dir = get_global_config_dir()
            global_dir.mkdir(parents=True, exist_ok=True)
            container_path = global_dir / f".{service_name}"
            toml_path = global_dir / f"{service_name}.toml"
        else:
            container_path = Path(path)
            toml_path = Path.cwd() / "deu.toml"

        validate_image(image)

        # Ensure image exists and detect shell
        pull_image(image)
        shell_type = detect_shell(image)
        if shell_type == ShellType.UNKNOWN:
            logger.error(
                "No shell (bash/sh) found in the image. Please use an image with a shell."
            )
            sys.exit(1)

        # Create the container directory
        container_path.mkdir(parents=True, exist_ok=True)

        # Create configurations
        workspace_path = container_path.parent
        compose_config = create_docker_compose_config(
            workspace_path, image, service_name
        )
        toml_config = create_toml_config(
            container_path, service_name, image, shell_type
        )

        # Write configurations
        write_config(container_path / "docker-compose.yml", compose_config, "yaml")
        write_config(toml_path, toml_config, "toml")

    except Exception as e:
        logger.error(f"Failed to create development container: {e}")
        sys.exit(1)


def generate_service_name() -> str:
    """Generate a random, friendly service name."""
    adjectives = [
        "happy",
        "sleepy",
        "daisy",
        "sunny",
        "lucky",
        "fuzzy",
        "cozy",
        "merry",
    ]
    nouns = ["dog", "cat", "bear", "fox", "bird", "fish", "tree", "star"]
    random_number = random.randint(1000, 9999)
    return f"{random.choice(adjectives)}_{random.choice(nouns)}_{random_number}"


def handle_activate(container_name: Optional[str] = None) -> None:
    """Activate the container shell."""
    config = find_container_config(container_name)
    if not config:
        if container_name:
            logger.error(f"Container '{container_name}' not found locally or globally.")
        else:
            logger.error(
                "No local container configuration found. Run 'deu init' first or specify a container name."
            )
        sys.exit(1)

    print(f"Trying to activate container shell for service '{config.service}'...")

    # Check if container is running, start it if not
    try:
        ps_output = subprocess.run(
            ["docker", "compose", "-f", str(config.path / "docker-compose.yml"), "ps"],
            check=True,
            capture_output=True,
            text=True,
        )
        if config.service not in ps_output.stdout:
            logger.info(f"Container '{config.service}' is not running. Starting it...")
            run_docker_compose_command(config.path, "up -d")
    except subprocess.CalledProcessError:
        logger.info(f"Container '{config.service}' is not running. Starting it...")
        run_docker_compose_command(config.path, "up -d")

    # Use appropriate shell based on configuration
    shell = "bash" if config.shell == ShellType.BASH else "sh"
    run_docker_compose_command(config.path, f"exec {config.service} {shell}")


def handle_background(container_name: Optional[str] = None) -> None:
    """Start the container in background mode."""
    config = find_container_config(container_name)
    if not config:
        if container_name:
            logger.error(f"Container '{container_name}' not found locally or globally.")
        else:
            logger.error(
                "No local container configuration found. Run 'deu init' first or specify a container name."
            )
        sys.exit(1)
    run_docker_compose_command(config.path, "up", detach=True)


def handle_logs(container_name: Optional[str] = None) -> None:
    """Show container logs."""
    config = find_container_config(container_name)
    if not config:
        if container_name:
            logger.error(f"Container '{container_name}' not found locally or globally.")
        else:
            logger.error(
                "No local container configuration found. Run 'deu init' first or specify a container name."
            )
        sys.exit(1)
    run_docker_compose_command(config.path, "logs -f")


def handle_stop(container_name: Optional[str] = None) -> None:
    """Stop the container."""
    config = find_container_config(container_name)
    if not config:
        if container_name:
            logger.error(f"Container '{container_name}' not found locally or globally.")
        else:
            logger.error(
                "No local container configuration found. Run 'deu init' first or specify a container name."
            )
        sys.exit(1)
    run_docker_compose_command(config.path, "stop")


def handle_rm(container_name: Optional[str] = None) -> None:
    """Remove the container."""
    config = find_container_config(container_name)
    if not config:
        if container_name:
            logger.error(f"Container '{container_name}' not found locally or globally.")
        else:
            logger.error(
                "No local container configuration found. Run 'deu init' first or specify a container name."
            )
        sys.exit(1)
    run_docker_compose_command(config.path, "down")


def handle_list() -> None:
    """List all available containers (local and global)."""
    # Check for local container
    local_toml = Path.cwd() / "deu.toml"
    if local_toml.exists():
        try:
            config = toml.loads(local_toml.read_text())
            print("Local container:")
            print(f"  {config['container']['service']} (local)")
        except Exception as e:
            logger.error(f"Failed to read local configuration: {e}")

    # Check global containers
    global_dir = get_global_config_dir()
    if global_dir.exists():
        global_containers = []
        for toml_file in global_dir.glob("*.toml"):
            try:
                config = toml.loads(toml_file.read_text())
                global_containers.append(config["container"]["service"])
            except Exception:
                continue

        if global_containers:
            print("\nGlobal containers:")
            for container in sorted(global_containers):
                print(f"  {container} (global)")
        elif not local_toml.exists():
            print("No containers found (local or global).")


def handle_delete(container_name: Optional[str] = None) -> None:
    """Stop and remove a container."""
    config = find_container_config(container_name)
    if not config:
        if container_name:
            logger.error(f"Container '{container_name}' not found locally or globally.")
        else:
            logger.error(
                "No local container configuration found. Run 'deu init' first or specify a container name."
            )
        sys.exit(1)

    # Stop the container first
    try:
        run_docker_compose_command(config.path, "stop")
    except Exception as e:
        logger.warning(f"Failed to stop container: {e}")

    # Then remove it
    try:
        run_docker_compose_command(config.path, "down")
        logger.info(f"Container '{config.service}' has been deleted.")
    except Exception as e:
        logger.error(f"Failed to remove container: {e}")
        sys.exit(1)


def handle_examples() -> None:
    """Show common usage examples."""
    print(
        """
Common DEU Usage Examples:

1. Create a new development container:
   deu init --image ubuntu:24.04
   deu init -g --image python:3.11
   deu init --image node:20 --service my_node_app

2. Work with containers:
   deu activate              # Start shell in local container
   deu activate my_container # Start shell in specific container
   deu background            # Start container in background
   deu logs                  # View container logs

3. Manage containers:
   deu stop                  # Stop local container
   deu rm                    # Remove local container
   deu delete                # Stop and remove local container
   deu list                  # Show all available containers

4. Global containers:
   deu init -g --image ubuntu:24.04  # Create global container
   deu activate my_global            # Use global container
   deu delete my_global              # Remove global container

5. Container naming:
   deu init --image ubuntu:24.04 --service my_dev
   deu activate my_dev
   deu delete my_dev
"""
    )


def main() -> None:
    """Main entry point for the script."""
    parser = argparse.ArgumentParser(
        description="DEU - Docker Environment Utility\n\n"
        "A simple utility to create and manage development containers.\n"
        "Use 'deu examples' to see common usage patterns.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    subparsers = parser.add_subparsers(dest="command", help="Command to execute")

    # Init command
    init_parser = subparsers.add_parser(
        "init",
        help="Initialize a new development container",
        description="Create a new development container with the specified image.\n"
        "Example: deu init --image ubuntu:24.04",
    )
    init_parser.add_argument("path", nargs="?", help="Path to the .container directory")
    init_parser.add_argument(
        "--image", required=True, help="Docker image to use (e.g., ubuntu:24.04)"
    )
    init_parser.add_argument(
        "--service",
        default=generate_service_name(),
        help="Name of the container service",
    )
    init_parser.add_argument(
        "-g",
        "--global",
        dest="is_global",
        action="store_true",
        help="Create container in global config (~/.config/deu/)",
    )

    # Activate command
    activate_parser = subparsers.add_parser(
        "activate",
        help="Activate container shell",
        description="Start an interactive shell in the container.\n"
        "Example: deu activate my_container",
    )
    activate_parser.add_argument(
        "container_name", nargs="?", help="Name of the container to activate"
    )

    # Background command
    background_parser = subparsers.add_parser(
        "background",
        help="Start container in background",
        description="Start the container in detached mode.\n"
        "Example: deu background my_container",
    )
    background_parser.add_argument(
        "container_name", nargs="?", help="Name of the container to start"
    )

    # Logs command
    logs_parser = subparsers.add_parser(
        "logs",
        help="Show container logs",
        description="View container logs in follow mode.\n"
        "Example: deu logs my_container",
    )
    logs_parser.add_argument(
        "container_name", nargs="?", help="Name of the container to show logs for"
    )

    # Stop command
    stop_parser = subparsers.add_parser(
        "stop",
        help="Stop container",
        description="Stop a running container.\n" "Example: deu stop my_container",
    )
    stop_parser.add_argument(
        "container_name", nargs="?", help="Name of the container to stop"
    )

    # Remove command
    rm_parser = subparsers.add_parser(
        "rm",
        help="Remove container",
        description="Remove a container and its configuration.\n"
        "Example: deu rm my_container",
    )
    rm_parser.add_argument(
        "container_name", nargs="?", help="Name of the container to remove"
    )

    # List command
    list_parser = subparsers.add_parser(
        "list",
        help="List all available containers",
        description="Show all containers (local and global).\n" "Example: deu list",
    )

    # Delete command
    delete_parser = subparsers.add_parser(
        "delete",
        help="Stop and remove a container",
        description="Stop and remove a container in one command.\n"
        "Example: deu delete my_container",
    )
    delete_parser.add_argument(
        "container_name", nargs="?", help="Name of the container to delete"
    )

    # Examples command
    subparsers.add_parser(
        "examples",
        help="Show common usage examples",
        description="Display common usage patterns and examples.",
    )

    args = parser.parse_args()

    if args.command == "init":
        create_dev_container(
            args.path or ".container", args.image, args.service, args.is_global
        )
    elif args.command == "activate":
        handle_activate(args.container_name)
    elif args.command == "background":
        handle_background(args.container_name)
    elif args.command == "logs":
        handle_logs(args.container_name)
    elif args.command == "stop":
        handle_stop(args.container_name)
    elif args.command == "rm":
        handle_rm(args.container_name)
    elif args.command == "list":
        handle_list()
    elif args.command == "delete":
        handle_delete(args.container_name)
    elif args.command == "examples":
        handle_examples()
    else:
        parser.print_help()
        sys.exit(1)


if __name__ == "__main__":
    main()