#!/usr/bin/python3

import importlib.machinery
import importlib.util
import json
import os
from pathlib import Path
import tempfile
import unittest
from unittest import mock


PROJECT_ROOT = Path(__file__).resolve().parents[1]
HELPER_PATH = PROJECT_ROOT / "scripts" / "quick-setup-driver-helper"

loader = importlib.machinery.SourceFileLoader("quick_setup_driver_helper", str(HELPER_PATH))
spec = importlib.util.spec_from_loader(loader.name, loader)
driver_helper = importlib.util.module_from_spec(spec)
loader.exec_module(driver_helper)


def manifest(**overrides):
    value = {
        "id": "example",
        "name": "Example Driver",
        "phase": "post_wifi",
        "requires_reboot": True,
        "check_script": "check",
        "install_script": "install",
        "success_script": "verify",
        "error": None,
    }
    value.update(overrides)
    return value


class NormalizeScriptTests(unittest.TestCase):
    def test_accepts_strings(self):
        self.assertEqual(driver_helper.normalize_script("one\ntwo"), "one\ntwo")

    def test_joins_script_arrays(self):
        self.assertEqual(driver_helper.normalize_script(["one", "two"]), "one\ntwo")

    def test_rejects_non_string_values(self):
        with self.assertRaisesRegex(ValueError, "string or an array"):
            driver_helper.normalize_script({"command": "true"})
        with self.assertRaisesRegex(ValueError, "only contain strings"):
            driver_helper.normalize_script(["true", 1])


class LoadManifestTests(unittest.TestCase):
    def test_loads_json_files_in_name_order_and_ignores_other_files(self):
        with tempfile.TemporaryDirectory() as directory:
            Path(directory, "z.json").write_text(
                json.dumps({
                    "id": "z",
                    "phase": "post_wifi",
                    "check_script": ["line 1", "line 2"],
                    "install_script": "install",
                    "success_script": "verify",
                }),
                encoding="utf-8",
            )
            Path(directory, "a.json").write_text(
                json.dumps({
                    "id": "a",
                    "phase": "pre_wifi",
                    "check_script": "check",
                    "install_script": "install",
                    "success_script": "verify",
                }),
                encoding="utf-8",
            )
            Path(directory, "README.txt").write_text("ignored", encoding="utf-8")

            loaded = driver_helper.load_driver_files(directory)

        self.assertEqual([item["id"] for item in loaded], ["a", "z"])
        self.assertEqual(loaded[1]["check_script"], "line 1\nline 2")

    def test_reports_invalid_json_without_stopping_other_manifests(self):
        with tempfile.TemporaryDirectory() as directory:
            Path(directory, "bad.json").write_text("{", encoding="utf-8")
            Path(directory, "good.json").write_text(
                json.dumps({
                    "id": "good",
                    "phase": "pre_wifi",
                    "check_script": "check",
                    "install_script": "install",
                    "success_script": "verify",
                }),
                encoding="utf-8",
            )

            loaded = driver_helper.load_driver_files(directory)

        self.assertEqual(len(loaded), 2)
        self.assertIn("Failed to load manifest", loaded[0]["error"])
        self.assertEqual(loaded[1]["id"], "good")

    def test_skips_disabled_manifests(self):
        with tempfile.TemporaryDirectory() as directory:
            Path(directory, "disabled.json").write_text(
                json.dumps({
                    "id": "disabled",
                    "enabled": False,
                    "phase": "pre_wifi",
                    "check_script": "check",
                    "install_script": "install",
                    "success_script": "verify",
                }),
                encoding="utf-8",
            )

            loaded = driver_helper.load_driver_files(directory)

        self.assertEqual(loaded, [])

    def test_missing_directory_is_an_empty_manifest_set(self):
        self.assertEqual(driver_helper.load_driver_files("/definitely/not/here"), [])


