#!/usr/bin/python3

import json
import os
import sys
from pathlib import Path

import tomlkit

BASE_REPO_DIR: Path = Path("/var/lib/slpkg/repos")
REPOSITORIES_TOML_PATH: Path = Path("/etc/slpkg/repositories.toml")

# ---------------------------------------------------------------------------
# Flags available for each command (and its aliases).
# Build from shared flag groups to reduce repetition.
# ---------------------------------------------------------------------------

_DBG = ["-d", "--debug"]
_NLK = ["-k", "--no-lock"]
_COL = ["-C", "--color"]
_REP = ["-o", "--repository"]
_YES = ["-y", "--yes"]
_QUI = ["-q", "--quiet"]
_CHK = ["-c", "--check", "--dry-run"]
_JSN = ["-j", "--json"]
_PAR = ["-P", "--parallel"]
_DLG = ["-D", "--dialog"]
_SEL = ["-S", "--select"]
_NCI = ["-i", "--ignore-case"]
_PKV = ["-p", "--pkg-version"]
_LIN = ["-n", "--lines"]
_NRS = ["-N", "--no-resolve"]
_SKI = ["-I", "--skip-installed"]
_TAG = ["-T", "--repo-tag"]
_EDT = ["-e", "--edit"]  # edit_sbo_parser (takes NAME) and edit_parser (store_true)

# Canonical-command → flag list (aliases share the same list via expansion below)
COMMAND_FLAGS: dict[str, list[str]] = {
    "update": _DBG + _NLK + _COL + _REP + _CHK,
    "upgrade": (
        _DBG + _NLK + _COL + _REP + _CHK + _YES + _NRS + _QUI + _PAR
        + ["-r", "--reinstall"] + ["-n", "--install-new"] + ["-T", "--ignore-tag"]
        + _NCI + _SEL + _DLG + ["--repos-priority"]
    ),
    "config": _DBG + _NLK + _DLG + _EDT + ["--validate"],
    "repo-info": _DBG + _NLK + _COL + _REP + _JSN + ["-b", "--benchmark"],
    "repo": (
        _DBG + _NLK + _COL + _QUI + _JSN
        + ["--enable", "--disable", "--mirror", "--default", "--priority", "--clear-priority"]
    ),
    "stats": _DBG + _NLK + _COL + _JSN + ["-t", "--top"],
    "rollback": _DBG + _NLK + _COL + _YES,
    "clean": _DBG + _NLK,
    "clean-system": _DBG + _NLK + _COL + _YES + _QUI + ["-A", "--autoremove"],
    "health-check": _DBG + _NLK + _COL + _QUI,
    "audit": _DBG + _NLK + _COL,
    "self-check": _DBG + _NLK,
    "verify": _DBG + _NLK + _COL + _YES + _QUI + ["-f", "--fix"],
    "build": (
        _DBG + _NLK + _COL + _REP + _YES + _NRS + _SKI + _QUI + _PAR
        + _NCI + _SEL + _DLG + _EDT + _TAG
    ),
    "reinstall": (
        _DBG + _NLK + _COL + _REP + _YES + _NRS + _QUI + _PAR
        + _NCI + _SEL + _DLG + _EDT + _TAG
    ),
    "install": (
        _DBG + _NLK + _COL + _REP + _YES + _CHK + ["-r", "--reinstall"] + _NRS
        + _SKI + _QUI + _PAR + _NCI + _SEL + _DLG + _EDT + _TAG + ["--repos-priority"]
    ),
    "rebuild": _DBG + _NLK + _COL + _REP + _YES + _CHK + _QUI + _PAR + _SEL + _DLG + _TAG,
    "install-pkg": (
        _DBG + _NLK + _COL + _YES + ["-r", "--reinstall"] + _QUI + ["-X", "--extra-args"]
    ),
    "remove": _DBG + _NLK + _COL + _YES + _CHK + _NRS + _SEL + _DLG + _QUI,
    "download": _DBG + _NLK + _COL + _REP + _YES + _NCI + _SEL + _DLG + ["-w", "--output-dir"],
    "export": _DBG + _NLK + _COL + ["-n", "--no-deps"],
    "import": (
        _DBG + _NLK + _COL + _YES + ["-r", "--reinstall"] + _NRS + _SKI
        + _QUI + _PAR + _EDT + _TAG
    ),
    "list": _DBG + _NLK + _COL + _NCI + ["-t", "--description"] + _SEL + _DLG + _JSN,
    "list-files": _DBG + _NLK + _COL + _NCI + _QUI,
    "info": _DBG + _NLK + _COL + _REP + _NCI + _SEL + _DLG + _JSN,
    "search": (
        _DBG + _NLK + _COL + _REP + _NCI + _PKV + ["-t", "--description"]
        + _SEL + _DLG + _QUI + _JSN
    ),
    "file-search": _DBG + _NLK + _COL + _REP + _NCI + _QUI + ["-L", "--local"],
    "reverse": (
        _DBG + _NLK + _COL + _REP + _NCI + _QUI + _PKV + ["-F", "--full-reverse"] + _SEL + _DLG
    ),
    "depends": _DBG + _NLK + _COL + _REP + _NCI + _QUI + _PKV + _SEL + _DLG,
    "changelog": (
        _DBG + _NLK + _COL + _REP + _EDT + _LIN + ["-s", "--short"]
        + ["--date"] + ["-g", "--pager"] + ["-Q", "--query"]
    ),
    "history": _DBG + _NLK + _COL + _LIN + _JSN + ["-x", "--clear", "--package"],
    "orphans": _DBG + _NLK + _COL + _YES + _QUI + ["-A", "--autoremove"],
    "why": _DBG + _NLK + _COL + _QUI + ["-t", "--transitive"] + _PKV,
    "mark": _DBG + _NLK + _COL + _QUI + ["-u", "--unmark"],
    "pin": _DBG + _NLK + _COL + _QUI + ["-u", "--unpin"],
    "check-broken": _DBG + _NLK + _COL + _QUI,
    "blacklist": _DBG + _NLK + _COL + _QUI + ["-r", "--remove"],
    "version": _DBG + _NLK + [],
    "help": _DBG + _NLK + [],
}

