path: root/dots_manager/.main.py

#!/usr/bin/env python3

import sys

sys.dont_write_bytecode = True

from dataclasses import dataclass
from pathlib import Path
from typing import (
    Dict,
    Generator,
    List,
    Optional,
    Any,
    Literal,
    ParamSpec,
    Tuple,
    Callable,
    TypeVar,
)
import argparse
import json
import logging
import os
import subprocess
import sys
import shutil
from concurrent.futures import Executor, ThreadPoolExecutor, as_completed
from functools import reduce
import jinja2
from kawaii_logger import setup_logger

R = TypeVar("R")
T = TypeVar("T")
P = ParamSpec("P")


@dataclass(frozen=True)
class Config:
    default_target_dir: Path = Path.home()
    default_source_dir: Path = Path(__file__).parent / "dots"
    default_compiled_dir: Path = Path("./.compiled_dotfiles/")

    script_dir: Path = Path("home/scripts")
    contexts: Path = Path("contexts.json")

    template_extension: str = ".j2"
    max_workers: int = (os.cpu_count() or 1) * 2


@dataclass(frozen=True)
class Environment:
    platform: str
    system_name: str
    context: Dict[str, Any]
    logger: logging.Logger


def main() -> None:
    args = parse_arguments()
    env = initialize_environment(args)
    if not env:
        sys.exit(1)
    if not any((args.compile, args.stow, args.clean)):
        env.logger.error("you gotta specify at least one action nya~ (⁎⁍̴̆‾⁍̴̆⁎)")
        sys.exit(1)

    if args.compile and not compile_dotfiles(Path(args.source), Path(args.comp), env):
        env.logger.error("uh oh! failed to compile dotfiles (ɐ•́ﾟ•̀ɐ)")
        sys.exit(1)

    if args.stow and not apply_stow_operation_to_packages(
        Path(args.comp), Path(args.target), "--no-folding", env
    ):
        env.logger.error("failed to stow dotfiles... nyaaa (╥ﾟ╥)")
        sys.exit(1)

    if args.clean and not apply_stow_operation_to_packages(
        Path(args.comp), Path(args.target), "-D", env
    ):
        env.logger.error("couldn’t clean dotfiles... sobs (ɐ•́ﾟ•̀ɐ)")
        sys.exit(1)

    env.logger.info("yay~ all done!! ₕᵒ. .ᵒₕ♡")
    sys.exit(0)


def initialize_environment(args) -> Environment:
    logger = setup_logger(args.verbose)

    scripts = args.source / Config.script_dir
    platform = run_shell_command([str(scripts / "platform.sh")], logger)
    if not platform:
        raise ValueError("failed to determine platform... ")

    os.environ["PLATFORM"] = platform
    system_name = run_shell_command([str(scripts / "system_name.sh")], logger)
    if not system_name:
        raise ValueError("failed to determine system name... ")

    context = load_context(platform, system_name, Config.contexts, logger)
    return Environment(platform, system_name, context, logger)


def render_template_to(
    template: Path, destination: Path, env: Environment
) -> Optional[Path]:
    env.logger.debug(f"rendering template {template} to {destination} ✧*ฺ")
    if not is_template(template):
        env.logger.error(f"template {template} is nyot a vawid tempwate D:")
        return None

    env.logger.debug(f"rendering template {template} ✧*ฺ")
    jinja_env = jinja2.Environment(
        loader=jinja2.BaseLoader,
        undefined=jinja2.StrictUndefined,
        trim_blocks=True,
        lstrip_blocks=True,
    )
    rendered = render(template, jinja_env, env)
    if not rendered:
        env.logger.error(f"failed to render template {template}")
        return None

    env.logger.debug(f"rendered {template} successfully~ :D ✧*ฺ")
    destination.write_text(rendered)
    return destination


def compile_dotfiles(source: Path, output: Path, env: Environment) -> bool:
    env.logger.info(f"compiling {source} to {output} ✨")
    output.mkdir(exist_ok=True, parents=True)
    shutil.copytree(source, output, dirs_exist_ok=True)
    env.logger.debug(f"copied {source} to {output} ₰˜⋉♡")

    tasks = parallelize(
        lambda template: is_some(
            render_template_to,
            template,
            template.with_suffix(""),
            env,
        )
        and (template.unlink() is None),
        find_templates_under(output),
        env,
    )
    return all(tasks)


def render(
    source: Path, jinja_env: jinja2.Environment, env: Environment
) -> Optional[str]:
    try:
        env.logger.debug(f"rendering template {source} ✿.｡.:･")
        return jinja_env.from_string(source.read_text()).render(**env.context)
    except Exception as e:
        env.logger.error(f"couldn’t render {source}: {e} (；⌓̀_⌓́)")
        return None


