path: root/dots.py

#!/usr/bin/env python3

from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Optional
import argparse
import json
import logging
import os
import subprocess
import sys
import shutil
from concurrent.futures import Executor, ThreadPoolExecutor, as_completed
import jinja2


@dataclass(frozen=True)
class Config:
    SCRIPT_DIR: Path = Path(__file__).parent / "home/scripts"
    COMPILED_DIR: Path = Path(".compiled_dotfiles")
    CONTEXTS: Path = Path("context.json")
    DEFAULT_TARGET_DIR: Path = Path.home()
    DEFAULT_SOURCE_DIR: Path = Path(".")
    MAX_WORKERS: int = (os.cpu_count() or 1) * 2


@dataclass(frozen=True)
class Environment:
    platform: str
    device_name: str
    context: Dict
    verbose: bool
    logger: logging.Logger


def main() -> None:
    args = parse_arguments()
    env = initialize_environment(args.verbose)

    if not (args.compile or args.stow or args.clean):
        env.logger.error("At least one of --compile, --stow, or --clean must be specified")
        sys.exit(1)

    if args.compile:
        env.logger.debug(f"Compiling {args.source} -> {args.comp}")
        if not compile_dotfiles(args.source, args.comp, env):
            env.logger.error("Failed to compile dotfiles")
            sys.exit(1)

    if args.stow:
        env.logger.debug(f"Stowing {(args.target)}")
        if not stow_dotfiles(args.comp, args.target, env, clean=False):
            env.logger.error("Failed to stow dotfiles")
            sys.exit(1)

    if args.clean:
        env.logger.debug(f"Cleaning dotfiles from {(args.target)}")
        if not stow_dotfiles(args.comp, args.target, env, clean=True):
            env.logger.error("Failed to clean dotfiles")
            sys.exit(1)

    env.logger.info("Done!")
    sys.exit(0)


def initialize_environment(verbose: bool) -> Environment:
    logger = setup_logging(verbose)
    platform = get_platform()
    device_name = get_device_name()
    context = {
        **load_context(platform, device_name, Config.CONTEXTS),
        "platform": platform,
        "device_name": device_name,
    }

    return Environment(platform, device_name, context, verbose, logger)


def copy_with_templates_rendered(executor: Executor, source: Path, destination: Path, env: Environment) -> bool:
    shutil.copytree(source, destination, dirs_exist_ok=True)
    env.logger.debug(f"copytree {source} -> {destination}")

    templates = [t for t in destination.glob("**/*.j2") if t.is_file()]
    if not templates:
        env.logger.debug(f"{source} has no templates")
        return True

    for template in templates:
        env.logger.debug(f"submitting work to render {source} . {templates}")
        executor.submit(replace_with_rendered_template, template, env)

    return True

def replace_with_rendered_template(template: Path, env: Environment) -> bool:
    env.logger.debug(f"rendering {template}")
    jinja_env = jinja2.Environment(
        undefined=jinja2.StrictUndefined, trim_blocks=True, lstrip_blocks=True
    )
    with open(template, "w") as t:
        t.write(render(template, jinja_env, env) or "")
    env.logger.debug(f"stripping template suffix of {template}")
    template.rename(template.with_suffix(""))
    return True



def compile_dotfiles(source_dir: Path, target_dir: Path, env: Environment):
    target_dir.mkdir(exist_ok=True, parents=True)

    with ThreadPoolExecutor(max_workers=Config.MAX_WORKERS) as executor:
        dotfile_directories = [
            dir for dir in source_dir.iterdir() if dir.is_dir() and "." not in dir.name
        ]
        env.logger.debug(f"compiling dotfile directories {dotfile_directories}")
        futures = [
            executor.submit(
                copy_with_templates_rendered, executor, dotfiles, target_dir / dotfiles.name, env
            )
            for dotfiles in dotfile_directories
        ]
        env.logger.debug(f"constructed {len(futures)} tasks")
        return sum(1 for future in as_completed(futures) if future.result()) == len(
            futures
        )
    return False


def render(
    source: Path,
    jinja_env: jinja2.Environment,
    env: Environment,
) -> Optional[str]:
    try:
        env.logger.debug(f"Reading {source}")
        with open(source, "r") as f:
            template_content = f.read()
        env.logger.debug(f"Compiling template {source}")
        template = jinja_env.from_string(template_content)
        return template.render(**env.context)
    except Exception as e:
        env.logger.error(f"Error rendering template {source}: {e}")
        return None


def stow_package_task(
    stow_dir: Path, package: Path, target_dir: Path, env: Environment
) -> bool:
    try:
        return True
    except subprocess.SubprocessError as e:
        env.logger.error(f"Error stowing package {package.name}: {e}")
        return False


def stow_dotfiles(source_dir: Path, target_dir: Path, env: Environment, clean: bool = False) -> bool:
    if not stow_installed():
        return False

    packages = [d for d in source_dir.iterdir() if d.is_dir()]
    stow_op = ["-D"] if clean else ["--no-folding"]

    for pkg in packages:
        run_shell_command(["stow", *stow_op, pkg.name])
    return True


def stow_installed() -> bool:
    return "" != run_shell_command(["stow", "--version"])


def run_shell_command(cmd: List[str]) -> str:
    result = subprocess.run(cmd, capture_output=True, text=True, check=True)
    return result.stdout.strip()


def get_platform() -> str:
    return run_shell_command([str(Config.SCRIPT_DIR / "platform.sh")])


def get_device_name() -> str:
    return run_shell_command([str(Config.SCRIPT_DIR / "system_name.sh")])


def load_context(platform: str, device_name: str, context_file: Path) -> Dict:
    try:
        with open(context_file) as f:
            contexts = json.load(f)

        if platform not in contexts:
            return {}

        return contexts[platform].get(
            device_name, contexts[platform].get("default", {})
        )
    except (FileNotFoundError, json.JSONDecodeError):
        return {}


def parse_arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Dotfiles manager")
    parser.add_argument("--compile", action="store_true", help="Compile dotfiles")
    parser.add_argument("--stow", action="store_true", help="Stow compiled dotfiles")
    parser.add_argument("--clean", action="store_true", help="Clean (remove) stowed dotfiles")
    parser.add_argument(
        "--source",
        help=f"Target directory for dotfiles (default: {Config.DEFAULT_SOURCE_DIR})",
        default=Config.DEFAULT_SOURCE_DIR,
    )
    parser.add_argument(
        "--comp",
        help=f"Target directory for compiled templates (default: {Config.COMPILED_DIR})",
        default=Config.COMPILED_DIR,
    )
    parser.add_argument(
        "--target",
        help=f"Target directory for stow (default: {Config.DEFAULT_TARGET_DIR})",
        default=Config.DEFAULT_TARGET_DIR,
    )
    parser.add_argument(
        "--verbose", "-v", action="store_true", help="Enable verbose logging"
    )
    return parser.parse_args()


def setup_logging(verbose: bool) -> logging.Logger:
    logging.basicConfig(
        level=logging.DEBUG if verbose else logging.INFO,
        format="%(asctime)s - %(levelname)s - %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )
    return logging.getLogger("dotfiles")


if __name__ == "__main__":
    main()