class DriverDirectoryTests(unittest.TestCase):
    @mock.patch.object(driver_helper.os, "geteuid", return_value=0)
    def test_root_always_uses_trusted_system_directory(self, _geteuid):
        with mock.patch.dict(os.environ, {driver_helper.DEV_DRIVERS_DIR_ENV: "/tmp/untrusted"}):
            self.assertEqual(
                driver_helper.resolve_drivers_dir(),
                driver_helper.SYSTEM_DRIVERS_DIR,
            )

    @mock.patch.object(driver_helper.os, "geteuid", return_value=1000)
    def test_non_root_can_use_development_override(self, _geteuid):
        with mock.patch.dict(os.environ, {driver_helper.DEV_DRIVERS_DIR_ENV: "/tmp/drivers"}):
            self.assertEqual(driver_helper.resolve_drivers_dir(), "/tmp/drivers")


class ProcessManifestTests(unittest.TestCase):
    def test_skips_manifests_for_another_phase(self):
        self.assertIsNone(driver_helper.process_manifest(manifest(), "pre_wifi"))

    def test_missing_script_is_a_failure_without_execution(self):
        with mock.patch.object(driver_helper, "run_script") as run_script:
            result = driver_helper.process_manifest(manifest(install_script=""), "post_wifi")

        self.assertEqual(result["status"], "failed")
        self.assertEqual(result["message"], "Missing install_script.")
        run_script.assert_not_called()

    def test_not_needed_stops_after_detection(self):
        with mock.patch.object(
            driver_helper,
            "run_script",
            return_value=(1, "not this hardware", ""),
        ) as run_script:
            result = driver_helper.process_manifest(manifest(), "post_wifi")

        self.assertEqual(result["status"], "not_needed")
        self.assertEqual(result["details"], "not this hardware")
        run_script.assert_called_once_with("check")

    def test_unexpected_detection_status_is_a_failure(self):
        with mock.patch.object(driver_helper, "run_script", return_value=(2, "", "probe failed")):
            result = driver_helper.process_manifest(manifest(), "post_wifi")

        self.assertEqual(result["status"], "failed")
        self.assertEqual(result["details"], "probe failed")

    def test_install_failure_does_not_run_verification(self):
        with mock.patch.object(
            driver_helper,
            "run_script",
            side_effect=[(0, "", ""), (1, "install output", "install error")],
        ) as run_script:
            result = driver_helper.process_manifest(manifest(), "post_wifi")

        self.assertEqual(result["message"], "Installation failed.")
        self.assertEqual(result["details"], "install error\n\ninstall output")
        self.assertEqual(run_script.call_count, 2)

    def test_success_reports_reboot_and_verification_output(self):
        stages = []
        with mock.patch.object(
            driver_helper,
            "run_script",
            side_effect=[(0, "", ""), (0, "", ""), (0, "verified", "")],
        ):
            result = driver_helper.process_manifest(manifest(), "post_wifi", stages.append)

        self.assertEqual(result["status"], "success")
        self.assertTrue(result["requires_reboot"])
        self.assertIn("reboot", result["message"].lower())
        self.assertEqual(result["details"], "verified")
        self.assertEqual(stages, ["checking", "installing", "verifying", "complete"])


class RunScriptTests(unittest.TestCase):
    @mock.patch.object(driver_helper.subprocess, "run")
    def test_uses_absolute_bash_and_a_restricted_environment(self, subprocess_run):
        subprocess_run.return_value = mock.Mock(returncode=0, stdout=" output \n", stderr="")
        with mock.patch.dict(os.environ, {"HOME": "/home/test", "LD_PRELOAD": "/tmp/evil"}):
            result = driver_helper.run_script("true")

        self.assertEqual(result, (0, "output", ""))
        args, kwargs = subprocess_run.call_args
        self.assertEqual(args[0], ["/usr/bin/bash", "-c", "true"])
        self.assertEqual(kwargs["env"]["PATH"], driver_helper.SAFE_PATH)
        self.assertEqual(kwargs["env"]["HOME"], "/home/test")
        self.assertNotIn("LD_PRELOAD", kwargs["env"])


if __name__ == "__main__":
    unittest.main()