# Aliases → canonical command name
_ALIASES: dict[str, str] = {
    "u": "update", "upd": "update",
    "U": "upgrade", "upg": "upgrade",
    "C": "config", "conf": "config",
    "ri": "repo-info", "rinf": "repo-info",
    "rp": "repo", "rep": "repo",
    "st": "stats", "stt": "stats",
    "rb": "rollback", "und": "rollback",
    "c": "clean", "cln": "clean",
    "cs": "clean-system", "csy": "clean-system",
    "hc": "health-check", "hchk": "health-check",
    "ad": "audit", "adt": "audit",
    "sc": "self-check", "schk": "self-check",
    "V": "verify", "vfy": "verify",
    "b": "build", "bld": "build",
    "re": "reinstall", "rein": "reinstall",
    "i": "install", "inst": "install",
    "rbd": "rebuild", "rbld": "rebuild",
    "ip": "install-pkg", "ipkg": "install-pkg",
    "R": "remove", "rmv": "remove",
    "dl": "download", "dwn": "download",
    "ex": "export", "exp": "export",
    "im": "import", "imp": "import",
    "l": "list", "lst": "list",
    "lf": "list-files", "lfl": "list-files",
    "inf": "info", "show": "info",
    "s": "search", "ser": "search",
    "f": "file-search", "fs": "file-search",
    "r": "reverse", "rdeps": "reverse",
    "d": "depends", "dep": "depends",
    "n": "changelog", "chg": "changelog",
    "hi": "history", "hst": "history",
    "or": "orphans", "orp": "orphans",
    "w": "why", "wy": "why",
    "mk": "mark", "mrk": "mark",
    "pn": "pin", "pnh": "pin",
    "cb": "check-broken", "cbr": "check-broken",
    "bl": "blacklist", "blk": "blacklist",
    "v": "version", "ver": "version",
    "h": "help", "hlp": "help",
}

# Flags that consume the next word as a value.
# After seeing one of these as prev_word, the current word is a value, not a flag.
VALUE_FLAGS: set[str] = {
    "-o", "--repository",
    "-C", "--color",
    "-n", "--lines",
    "-T", "--repo-tag",
    "-e", "--edit",        # edit_sbo_parser version takes NAME
    "-Q", "--query",
    "--date",
    "-w", "--output-dir",
    "-t", "--top",         # stats: takes N
    "--package",           # history: takes NAME
    "-X", "--extra-args",
    "--repos-priority",    # nargs="+": first value is a repo name
    "--enable",
    "--disable",
    "--mirror",
    "--default",
    "--priority",
}

# Flags whose value should be a repository name
REPO_VALUE_FLAGS: set[str] = {
    "-o", "--repository",
    "--enable", "--disable",
    "--default", "--priority",
    "--repos-priority",
    "--mirror",   # first arg is repo name (second is URL, but we can't distinguish)
}

