#!/usr/bin/env python3
import tomllib
import shutil
import tarfile
import tempfile
from datetime import datetime
from pathlib import Path
from enum import Enum


CONFIG_FILE = Path(__file__).parent / "config.toml"

def create_tool_bundle(bundle_name, entries):
    """
    Create a tar.gz bundle for a single tool.
    `entries` should be a list of entries for that tool only.
    """
    tempdir = Path(tempfile.mkdtemp(prefix="dotfile_manager"))

    bundle_dir = Path(f"{tempdir}/dotfiles-{bundle_name}")
    bundle_dir.mkdir()

    # Create extract.sh
    extract_script = bundle_dir / "extract.sh"
    with extract_script.open("w") as f:
        f.write("#!/bin/sh\n")
        f.write("set -e\n\n")
        f.write("# Extract dotfiles for {}\n".format(bundle_name))

        for entry in entries:
            src = Path(entry["repo_path"]).name
            dest = entry["system_path"].replace("$HOME", str(Path.home()))
            backup = f"{dest}.old"

            f.write(f"echo 'Installing {entry['repo_path']} -> {dest}'\n")
            f.write(f"if [ -e '{dest}' ]; then\n")
            f.write(f"  mv '{dest}' '{backup}'\n")
            f.write("fi\n")
            if entry["type"] == "dir":
                f.write(f"cp -r '{src}' '{dest}'\n")
            else:
                f.write(f"mkdir -p $(dirname '{dest}')\n")
                f.write(f"cp '{src}' '{dest}'\n")
            f.write("\n")

    extract_script.chmod(0o755)

    # Copy files into bundle_dir
    for entry in entries:
        src_path = Path(entry["repo_path"])
        dest_path = bundle_dir / src_path.name
        if src_path.is_dir():
            shutil.copytree(src_path, dest_path)
        else:
            dest_path.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(src_path, dest_path)

    # Create tar.gz
    tar_path = Path(f"dotfiles-{bundle_name}.tar.gz")
    with tarfile.open(tar_path, "w:gz") as tar:
        tar.add(bundle_dir, arcname=bundle_dir.name)

    # Cleanup temp directory
    shutil.rmtree(bundle_dir)

    print(f"Bundle created: {tar_path}")


class ToolType(str, Enum):
    FILE = "file"
    DIR = "dir"


def parse_dotfiles_toml(path: Path):
    """
    Parse a dotfiles TOML and return a flat list of entries.
    Rules:
      - Every entry (top-level or nested) must have a 'type' (file or dir)
      - Nested entries belong to the same tool
    Returns a list of dicts with keys: tool, repo_path, system_path, type, ...
    """
    with open(path, "rb") as f:
        cfg = tomllib.load(f)

    all_entries = []
    tools = cfg.get("tools", {})

    def process_table(tool_name, table):
        """
        Process a tool table: flatten nested entries, enforce 'type'.
        """
        # If table itself has 'type', treat it as a single entry
        if "type" in table:
            if table["type"] not in ("file", "dir"):
                raise ValueError(f"Invalid type for tool '{tool_name}': {table['type']}")
            entry = dict(table)
            entry["tool"] = tool_name
            all_entries.append(entry)
            return  # stop recursion here

        # Otherwise recurse into nested tables
        for key, value in table.items():
            if isinstance(value, dict):
                process_table(tool_name, value)
            else:
                raise ValueError(f"Invalid entry in tool '{tool_name}': {key}={value}")

    for tool_name, tool_data in tools.items():
        if not isinstance(tool_data, dict):
            raise ValueError(f"Tool '{tool_name}' must be a table")
        process_table(tool_name, tool_data)

    return all_entries


def filter_entries(entries, targets=None):
    if not targets:
        return entries

    available_tools = {e["tool"] for e in entries}
    unknown = [t for t in targets if t not in available_tools]
    if unknown:
        raise ValueError(f"Unknown tool(s) specified: {', '.join(unknown)}")

    return [e for e in entries if e["tool"] in targets]


def install(entries):
    """
    Install all entries from repo_path → system_path.
    If system_path exists, move it to .old before overwriting.
    """
    for entry in entries:
        repo_path = Path(entry["repo_path"]).expanduser()
        system_path = Path(entry["system_path"]).expanduser()
        backup_path = system_path.with_name(system_path.name + ".old")

        # Backup existing file/folder
        if system_path.exists():
            if backup_path.exists():
                shutil.rmtree(backup_path) if backup_path.is_dir() else backup_path.unlink()
            system_path.rename(backup_path)

        # Copy from repo to system
        if entry["type"] == "dir":
            shutil.copytree(repo_path, system_path)
        else:  # file
            system_path.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(repo_path, system_path)


def store(entries):
    """
    Store all entries: system_path → repo_path (strict sync).
    - For directories: remove repo folder first
    - For files: overwrite
    """
    for entry in entries:
        repo_path = Path(entry["repo_path"]).expanduser()
        system_path = Path(entry["system_path"]).expanduser()

        if entry["type"] == "dir":
            if repo_path.exists():
                shutil.rmtree(repo_path)
            shutil.copytree(system_path, repo_path)
        else:  # file
            repo_path.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy2(system_path, repo_path)


def restore(entries):
    """
    Restore all entries from .old → system_path.
    """
    for entry in entries:
        system_path = Path(entry["system_path"]).expanduser()
        backup_path = system_path.with_name(system_path.name + ".old")

        if backup_path.exists():
            if system_path.exists():
                if system_path.is_dir():
                    shutil.rmtree(system_path)
                else:
                    system_path.unlink()
            backup_path.rename(system_path)


def clean(entries):
    """
    Remove all .old backups created during install.
    """
    for entry in entries:
        system_path = Path(entry["system_path"]).expanduser()
        backup_path = system_path.with_name(system_path.name + ".old")

        if backup_path.exists():
            if backup_path.is_dir():
                shutil.rmtree(backup_path)
            else:
                backup_path.unlink()


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Manage dotfiles from TOML config")
    parser.add_argument(
        "action",
        choices=["install", "store", "restore", "clean", "bundle"],
        help="Action to perform"
    )
    parser.add_argument(
        "targets",
        nargs="*",
        help="Tools to act on (default: all)"
    )

    args = parser.parse_args()

    entries = parse_dotfiles_toml(CONFIG_FILE)
    entries = filter_entries(entries, args.targets)

    if args.action == "install":
        install(entries)
    elif args.action == "store":
        store(entries)
    elif args.action == "restore":
        restore(entries)
    elif args.action == "clean":
        clean(entries)
    elif args.action == "bundle":
        tools = "-".join(sorted(list({entry['tool'] for entry in entries})))
        bundle_name = f"{datetime.now().strftime('%Y-%m-%d-%H-%M')}-{tools}"
        create_tool_bundle(bundle_name, entries)