#!/usr/bin/env python3

from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Optional, Any
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


@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
    system_name: str
    context: Dict[str, Any]
    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("you gotta specify at least one action nya~ (⁎⁍̴̆‾⁍̴̆⁎)")
        sys.exit(1)

    if args.compile:
        env.logger.info(f"compiling {args.source} to {args.comp} ✨✧˖°")
        if 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:
        if not stow_dotfiles(Path(args.comp), Path(args.target), env, clean=False):
            env.logger.error("failed to stow dotfiles... nyaaa (╥ﾟ╥)")
            sys.exit(1)

    if args.clean:
        env.logger.info(f"cleaning dotfiles from {args.target} (⌟‾╥ ‾╥)°")
        if not stow_dotfiles(Path(args.comp), Path(args.target), env, clean=True):
            env.logger.error("couldn’t clean dotfiles... sobs (ɐ•́ﾟ•̀ɐ)")
            sys.exit(1)

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


def initialize_environment(verbose: bool) -> Environment:
    logger = setup_logger(verbose)
    platform = get_platform()
    system_name = get_system_name()
    context = {
        **merge_dicts(
            load_context(platform, system_name, Config.contexts, logger),
            {"platform": platform, "system_name": system_name},
        )
    }
    return Environment(platform, system_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"copied {source} to {destination} ₰˜⋉♡")

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

    for template in templates:
        env.logger.debug(f"submitting template render for {template} ₰˜൨൨")
        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 {template} ✧*ฺ")
    jinja_env = jinja2.Environment(
        loader=jinja2.BaseLoader,
        undefined=jinja2.StrictUndefined,
        trim_blocks=True,
        lstrip_blocks=True,
    )
    rendered = render(template, jinja_env, env)
    with open(template, "w") as t:
        t.write(rendered)
    env.logger.debug(f"removing .j2 suffix from {template.absolute()} ✧˖°")
    template.rename(template.with_suffix(""))
    return True


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

    with ThreadPoolExecutor(max_workers=Config.max_workers) as executor:
        dotfile_dirs = list_dotfiles(source_dir)
        env.logger.debug(f"found dotfile dirs: {dotfile_dirs}")

        futures = [
            executor.submit(
                copy_with_templates_rendered,
                executor,
                d,
                target_dir / d.name,
                env,
            )
            for d in dotfile_dirs
        ]
        env.logger.info(f"submitted {len(futures)} tasks to executor ₰˜.༄")

        return sum(1 for f in as_completed(futures) if f.result()) == len(futures)


def render(source: Path, jinja_env: jinja2.Environment, env: Environment) -> Optional[str]:
    try:
        with open(source, "r") as f:
            content = f.read()
        env.logger.debug(f"reading template {source} ✿.｡.:･")
        template = jinja_env.from_string(content)
        env.logger.debug(f"rendered template from {source} ~ nyaaa :3")
        return template.render(**env.context)
    except Exception as e:
        env.logger.error(f"couldn’t render {source}: {e} (；⌓̀_⌓́)")
        return None


def stow_dotfiles(source_dir: Path, target_dir: Path, env: Environment, clean: bool = False) -> bool:
    if not stow_installed():
        env.logger.error("stow not installed (╥ﾟ╥)")
        return False

    packages = list_dotfiles(source_dir)
    stow_cmd = ["-D"] if clean else ["--no-folding"]

    for pkg in packages:
        env.logger.info(f"running stow {stow_cmd} for {pkg.name} ₰˜݆༿")
        run_shell_command(["stow", "-d", source_dir, "-t", target_dir, *stow_cmd, 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_system_name() -> str:
    os.environ["PLATFORM"] = get_platform()
    return run_shell_command([str(Config.script_dir / "system_name.sh")])

def list_dotfiles(p: Path) -> list[Path]:
    denylist = [".", "__"]
    return [d for d in p.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 merge_dicts(defaults, system_config)
    except (FileNotFoundError, json.JSONDecodeError) as e:
        logger.error(f"error loading context: {e} ⋆ฺ°☁｡⋆ฺ °★ °｡")
        return {}


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="compile dotfiles ✨")
    parser.add_argument("--stow", action="store_true", help="stow compiled dotfiles ♡")
    parser.add_argument("--clean", action="store_true", help="clean stowed dotfiles ˚°")
    parser.add_argument("--source", default=Config.default_source_dir, help="source dir for dotfiles")
    parser.add_argument("--comp", default=Config.compiled_dir, help="compiled template output dir")
    parser.add_argument("--target", default=Config.default_target_dir, help="stow target directory")
    parser.add_argument("--verbose", "-v", action="store_true", help="enable verbose logging ✧*ฺ", default=False)
    return parser.parse_args()


if __name__ == "__main__":
    main()