#!/usr/bin/env python3

import argparse
import subprocess
import json
from pathlib import Path
import shlex
import shutil

DOTFILES_DIR = Path(__file__).parent
SETUPS_DIR = DOTFILES_DIR / "setups"
CONFIG_DIR = DOTFILES_DIR / "config"
CONFIG_PATH = DOTFILES_DIR / "config.json"

def load_config():
    if not CONFIG_PATH.exists():
        raise FileNotFoundError(f"Configuration file not found: {CONFIG_PATH}")

    with open(CONFIG_PATH, "r") as f:
        data = json.load(f)

        if not isinstance(data, dict):
            raise ValueError("JSON must be an object")

        if "environments" not in data:
            raise ValueError("Missing required field: 'environments'")

        if not isinstance(data["environments"], dict):
            raise ValueError("'environments' must be an object")

        if "template" in data and not isinstance(data["template"], dict):
            raise ValueError("'template' must be an object if present")

        return data

def get_environment_packages(config, env, search_package=None):
    env_entries = config["environments"].get(env, [])
    if not env_entries:
        raise TypeError(f"Environment {env} was not found or it is empty")

    template_config = config.get("template", {})
    packages = []
    for entry in env_entries:
        if isinstance(entry, str):
            entry = { "package": entry }

        package_name = entry.pop("package", None)
        if package_name is None:
            raise TypeError(f"The following entry is missing `package` field: {entry}")

        package = { "name": package_name }
        if package_name in template_config and not entry.get("ignore-template", False):
            template_entry = template_config[package_name]
            package = { **package, **template_entry, **entry }
        else:
            package.update(entry)

        package.pop("ignore-template", None)

        if "link" in package:
            link_from = package["link"].get("from")
            link_to = package["link"].get("to")
            if not isinstance(link_from, str) or not isinstance(link_to, str):
                raise ValueError("`link` should follow the structure: `{ from: str, to: str }`")

            if len(link_from.split("/")) != 2:
                raise ValueError("`link.from` should be '<env>/<package>'")

            package["link"] = {
                "from": Path(CONFIG_DIR / link_from).expanduser(),
                "to": Path(link_to).expanduser()
            }

        if search_package == None:
            packages.append(package)
        else:
            if package["name"] == search_package:
                packages.append(package)
                break

    return packages

def force_delete(path):
    if path.is_file() or path.is_symlink():
        path.unlink()
    elif path.is_dir():
        shutil.rmtree(path)

def link_environment(config, env, **kwargs):
    options = {
        "package": kwargs.get("package"),
        "copy": kwargs.get("copy", False),
        "force": kwargs.get("force", False)
    }

    packages = get_environment_packages(config, env, search_package=options["package"])
    for package in packages:
        print(f"[{package['name']}]")

        if "link" not in package:
            print("\t> Skipped: No link entry")
            continue

        src = package["link"]["from"]
        dest = package["link"]["to"]

        if dest.exists():
            if options["force"]:
                force_delete(dest)
                print(f"\t> Deleted: {dest}")
            else:
                print(f"\t> Skipped: Already exists {dest}")
                continue

        if options["copy"]:
            dest.parent.mkdir(parents=True, exist_ok=True)
            shutil.copy(src, dest)
            print(f"\t> Copied: {src} -> {dest}")
        else:
            dest.symlink_to(src)
            print(f"\t> Symlinked: {src} -> {dest}")

        if "post-link" in package:
            command = package["post-link"]
            subprocess.run(shlex.split(command), check=True)
            print(f"\t> Post-link executed: `{command}`")

def install_environment(config, env, **kwargs):
    options = {
        "package": kwargs.get("package"),
    }

    packages = get_environment_packages(config, env, search_package=options["package"])
    for package in packages:
        print(f"[{package['name']}]")

        if "install" not in package:
            print("\t> Skipped: No install entry")
            continue

        install_command = package.get("install")
        if install_command:
            subprocess.run(shlex.split(install_command), check=True)
            print(f"\t> Installed: `{install_command}`")

        if "post-install" in package:
            postinstall_command = package["post-install"]
            subprocess.run(shlex.split(postinstall_command), check=True)
            print(f"\t> Post-install executed: `{postinstall_command}`")

def setup_environment(config, env, **kwargs):
    print(f"[{env}]")
    options = {
        "extra_args": kwargs.get("extra"),
    }

    setup_script = SETUPS_DIR / f"{env}.sh"
    if setup_script.exists():
        cmd = ["bash", str(setup_script)]
        if options["extra_args"]:
            cmd.extend(shlex.split(options["extra_args"]))  # Split extra args safely
        subprocess.run(cmd, check=True)
        print(f"\t> Setup script executed: {setup_script} {options['extra_args'] or ''}")
    else:
        print(f"\t> No setup script found for {env}")

def main():
    config = load_config()

    config_envs = list(config["environments"].keys())
    setup_envs = [script.stem for script in SETUPS_DIR.glob("*.sh")]

    parser = argparse.ArgumentParser(description="Dotfile & System Setup Manager")
    subparsers = parser.add_subparsers(dest="command", required=True)

    subparser = subparsers.add_parser("link", help="Link configs")
    subparser.add_argument("env", choices=config_envs)
    subparser.add_argument("-p", "--package")
    subparser.add_argument("-f", "--force", action="store_true")
    subparser.add_argument("--copy", action="store_true")

    subparser = subparsers.add_parser("install", help="Install packages")
    subparser.add_argument("env", choices=config_envs)
    subparser.add_argument("-p", "--package")

    setup_parser = subparsers.add_parser("setup", help="Run setup script")
    setup_parser.add_argument("env", choices=setup_envs)
    setup_parser.add_argument("--extra")

    args = parser.parse_args()
    command_actions = {"link": link_environment, "install": install_environment, "setup": setup_environment}
    command_actions[args.command](config, **vars(args))

if __name__ == "__main__":
    main()