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 main() -> None:
    """Main entry point for the script."""
    parser = argparse.ArgumentParser(
        description="Development container utility",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    subparsers = parser.add_subparsers(dest="command", help="Command to execute")

    # Init command
    init_parser = subparsers.add_parser("init", help="Initialize a new container")
    init_parser.add_argument("path", nargs="?", help="Path to the .container directory")
    init_parser.add_argument("--image", required=True, help="Docker image to use")
    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",
    )

    # Activate command
    activate_parser = subparsers.add_parser("activate", help="Activate container shell")
    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"
    )
    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")
    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")
    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")
    rm_parser.add_argument(
        "container_name", nargs="?", help="Name of the container to remove"
    )

    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)
    else:
        parser.print_help()
        sys.exit(1)


if __name__ == "__main__":
    main()