#!/usr/bin/python3
"""Regression coverage for the preinstalled Broadcom wl driver contract."""

from __future__ import annotations

import json
import os
from pathlib import Path
import subprocess
import tempfile
import unittest


PROJECT_ROOT = Path(__file__).resolve().parents[1]
MANIFEST_PATH = PROJECT_ROOT / "data" / "drivers" / "broadcom-wifi.json"


class BroadcomDriverManifestTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.manifest = json.loads(MANIFEST_PATH.read_text(encoding="utf-8"))
        cls.check_script = "\n".join(cls.manifest["check_script"])
        cls.install_script = "\n".join(cls.manifest["install_script"])
        cls.success_script = "\n".join(cls.manifest["success_script"])

    def test_is_enabled(self):
        self.assertIs(self.manifest["enabled"], True)
        self.assertEqual(self.manifest["phase"], "pre_wifi")

    def test_detection_uses_the_installed_wl_modaliases(self):
        self.assertIn('pci:v000014E4d*bc02sc80i*', self.check_script)
        self.assertIn('modinfo -k "$(uname -r)" wl', self.check_script)
        self.assertIn('modprobe --resolve-alias "$modalias"', self.check_script)
        self.assertNotIn("grep -qE", self.check_script)

    def test_checks_for_the_preinstalled_wl_module_before_unblocking_it(self):
        self.assertIn('modinfo -k "$(uname -r)" wl', self.install_script)
        self.assertLess(
            self.install_script.index('modinfo -k "$(uname -r)" wl'),
            self.install_script.index("rm -f /etc/modprobe.d/broadcom-wl-blacklist.conf"),
        )

    def test_reports_a_clear_error_when_the_image_does_not_contain_wl(self):
        self.assertIn("does not contain a prebuilt Broadcom wl module", self.install_script)

    def test_verifies_wl_module_metadata_after_loading(self):
        self.assertIn('modinfo -k "$(uname -r)" wl', self.success_script)

    def test_unloads_every_rpmfusion_conflict_before_wl(self):
        unload_loop = (
            "for module in b43 brcmsmac brcmfmac_wcc brcmfmac_cyw "
            "brcmfmac_bca brcmfmac bcma ssb; do"
        )

        self.assertIn(unload_loop, self.install_script)
        self.assertIn('if ! modprobe -r "$module" 2>/dev/null; then', self.install_script)
        self.assertNotIn('modprobe -r "$module" 2>/dev/null || true', self.install_script)
        self.assertLess(
            self.install_script.index(unload_loop),
            self.install_script.index("modprobe wl"),
        )

    def test_restores_the_boot_block_if_wl_activation_fails(self):
        self.assertIn("restore_broadcom_block()", self.install_script)
        self.assertIn(
            "trap restore_broadcom_block EXIT",
            self.install_script,
        )
        self.assertIn(
            "trap 'exit 1' HUP INT TERM",
            self.install_script,
        )
        self.assertIn(
            "ln -sfn /dev/null /etc/modprobe.d/broadcom-wl-blacklist.conf",
            self.install_script,
        )
        self.assertIn('"blacklist wl"', self.install_script)
        self.assertIn('"install wl /bin/false"', self.install_script)
        self.assertIn('modprobe -r wl 2>/dev/null || true', self.install_script)
        self.assertIn('removed_modules="$module $removed_modules"', self.install_script)
        self.assertIn('for module in $removed_modules; do', self.install_script)
        self.assertLess(
            self.install_script.index("trap restore_broadcom_block EXIT"),
            self.install_script.index("rm -f /etc/modprobe.d/broadcom-wl-blacklist.conf"),
        )
        self.assertLess(
            self.install_script.index("modprobe wl"),
            self.install_script.index('bound_module=${bound_module##*/}'),
        )
        self.assertLess(
            self.install_script.index('bound_module=${bound_module##*/}'),
            self.install_script.rindex("trap - EXIT HUP INT TERM"),
        )

    def run_check_script(
        self,
        *,
        modalias: str,
        resolved_modules: str,
        wl_available: bool = True,
    ):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            devices_dir = root / "devices"
            fake_bin = root / "bin"
            device_dir = devices_dir / "0000:01:00.0"
            device_dir.mkdir(parents=True)
            fake_bin.mkdir()
            (device_dir / "modalias").write_text(modalias + "\n", encoding="utf-8")

            commands = {
                "uname": "printf '%s\\n' test-kernel",
                "modinfo": 'test "$WL_AVAILABLE" = 1',
                "modprobe": (
                    "if [[ \"$1\" = --resolve-alias ]]; then "
                    "printf '%s\\n' \"$RESOLVED_MODULES\"; fi"
                ),
            }
            for name, body in commands.items():
                path = fake_bin / name
                path.write_text(f"#!/usr/bin/bash\n{body}\n", encoding="utf-8")
                path.chmod(0o755)

            env = os.environ.copy()
            env.update(
                {
                    "PATH": f"{fake_bin}:/usr/bin:/bin",
                    "QUICK_SETUP_PCI_DEVICES_DIR": str(devices_dir),
                    "RESOLVED_MODULES": resolved_modules,
                    "WL_AVAILABLE": "1" if wl_available else "0",
                }
            )
            return subprocess.run(
                ["/usr/bin/bash", "-c", self.check_script],
                text=True,
                capture_output=True,
                env=env,
                check=False,
            )

    def test_detection_accepts_only_broadcom_devices_resolved_to_wl(self):
        broadcom = "pci:v000014E4d000043A0sv0000106Bsd00000134bc02sc80i00"
        intel = "pci:v00008086d00002724sv00008086sd00000080bc02sc80i00"

        self.assertEqual(
            self.run_check_script(modalias=broadcom, resolved_modules="wl").returncode,
            0,
        )
        self.assertEqual(
            self.run_check_script(modalias=broadcom, resolved_modules="brcmfmac").returncode,
            1,
        )
        self.assertEqual(
            self.run_check_script(modalias=intel, resolved_modules="wl").returncode,
            1,
        )

    def test_missing_wl_on_broadcom_reaches_the_actionable_install_failure(self):
        broadcom = "pci:v000014E4d000043A0sv0000106Bsd00000134bc02sc80i00"

        result = self.run_check_script(
            modalias=broadcom,
            resolved_modules="",
            wl_available=False,
        )

        self.assertEqual(result.returncode, 0)
        self.assertIn("wl module is unavailable", result.stdout)

    def test_missing_wl_does_not_claim_a_brcmfmac_device(self):
        brcmfmac_device = "pci:v000014E4d000043C5sv0000106Bsd00000134bc02sc80i00"

        result = self.run_check_script(
            modalias=brcmfmac_device,
            resolved_modules="",
            wl_available=False,
        )

        self.assertEqual(result.returncode, 1)
        self.assertIn("Deferring to in-kernel drivers", result.stdout)

    def run_install_script(
        self,
        *,
        fail_wl: bool = False,
        fail_unload: str = "",
        bind_wl: bool = True,
    ):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            config_dir = root / "modprobe.d"
            devices_dir = root / "devices"
            fake_bin = root / "bin"
            device_dir = devices_dir / "0000:01:00.0"
            driver_dir = device_dir / "driver"
            modules_dir = root / "modules"
            brcmfmac_target = modules_dir / "brcmfmac"
            wl_target = modules_dir / "wl"
            config_dir.mkdir()
            fake_bin.mkdir()
            driver_dir.mkdir(parents=True)
            modules_dir.mkdir()
            for module in (
                "b43",
                "brcmsmac",
                "brcmfmac_wcc",
                "brcmfmac_cyw",
                "brcmfmac_bca",
                "brcmfmac",
                "bcma",
                "ssb",
                "wl",
            ):
                (modules_dir / module).mkdir()
            (device_dir / "modalias").write_text(
                "pci:v000014E4d000043A0sv0000106Bsd00000134bc02sc80i00\n",
                encoding="utf-8",
            )
            module_link = driver_dir / "module"
            module_link.symlink_to(brcmfmac_target, target_is_directory=True)

            vendor_mask = config_dir / "broadcom-wl-blacklist.conf"
            wl_block = config_dir / "default-disable-broadcom-wl.conf"
            vendor_mask.symlink_to("/dev/null")
            wl_block.write_text(
                "blacklist wl\ninstall wl /bin/false\n",
                encoding="utf-8",
            )

            commands = {
                "uname": "printf '%s\\n' test-kernel",
                "modinfo": "exit 0",
                "depmod": "exit 0",
                "modprobe": (
                    "printf '%s\\n' \"$*\" >> \"$MODPROBE_LOG\"\n"
                    "if [[ \"$1\" = --resolve-alias ]]; then printf '%s\\n' wl; exit 0; fi\n"
                    "if [[ -n \"$FAIL_UNLOAD\" && \"$*\" = \"-r $FAIL_UNLOAD\" ]]; "
                    "then exit 1; fi\n"
                    "if [[ \"$*\" = '-r brcmfmac' ]]; then rm -f \"$DEVICE_MODULE_LINK\"; fi\n"
                    "if [[ \"$*\" = '-r wl' && \"$(basename \"$(readlink -f \"$DEVICE_MODULE_LINK\" 2>/dev/null)\")\" = wl ]]; "
                    "then rm -f \"$DEVICE_MODULE_LINK\"; fi\n"
                    "if [[ \"$FAIL_WL\" = 1 && \"$*\" = wl ]]; then exit 1; fi\n"
                    "if [[ \"$BIND_WL\" = 1 && \"$*\" = wl ]]; then "
                    "ln -sfn \"$WL_MODULE_TARGET\" \"$DEVICE_MODULE_LINK\"; fi\n"
                    "if [[ \"$*\" = brcmfmac ]]; then "
                    "ln -sfn \"$BRCMFMAC_MODULE_TARGET\" \"$DEVICE_MODULE_LINK\"; fi"
                ),
            }
            for name, body in commands.items():
                path = fake_bin / name
                path.write_text(f"#!/usr/bin/bash\n{body}\n", encoding="utf-8")
                path.chmod(0o755)

            log_path = root / "modprobe.log"
            env = os.environ.copy()
            env.update(
                {
                    "PATH": f"{fake_bin}:/usr/bin:/bin",
                    "QUICK_SETUP_PCI_DEVICES_DIR": str(devices_dir),
                    "QUICK_SETUP_MODULES_DIR": str(modules_dir),
                    "MODPROBE_LOG": str(log_path),
                    "DEVICE_MODULE_LINK": str(module_link),
                    "WL_MODULE_TARGET": str(wl_target),
                    "BRCMFMAC_MODULE_TARGET": str(brcmfmac_target),
                    "FAIL_WL": "1" if fail_wl else "0",
                    "FAIL_UNLOAD": fail_unload,
                    "BIND_WL": "1" if bind_wl else "0",
                }
            )
            script = self.install_script.replace("/etc/modprobe.d", str(config_dir))
            result = subprocess.run(
                ["/usr/bin/bash", "-c", script],
                text=True,
                capture_output=True,
                env=env,
                check=False,
            )

            return {
                "returncode": result.returncode,
                "vendor_mask_exists": vendor_mask.is_symlink(),
                "vendor_mask_target": os.readlink(vendor_mask) if vendor_mask.is_symlink() else None,
                "wl_block": wl_block.read_text(encoding="utf-8") if wl_block.exists() else None,
                "modprobe_log": log_path.read_text(encoding="utf-8"),
                "bound_module": module_link.resolve().name if module_link.is_symlink() else None,
            }

    def test_successful_activation_unblocks_and_loads_wl(self):
        result = self.run_install_script()

        self.assertEqual(result["returncode"], 0)
        self.assertFalse(result["vendor_mask_exists"])
        self.assertIsNone(result["wl_block"])
        for module in (
            "b43",
            "brcmsmac",
            "brcmfmac_wcc",
            "brcmfmac_cyw",
            "brcmfmac_bca",
            "brcmfmac",
            "bcma",
            "ssb",
        ):
            self.assertIn(f"-r {module}", result["modprobe_log"].splitlines())
        self.assertIn("wl", result["modprobe_log"].splitlines())
        self.assertEqual(result["bound_module"], "wl")

    def assert_boot_block_restored(self, result):
        self.assertNotEqual(result["returncode"], 0)
        self.assertTrue(result["vendor_mask_exists"])
        self.assertEqual(result["vendor_mask_target"], "/dev/null")
        self.assertEqual(result["wl_block"], "blacklist wl\ninstall wl /bin/false\n")
        self.assertEqual(result["bound_module"], "brcmfmac")
        self.assertIn("-r wl", result["modprobe_log"].splitlines())

    def test_failed_activation_restores_the_boot_block(self):
        self.assert_boot_block_restored(self.run_install_script(fail_wl=True))

    def test_failed_conflict_unload_restores_the_boot_block(self):
        self.assert_boot_block_restored(self.run_install_script(fail_unload="brcmfmac"))

    def test_unbound_wl_restores_the_boot_block(self):
        self.assert_boot_block_restored(self.run_install_script(bind_wl=False))


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