#!/usr/bin/python3

# Runs Quick Setup driver manifests. When invoked via pkexec the polkit action
# io.stillhq.quicksetup.run-driver-helper allows this without admin auth, so
# this helper must never load driver files from anywhere except the trusted
# system directory. That directory lives on the bootc /usr tree, which is
# read-only unless the user has already obtained root, so trusting its contents
# is equivalent to trusting root.

import argparse
import json
import os
import subprocess
import sys

SYSTEM_DRIVERS_DIR = "/usr/share/quick-setup/drivers"
DEV_DRIVERS_DIR_ENV = "QUICK_SETUP_DEV_DRIVERS_DIR"

SAFE_PATH = "/usr/sbin:/usr/bin:/sbin:/bin"


def parse_args():
    parser = argparse.ArgumentParser(description="Run quick-setup driver manifests")
    parser.add_argument(
        "--phase",
        required=True,
        choices=["pre_wifi", "post_wifi"],
    )
    return parser.parse_args()


def resolve_drivers_dir():
    if os.geteuid() == 0:
        return SYSTEM_DRIVERS_DIR

    override = os.environ.get(DEV_DRIVERS_DIR_ENV)
    if override:
        return override

    return SYSTEM_DRIVERS_DIR


def normalize_script(value):
    if isinstance(value, str):
        return value

    if isinstance(value, list):
        lines = []
        for item in value:
            if not isinstance(item, str):
                raise ValueError("script arrays must only contain strings")
            lines.append(item)
        return "\n".join(lines)

    raise ValueError("script values must be a string or an array of strings")


def load_driver_files(drivers_dir):
    if not os.path.isdir(drivers_dir):
        return []

    entries = []
    for entry in sorted(os.listdir(drivers_dir)):
        if not entry.endswith(".json"):
            continue

        path = os.path.join(drivers_dir, entry)
        try:
            with open(path, "r", encoding="utf-8") as handle:
                manifest = json.load(handle)
        except Exception as exc:
            entries.append(
                {
                    "id": entry,
                    "name": entry,
                    "phase": None,
                    "requires_reboot": False,
                    "error": f"Failed to load manifest: {exc}",
                }
            )
            continue

        entries.append(
            {
                "id": manifest.get("id", entry[:-5]),
                "name": manifest.get("name", manifest.get("id", entry[:-5])),
                "phase": manifest.get("phase"),
                "requires_reboot": bool(manifest.get("requires_reboot", False)),
                "check_script": normalize_script(manifest.get("check_script", "")),
                "install_script": normalize_script(manifest.get("install_script", "")),
                "success_script": normalize_script(manifest.get("success_script", "")),
                "error": None,
            }
        )

    return entries


def run_script(script):
    env = {
        "PATH": SAFE_PATH,
        "LANG": "C.UTF-8",
        "LC_ALL": "C.UTF-8",
        "HOME": os.environ.get("HOME", "/root"),
    }
    completed = subprocess.run(
        ["/usr/bin/bash", "-c", script],
        text=True,
        capture_output=True,
        env=env,
    )
    return completed.returncode, completed.stdout.strip(), completed.stderr.strip()


def format_details(stdout_text, stderr_text):
    # Driver scripts surface their error message via stdout/stderr; we present
    # whichever is non-empty (stderr first, since shell errors land there).
    chunks = []
    if stderr_text:
        chunks.append(stderr_text)
    if stdout_text:
        chunks.append(stdout_text)
    return "\n\n".join(chunks)


def process_manifest(manifest, requested_phase):
    if manifest["error"] is not None:
        return {
            "driver_id": manifest["id"],
            "name": manifest["name"],
            "status": "failed",
            "message": "Quick Setup could not read this driver definition.",
            "details": manifest["error"],
            "requires_reboot": False,
        }

    if manifest["phase"] != requested_phase:
        return None

    for field_name in ("check_script", "install_script", "success_script"):
        if not manifest.get(field_name):
            return {
                "driver_id": manifest["id"],
                "name": manifest["name"],
                "status": "failed",
                "message": f"Missing {field_name}.",
                "details": "",
                "requires_reboot": False,
            }

    check_code, check_stdout, check_stderr = run_script(manifest["check_script"])
    if check_code == 1:
        return {
            "driver_id": manifest["id"],
            "name": manifest["name"],
            "status": "not_needed",
            "message": "Not needed on this hardware.",
            "details": format_details(check_stdout, check_stderr),
            "requires_reboot": False,
        }

    if check_code != 0:
        return {
            "driver_id": manifest["id"],
            "name": manifest["name"],
            "status": "failed",
            "message": "The detection script returned an unexpected status.",
            "details": format_details(check_stdout, check_stderr),
            "requires_reboot": False,
        }

    install_code, install_stdout, install_stderr = run_script(manifest["install_script"])
    if install_code != 0:
        return {
            "driver_id": manifest["id"],
            "name": manifest["name"],
            "status": "failed",
            "message": "Installation failed.",
            "details": format_details(install_stdout, install_stderr),
            "requires_reboot": False,
        }

    success_code, success_stdout, success_stderr = run_script(manifest["success_script"])
    if success_code != 0:
        return {
            "driver_id": manifest["id"],
            "name": manifest["name"],
            "status": "failed",
            "message": "The driver did not verify after installation.",
            "details": format_details(success_stdout, success_stderr),
            "requires_reboot": False,
        }

    message = "Installed successfully."
    if manifest["requires_reboot"]:
        message = "Installed successfully. Changes will apply on reboot."

    return {
        "driver_id": manifest["id"],
        "name": manifest["name"],
        "status": "success",
        "message": message,
        "details": format_details(success_stdout, success_stderr),
        "requires_reboot": manifest["requires_reboot"],
    }


def main():
    args = parse_args()

    drivers_dir = resolve_drivers_dir()
    manifests = load_driver_files(drivers_dir)
    results = []

    for manifest in manifests:
        result = process_manifest(manifest, args.phase)
        if result is not None:
            results.append(result)

    payload = {
        "phase": args.phase,
        "drivers_dir": drivers_dir,
        "results": results,
    }
    json.dump(payload, sys.stdout)
    sys.stdout.write("\n")
    return 0


if __name__ == "__main__":
    sys.exit(main())