def apply_stow_operation_to_packages(
    packages: Path,
    target: Path,
    stow_op: Literal["-D", "--no-folding"],
    env: Environment,
) -> bool:
    if not run_shell_command(["stow", "--version"], env.logger):
        env.logger.error("stow not installed D:")
        return False
    package_command = ["stow", "-d", str(packages), "-t", str(target), stow_op]

    _packages = list_stowable_packages(packages)
    env.logger.debug(f"found dotfile packages: {_packages}")
    commands = [package_command + [package.name] for package in _packages]

    env.logger.debug(f"stowing packages: {commands}")
    results = parallelize(
        lambda command: is_some(run_shell_command, command, env.logger),
        commands,
        env,
    )
    return len(commands) == sum(1 if x else 0 for x in results)


def parallelize(
    worker: Callable[[T], R],
    items: List[T],
    env: Environment,
    executor: Optional[Executor] = None,
) -> List[R]:
    if executor is None:
        executor = ThreadPoolExecutor(max_workers=Config.max_workers)
    with executor as executor:
        futures = [
            executor.submit(
                worker,
                item,
            )
            for item in items
        ]
        env.logger.info(f"submitted {len(futures)} tasks to executor ₰˜.༄")
        return [f.result() for f in as_completed(futures)]


def run_shell_command(cmd: List[str], logger: logging.Logger) -> Optional[str]:
    logger.debug(f"running command: {cmd}")
    try:
        return subprocess.run(
            cmd, capture_output=True, text=True, check=True
        ).stdout.strip()
    except subprocess.CalledProcessError as e:
        logger.error(f"command failed: {cmd}, {e.stderr}")
        return None


def is_template(path: Path) -> bool:
    return path.suffix == Config.template_extension


def find_templates_under(destination: Path) -> List[Path]:
    return [
        t
        for t in destination.glob("**/*" + Config.template_extension)
        if t.is_file() and is_template(t)
    ]


def list_stowable_packages(packages: Path) -> List[Path]:
    denylist = [".", "__"]
    return [
        d
        for d in packages.iterdir()
        if d.is_dir() and all(y not in d.name for y in denylist)
    ]


def load_context(
    platform: str, system_name: str, context_file: Path, logger: logging.Logger
) -> Dict[str, Any]:
    try:
        logger.info(f"reading context file: {context_file} ✧*:｡ﾟ✧")
        with open(context_file) as f:
            contexts = json.load(f)

        global_config = contexts.get("_global", {})
        platform_defaults = contexts.get(platform, {}).get("_default", {})
        defaults = merge_dicts(global_config, platform_defaults)

        system_config = contexts.get(platform, {}).get(system_name, {})
        if not system_config:
            logger.warning(
                f"couldn’t find system-specific config for {platform}.{system_name} (ɐ•́ﾟ•̀ɐ)"
            )

        return {
            "platform": platform,
            "system_name": system_name,
            **merge_dicts(defaults, system_config),
        }
    except (FileNotFoundError, json.JSONDecodeError) as e:
        logger.error(f"error loading context: {e} ⋆ฺ°☁｡⋆ฺ °★ °｡")
        return {}


def is_some(
    callable: Callable[P, Optional[T]], *args: P.args, **kwargs: P.kwargs
) -> Tuple[bool, Optional[T]]:
    result = callable(*args, **kwargs)
    return result is not None, result


def merge_dicts(*dicts: Dict[str, Any]) -> Dict[str, Any]:
    def merge(a: Dict[str, Any], b: Dict[str, Any]) -> Dict[str, Any]:
        out = dict(a)
        for k, v in b.items():
            if k in out and isinstance(out[k], dict) and isinstance(v, dict):
                out[k] = merge(out[k], v)
            else:
                out[k] = v
        return out

    return reduce(merge, dicts, {})


def parse_arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="cute dotfiles manager ✧˖°")

    parser.add_argument(
        "--compile", action="store_true", help="action: compile dotfiles"
    )
    parser.add_argument(
        "--source",
        type=Path,
        default=Config.default_source_dir,
        help=f"directory with stowable dotfiles. default '{Config.default_source_dir}'. :)",
    )
    parser.add_argument(
        "--comp",
        type=Path,
        default=Config.default_compiled_dir,
        help=f"compiled template output dir. default '{Config.default_compiled_dir}'. :3",
    )

    parser.add_argument(
        "--stow", action="store_true", help="action: stow compiled dotfiles"
    )
    parser.add_argument(
        "--clean", action="store_true", help="action: clean stowed dotfiles"
    )
    parser.add_argument(
        "--target",
        type=Path,
        default=Config.default_target_dir,
        help=f"stow target directory. default '{Config.default_target_dir}'. -.-",
    )

    parser.add_argument(
        "--verbose",
        "-v",
        action="store_true",
        help="enable verbose logging. default False. :D",
        default=False,
    )
    return parser.parse_args()


if __name__ == "__main__":
    main()