# Flags whose value should be an installed package name
PKG_VALUE_FLAGS: set[str] = {
    "--package",
}

# Flags whose value should be a directory path
DIR_VALUE_FLAGS: set[str] = {
    "-w", "--output-dir",
}

# Color choices
COLOR_CHOICES: list[str] = ["on", "off", "ON", "OFF"]


def get_enabled_repositories(only_default: bool = False) -> list[str]:
    """Reads repositories.toml and returns enabled repository names.
    If only_default is True, returns only the repository marked as DEFAULT.
    """
    if not REPOSITORIES_TOML_PATH.is_file():
        return []

    try:
        content = REPOSITORIES_TOML_PATH.read_text(encoding="utf-8")
        data = tomlkit.parse(content)
        enabled = []

        # Get default repository from [DEFAULT] REPOSITORY = "..."
        default_repo = None
        if "DEFAULT" in data and isinstance(data["DEFAULT"], dict):
            default_repo = data["DEFAULT"].get("REPOSITORY", "").lower()

        for key, value in data.items():
            if key == "DEFAULT":
                continue
            if isinstance(value, dict) and value.get("ENABLE") is True:
                repo_name = key.lower()
                enabled.append(repo_name)

        if only_default and default_repo:
            return [default_repo]

        return enabled
    except Exception:  # pylint: disable=[W0718]
        return []


def get_all_repositories() -> list[str]:
    """Returns all repository names (enabled or not) from repositories.toml."""
    if not REPOSITORIES_TOML_PATH.is_file():
        return []
    try:
        content = REPOSITORIES_TOML_PATH.read_text(encoding="utf-8")
        data = tomlkit.parse(content)
        return [key.lower() for key, value in data.items()
                if key != "DEFAULT" and isinstance(value, dict)]
    except Exception:  # pylint: disable=[W0718]
        return []


def get_all_packages(repository: str = "*") -> list[str]:
    """Returns a list of all package names from the specified repository data.json files."""
    packages: set[str] = set()
    repos_to_check = []

    if repository == "*" or not repository:
        # Try to get only the default repository first
        repos_to_check = get_enabled_repositories(only_default=True)
        # If no repository is marked as DEFAULT, fall back to all enabled
        if not repos_to_check:
            repos_to_check = get_enabled_repositories()
    else:
        repos_to_check = [repository.lower()]

    for repo in repos_to_check:
        json_path = BASE_REPO_DIR / repo / "data.json"
        if json_path.is_file():
            try:
                with json_path.open("r", encoding="utf-8") as f:
                    data = json.load(f)
                    packages.update(data.keys())
            except Exception:  # pylint: disable=[W0718]
                continue

    return sorted(list(packages))


def get_installed_packages(prefix: str = "") -> list[str]:
    """Returns a list of installed package names matching the prefix."""
    pkg_dir = Path("/var/log/packages")
    if not pkg_dir.is_dir():
        return []

    pkgs = []
    for pkg_path in pkg_dir.iterdir():
        name = pkg_path.name
        # Simple heuristic to extract the package name from full Slackware version string
        # Match name-version-arch-build
        parts = name.rsplit("-", 3)
        if len(parts) >= 1:
            pkg_name = parts[0]
            if not prefix or pkg_name.startswith(prefix):
                pkgs.append(pkg_name)
    return sorted(list(set(pkgs)))


def get_blacklisted_packages() -> list[str]:
    """Reads blacklist.toml and returns currently blacklisted packages."""
    path = Path("/etc/slpkg/blacklist.toml")
    if not path.is_file():
        return []
    try:
        content = path.read_text(encoding="utf-8")
        data = tomlkit.parse(content)
        return data.get("PACKAGES", [])  # type: ignore
    except Exception:  # pylint: disable=[W0718]
        return []


