import re
import os
import shutil
import stat
import subprocess
import tempfile
import unittest
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
NIX_FILE = ROOT / "hosts/asusmini/auto-update.nix"
BASH = shutil.which("bash") or "/run/current-system/sw/bin/bash"


FAKE = r'''#!/usr/bin/env python3
import os
import shutil
import sys
from pathlib import Path

name = Path(sys.argv[0]).name
args = sys.argv[1:]
log = Path(os.environ["FAKE_LOG"])
with log.open("a") as f:
    f.write(name + " " + " ".join(args) + "\n")

def fail(flag):
    return os.environ.get(flag) == "1"

if name == "flock":
    pass
elif name == "mktemp":
    p = Path(os.environ["TMP_ROOT"]) / "lock-snapshot"
    p.touch()
    print(p)
elif name == "cp":
    shutil.copyfile(args[-2], args[-1])
elif name == "rm":
    for arg in args:
        if not arg.startswith("-"):
            Path(arg).unlink(missing_ok=True)
elif name == "date":
    print("fixture-date")
elif name == "git":
    operation = " ".join(args)
    if "diff --cached --quiet" in operation and fail("STAGED_CHANGES"):
        raise SystemExit(1)
    if "diff --quiet -- *.nix" in operation and fail("NIX_CHANGES"):
        raise SystemExit(1)
    if "diff --quiet -- flake.lock" in operation:
        raise SystemExit(0 if not Path("flake.lock").read_text().startswith("updated") else 1)
    if args[:1] == ["pull"]:
        if fail("PULL_FAIL"):
            raise SystemExit(1)
        Path("flake.lock").write_text("pulled lock\n")
    elif args[:1] == ["push"] and "--dry-run" not in args and fail("PUBLISH_FAIL"):
        raise SystemExit(1)
    elif args[:1] == ["commit"]:
        if fail("COMMIT_FAIL"):
            raise SystemExit(1)
        Path("flake.lock").write_text("committed lock\n")
elif name == "nix":
    if args[:2] == ["flake", "update"]:
        Path("flake.lock").write_text("updated lock\n")
    elif args[:1] == ["build"]:
        if fail("BUILD_FAIL"):
            raise SystemExit(1)
        print(os.environ["TOPLEVEL"])
elif name == "nix-env":
    if fail("ENV_SET_FAIL"):
        raise SystemExit(1)
elif name == "systemctl":
    if args[:2] == ["is-active", "--quiet"] and fail("HEALTH_FAIL"):
        raise SystemExit(1)
    if args[:1] == ["stop"] and fail("STOP_FAIL"):
        raise SystemExit(1)
elif name == "curl":
    if fail("HEALTH_FAIL"):
        raise SystemExit(1)
elif name == "jq":
    sys.stdin.read()
    if fail("HEALTH_FAIL"):
        raise SystemExit(1)
else:
    Path(os.environ["UNKNOWN_LOG"]).open("a").write(name + "\n")
    raise SystemExit(127)
'''


