import os
import subprocess
import tempfile
import unittest

REPO_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
SCRIPT_PATH = os.path.join(REPO_DIR, "cloud-uptime", "recover-after-rebuild.sh")

class TestRecoverAfterRebuild(unittest.TestCase):
    def test_script_exists_and_executable(self):
        self.assertTrue(os.path.exists(SCRIPT_PATH), f"Script missing: {SCRIPT_PATH}")
        self.assertTrue(os.access(SCRIPT_PATH, os.X_OK), "Script not executable")

    def test_bash_syntax_check(self):
        proc = subprocess.run(["bash", "-n", SCRIPT_PATH], capture_output=True, text=True)
        self.assertEqual(proc.returncode, 0, f"Bash syntax error: {proc.stderr}")

    def test_dry_run_execution(self):
        proc = subprocess.run([SCRIPT_PATH, "--dry-run"], capture_output=True, text=True)
        self.assertEqual(proc.returncode, 0, f"Dry-run failed: {proc.stderr}")
        self.assertIn("running in DRY-RUN mode", proc.stdout)
        self.assertIn("machine:", proc.stdout)

    def test_machine_env_override(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            ws_tunnel = os.path.join(tmpdir, "workspace", "tunnel")
            os.makedirs(ws_tunnel, exist_ok=True)
            env_file = os.path.join(ws_tunnel, "machine.env")
            with open(env_file, "w") as f:
                f.write("MUSE_MACHINE=custom-test-node\nSSH_PORT=9922\nTERM_PORT=8877\n")
            
            env = os.environ.copy()
            env["HOME"] = tmpdir
            proc = subprocess.run([SCRIPT_PATH, "--dry-run"], env=env, capture_output=True, text=True)
            self.assertEqual(proc.returncode, 0, f"Run with env failed: {proc.stderr}")
            self.assertIn("custom-test-node", proc.stdout)
            self.assertIn("9922", proc.stdout)
            self.assertIn("8877", proc.stdout)

    def test_root_authorized_keys_preservation(self):
        with tempfile.TemporaryDirectory() as tmpdir:
            ws_tunnel = os.path.join(tmpdir, "workspace", "tunnel")
            hatch_ssh = os.path.join(tmpdir, "home", "hatch", ".ssh")
            root_ssh = os.path.join(tmpdir, "root", ".ssh")
            os.makedirs(ws_tunnel, exist_ok=True)
            os.makedirs(hatch_ssh, exist_ok=True)
            os.makedirs(root_ssh, exist_ok=True)

            backup_key_path = os.path.join(ws_tunnel, "root-authorized_keys")
            with open(backup_key_path, "w") as f:
                f.write("ssh-ed25519 AAAABACKUP1 root@backup\nssh-ed25519 AAASHARED common@shared\n")

            hatch_key_path = os.path.join(hatch_ssh, "authorized_keys")
            with open(hatch_key_path, "w") as f:
                f.write("ssh-ed25519 AAAAHATCH1 hatch@box\nssh-ed25519 AAASHARED common@shared\n")

            target_root_keys = os.path.join(root_ssh, "authorized_keys")
            bash_cmd = f"""
            if [ -f "{backup_key_path}" ]; then
                install -m 600 "{backup_key_path}" "{target_root_keys}"
            fi
            if [ -f "{hatch_key_path}" ]; then
                cat "{hatch_key_path}" >> "{target_root_keys}"
                sort -u "{target_root_keys}" -o "{target_root_keys}"
                chmod 600 "{target_root_keys}"
            fi
            """
            proc = subprocess.run(["bash", "-c", bash_cmd], capture_output=True, text=True)
            self.assertEqual(proc.returncode, 0, f"Key merge script failed: {proc.stderr}")

            self.assertTrue(os.path.exists(target_root_keys))
            with open(target_root_keys) as f:
                content = f.read()

            self.assertIn("root@backup", content)
            self.assertIn("hatch@box", content)
            self.assertIn("common@shared", content)
            # Ensure sort -u eliminated duplicate shared key
            self.assertEqual(content.count("common@shared"), 1)
            # Ensure permissions are 0600
            perms = oct(os.stat(target_root_keys).st_mode & 0o777)
            self.assertEqual(perms, "0o600")

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