def get_file_completions(prefix: str, command: str = "") -> list[str]:
    """Returns file and directory completions for a given prefix."""
    # Handle home directory expansion
    expanded_prefix = os.path.expanduser(prefix)

    if prefix.endswith("/"):
        directory = expanded_prefix
        base = ""
    else:
        directory = os.path.dirname(expanded_prefix) or "."
        base = os.path.basename(expanded_prefix)

    completions = []
    try:  # pylint: disable=[R1702]
        for entry in os.listdir(directory):
            if entry.startswith(base):
                full_path = os.path.join(directory, entry)

                # Determine the display path
                if prefix.startswith("/") or prefix.startswith("~") or "/" in prefix:
                    res_path = os.path.join(directory, entry)
                else:
                    res_path = entry

                if os.path.isdir(full_path):
                    completions.append(res_path + "/")
                else:
                    # Filter for common package extensions only for install-pkg
                    if command in ["install-pkg", "ip", "ipkg"]:
                        if any(entry.endswith(ext) for ext in [".txz", ".tgz", ".tlz", ".tbz"]):
                            completions.append(res_path)
                    else:
                        # For export/import/etc., show all files
                        completions.append(res_path)
    except Exception:  # pylint: disable=[W0718]
        pass
    return sorted(list(set(completions)))


def get_dir_completions(prefix: str) -> list[str]:
    """Returns only directory completions for a given prefix."""
    expanded_prefix = os.path.expanduser(prefix)
    if prefix.endswith("/"):
        directory = expanded_prefix
        base = ""
    else:
        directory = os.path.dirname(expanded_prefix) or "."
        base = os.path.basename(expanded_prefix)

    completions = []
    try:
        for entry in os.listdir(directory):
            if entry.startswith(base):
                full_path = os.path.join(directory, entry)
                if os.path.isdir(full_path):
                    if prefix.startswith("/") or prefix.startswith("~") or "/" in prefix:
                        completions.append(os.path.join(directory, entry) + "/")
                    else:
                        completions.append(entry + "/")
    except Exception:  # pylint: disable=[W0718]
        pass
    return sorted(completions)


def _resolve_command(args: list[str]) -> str:
    """Return the canonical command name found in args, or empty string."""
    for word in args[1:]:  # skip 'slpkg' itself
        if not word.startswith("-"):
            if word in COMMAND_FLAGS:
                return word
            if word in _ALIASES:
                return _ALIASES[word]
    return ""