class AutoUpdateTests(unittest.TestCase):
    def setUp(self):
        self.tmp = Path(tempfile.mkdtemp(prefix="auto-update-test-"))
        self.bin = self.tmp / "bin"
        self.bin.mkdir()
        self.log = self.tmp / "commands.log"
        self.unknown = self.tmp / "unknown.log"
        fake = self.bin / "fake.py"
        fake.write_text(FAKE)
        fake.chmod(fake.stat().st_mode | stat.S_IXUSR)
        for name in "flock mktemp cp rm date git nix nix-env systemctl curl jq".split():
            (self.bin / name).symlink_to(fake)

        self.repo = self.tmp / "repo"
        self.repo.mkdir()
        (self.repo / "flake.lock").write_text("initial lock\n")
        self.toplevel = self.tmp / "toplevel"
        (self.toplevel / "bin").mkdir(parents=True)
        switch = self.toplevel / "bin" / "switch-to-configuration"
        switch.write_text("#!/bin/sh\nif [ \"${ACTIVATE_FAIL:-0}\" = 1 ]; then exit 1; fi\nexit 0\n")
        switch.chmod(0o755)
        self.state = self.tmp / "state"
        (self.state).mkdir()
        (self.state / ".upgrade-verified").touch()

        script = NIX_FILE.read_text()
        body = script.split("script = ''", 1)[1].split("'';", 1)[0]
        body = re.sub(r"\$\{pkgs\.[^}]+\}/bin/", "", body)
        replacements = {
            "/run/lock/commonscomputing-update.lock": self.tmp / "update.lock",
            "/var/lib/knot/.upgrade-verified": self.state / ".upgrade-verified",
            "/nix/var/nix/profiles/system": self.tmp / "profile",
        }
        for old, new in replacements.items():
            body = body.replace(old, str(new))
        self.script = self.tmp / "auto-update.sh"
        self.script.write_text("#!/usr/bin/env bash\n" + body)
        self.script.chmod(0o755)
        self.env = os.environ.copy()
        self.env.update({"PATH": str(self.bin) + ":" + self.env["PATH"], "FAKE_LOG": str(self.log), "UNKNOWN_LOG": str(self.unknown), "TMP_ROOT": str(self.tmp), "TOPLEVEL": str(self.toplevel)})

    def tearDown(self):
        shutil.rmtree(self.tmp, ignore_errors=True)

    def run_update(self, **flags):
        env = self.env.copy()
        env.update({k: str(v) for k, v in flags.items()})
        return subprocess.run([BASH, str(self.script)], cwd=self.repo, env=env, text=True, capture_output=True)

    def commands(self):
        return self.log.read_text().splitlines() if self.log.exists() else []

    def assert_no_unknown(self):
        self.assertFalse(self.unknown.exists(), self.unknown.read_text() if self.unknown.exists() else "")

    def test_build_failure_restores_post_pull_lock_without_activation_or_commit(self):
        result = self.run_update(BUILD_FAIL=1)
        self.assertNotEqual(result.returncode, 0)
        self.assertEqual((self.repo / "flake.lock").read_text(), "pulled lock\n")
        commands = self.commands()
        self.assertFalse(any(x.startswith("nix-env") for x in commands))
        self.assertFalse(any(x.startswith("git commit") for x in commands))
        self.assert_no_unknown()

    def test_activation_failure_pauses_updates_without_rollback_or_publish(self):
        result = self.run_update(ACTIVATE_FAIL=1)
        self.assertNotEqual(result.returncode, 0)
        commands = self.commands()
        self.assertIn("systemctl stop auto-update.timer", commands)
        self.assertEqual(sum(x.startswith("nix-env") for x in commands), 1)
        self.assertFalse(any(x.startswith("git commit") or x == "git push" for x in commands))
        self.assertTrue((self.state / ".upgrade-verified").exists() is False)
        self.assert_no_unknown()

    def test_health_failure_pauses_updates_without_rollback_or_publish(self):
        result = self.run_update(HEALTH_FAIL=1)
        self.assertNotEqual(result.returncode, 0)
        commands = self.commands()
        self.assertIn("systemctl stop auto-update.timer", commands)
        self.assertEqual(sum(x.startswith("nix-env") for x in commands), 1)
        self.assertFalse(any(x.startswith("git commit") or x == "git push" for x in commands))
        self.assert_no_unknown()

    def test_happy_path_commits_and_pushes_after_health(self):
        result = self.run_update()
        self.assertEqual(result.returncode, 0, result.stderr)
        commands = self.commands()
        health = max(i for i, x in enumerate(commands) if x.startswith("curl "))
        commit = next(i for i, x in enumerate(commands) if x.startswith("git commit"))
        push = next(i for i, x in enumerate(commands) if x == "git push")
        self.assertLess(health, commit)
        self.assertLess(commit, push)
        self.assert_no_unknown()

    def test_publication_failure_retains_deployed_lock_without_rollback(self):
        result = self.run_update(PUBLISH_FAIL=1)
        self.assertNotEqual(result.returncode, 0)
        self.assertEqual((self.repo / "flake.lock").read_text(), "committed lock\n")
        commands = self.commands()
        self.assertFalse(any(x.startswith("nix-env") and i > commands.index("git push") for i, x in enumerate(commands)))
        self.assert_no_unknown()

    def test_staged_and_nix_edits_refuse_before_pull(self):
        for flag in ("STAGED_CHANGES", "NIX_CHANGES"):
            with self.subTest(flag=flag):
                result = self.run_update(**{flag: 1})
                self.assertNotEqual(result.returncode, 0)
                self.assertFalse(any(x.startswith("git pull") for x in self.commands()))
                self.assert_no_unknown()


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