def main() -> None:  # pylint: disable=[R0912,R0914,R0915]
    """Main function for generating shell completions."""
    try:
        # The wrapper scripts pass the current word index as the first argument,
        # followed by the array of words in the command line.
        if len(sys.argv) < 3:
            sys.exit(0)

        # Index of the current word being completed (0-based)
        cword_idx = int(sys.argv[1])
        # Array of words: ['slpkg', 'command', 'arg1', ...]
        words = sys.argv[2:]

        # Current word being typed (the prefix for completions)
        current_word = words[cword_idx] if cword_idx < len(words) else ""
        # The word immediately before the one being typed
        prev_word = words[cword_idx - 1] if cword_idx > 0 else ""
        # Index relative to 'slpkg' (0 = the word after 'slpkg', 1 = next arg, etc.)
        cword_index = cword_idx - 1

        # Use 'words' instead of 'args' for searching commands in the command line
        args = words

        all_slpkg_commands: list[str] = [
            "update", "u", "upd",
            "upgrade", "U", "upg",
            "config", "C", "conf",
            "repo-info", "ri", "rinf",
            "repo", "rp", "rep",
            "stats", "st", "stt",
            "rollback", "rb", "und",
            "clean", "c", "cln",
            "self-check", "sc", "schk",
            "build", "b", "bld",
            "install", "i", "inst",
            "reinstall", "re", "rein",
            "rebuild", "rbd", "rbld",
            "install-pkg", "ip", "ipkg",
            "remove", "R", "rmv",
            "download", "dl", "dwn",
            "export", "ex", "exp",
            "import", "im", "imp",
            "list", "l", "lst",
            "list-files", "lf", "lfl",
            "info", "inf", "show",
            "search", "s", "ser",
            "file-search", "f", "fs",
            "reverse", "r", "rdeps",
            "depends", "d", "dep",
            "why", "w", "wy",
            "orphans", "or", "orp",
            "mark", "mk", "mrk",
            "pin", "pn", "pnh",
            "check-broken", "cb", "cbr",
            "clean-system", "cs", "csy",
            "health-check", "hc", "hchk",
            "audit", "ad", "adt",
            "blacklist", "bl", "blk",
            "changelog", "n", "chg",
            "history", "hi", "hst",
            "verify", "V", "vfy",
            "version", "v", "ver",
            "help", "h", "hlp",
        ]

        # Commands that expect package names from repositories
        repo_commands: list[str] = [
            "install", "i", "inst",
            "reinstall", "re", "rein",
            "build", "b", "bld",
            "download", "dl", "dwn",
            "info", "inf", "show",
            "search", "s", "ser",
            "file-search", "f", "fs",
            "reverse", "r", "rdeps",
            "depends", "d", "dep",
        ]

        # Commands that expect currently installed package names
        installed_commands: list[str] = [
            "remove", "R", "rmv",
            "list-files", "lf", "lfl",
            "list", "l", "lst",
            "verify", "V", "vfy",
            "mark", "mk", "mrk",
            "pin", "pn", "pnh",
            "why", "w", "wy",
            "health-check", "hc", "hchk",
            "audit", "ad", "adt",
            "clean-system", "cs", "csy",
            "blacklist", "bl", "blk",
        ]

        file_commands: list[str] = [
            "install-pkg", "ip", "ipkg",
            "export", "ex", "exp",
            "import", "im", "imp",
        ]

        completions: list[str] = []
        prefix = current_word

        if cword_index == 0 or (
            not _resolve_command(args)
            and prev_word not in VALUE_FLAGS
            and not prefix.startswith("-")
        ):
            # Case: slpkg <TAB>  or  slpkg -o sbo <TAB>  (global flags before subcommand)
            completions = [c for c in all_slpkg_commands if c.startswith(prefix)]

        elif prev_word in ("-o", "--repository"):
            # Case: slpkg <CMD> -o <TAB> — always enabled repos
            repos = get_enabled_repositories()
            completions = [r for r in repos if r.startswith(prefix)]

        elif prev_word in REPO_VALUE_FLAGS - {"-o", "--repository"}:
            # --enable / --disable / --default / --priority / --repos-priority / --mirror
            # Suggest all repos (enabled or not) so user can enable/disable them
            repos = get_all_repositories()
            completions = [r for r in repos if r.startswith(prefix)]

        elif prev_word in PKG_VALUE_FLAGS:
            # --package (history): installed package name
            installed = get_installed_packages(prefix)
            completions = installed

        elif prev_word in DIR_VALUE_FLAGS:
            # -w / --output-dir: directory
            completions = get_dir_completions(prefix)

        elif prev_word in {"--color", "-C"}:
            completions = [c for c in COLOR_CHOICES if c.startswith(prefix)]

        elif prev_word in VALUE_FLAGS:
            # Other value-taking flags (--lines, --repo-tag, --query, --date, --edit,
            # --top, --extra-args): no useful completions
            completions = []

        elif prefix.startswith("-"):
            # Flag completion: find the active command and suggest its flags
            cmd = _resolve_command(args)
            if cmd:
                flags = COMMAND_FLAGS.get(cmd, [])
                completions = [f for f in flags if f.startswith(prefix)]

        elif prev_word in repo_commands or any(cmd in args for cmd in repo_commands):
            # Case: slpkg install <TAB>
            # Find if repository was specified with -o
            repo = "*"
            if "-o" in args:
                idx = args.index("-o")
                if idx + 1 < len(args):
                    repo = args[idx + 1]
            elif "--repository" in args:
                idx = args.index("--repository")
                if idx + 1 < len(args):
                    repo = args[idx + 1]

            # Get packages from the specified or default repository
            all_pkgs = get_all_packages(repo)
            completions = [p for p in all_pkgs if p.startswith(prefix)]

        elif prev_word in installed_commands or any(cmd in args for cmd in installed_commands):
            # Case: slpkg remove <TAB>
            if prev_word in ("-r", "--remove") and ("blacklist" in args or "bl" in args or "blk" in args):
                # Special case for blacklist removal
                blacklisted = get_blacklisted_packages()
                completions = [p for p in blacklisted if p.startswith(prefix)]
            else:
                installed = get_installed_packages(prefix)
                completions = installed

        elif prev_word in file_commands or any(cmd in args for cmd in file_commands):
            # Case: slpkg install-pkg <TAB>
            # Find the active command to apply proper filtering
            cmd = ""
            for c in file_commands:
                if c in args:
                    cmd = c
                    break
            completions = get_file_completions(prefix, cmd)

        elif any(c in args for c in ("help", "h", "hlp")):
            # Case: slpkg help <TAB> — suggest command names and category names
            _categories = ["repository", "packages", "query", "dependencies", "system"]
            _cmds = list(COMMAND_FLAGS.keys())
            completions = [s for s in _categories + _cmds if s.startswith(prefix)]

        for comp in completions:
            print(comp)

    except Exception:  # pylint: disable=[W0718]
        # Catch all exceptions silently, as completion scripts should not crash the shell.
        pass

    sys.exit(0)


if __name__ == "__main__":
    main()
