#!/usr/bin/env python3

import base64
import json
import subprocess
import sys
import tempfile
import unittest
from pathlib import Path


PROJECT_ROOT = Path(__file__).resolve().parents[1]
SDK_SRC = PROJECT_ROOT / "packages" / "jwebot-sdk" / "src"
KERNEL_SRC = PROJECT_ROOT / "packages" / "jwebot-kernel" / "src"
sys.path[:0] = [str(SDK_SRC), str(KERNEL_SRC)]

from jwebot_kernel import (  # noqa: E402
    ActivationError,
    ReleaseError,
    SignatureError,
    build_release,
    install_release,
    parse_manifest,
    rollback_release,
    sign_release_record,
    verify_installed_release_signature,
    verify_release_signature,
)


SSH_KEYGEN = Path("/usr/bin/ssh-keygen")
PUBLISHER = "test-publisher"


def manifest():
    return parse_manifest(
        {
            "id": "hello-command",
            "author": "tester",
            "description": "Say hello",
            "help": ["/hello"],
            "version": "1.0.0",
            "api_version": "1",
            "entrypoint": "hello_command:plugin",
            "type": "command",
            "capabilities": {
                "commands": ["/hello"],
                "storage": False,
                "network_profiles": [],
            },
        }
    )


def extension(root: Path, reply: str = "hello") -> Path:
    root.mkdir()
    (root / "src").mkdir()
    (root / "plugin.toml").write_text(
        'id = "hello-command"\n',
        encoding="utf-8",
    )
    (root / "src" / "hello_command.py").write_text(
        "\n".join(
            (
                "from jwebot_sdk import ExtensionResult, ReplyPlan",
                "",
                "def handle(event, arguments):",
                f"    return ExtensionResult(reply=ReplyPlan.text({reply!r}))",
            )
        ),
        encoding="utf-8",
    )
    return root


def publisher_key(root: Path, publisher: str = PUBLISHER):
    root.mkdir(parents=True, exist_ok=True)
    private_key = root / "publisher"
    result = subprocess.run(
        [
            str(SSH_KEYGEN),
            "-q",
            "-t",
            "ed25519",
            "-N",
            "",
            "-f",
            str(private_key),
        ],
        capture_output=True,
        check=False,
    )
    if result.returncode != 0:
        raise RuntimeError("cannot generate test publisher key")
    allowed_signers = root / "allowed_signers"
    public_key = private_key.with_suffix(".pub").read_text(
        encoding="utf-8"
    ).strip()
    allowed_signers.write_text(
        f"{publisher} {public_key}\n",
        encoding="utf-8",
    )
    return private_key, allowed_signers


@unittest.skipUnless(
    SSH_KEYGEN.is_file(),
    "/usr/bin/ssh-keygen is unavailable",
)
class ReleaseSigningTests(unittest.TestCase):
    def build(
        self,
        base: Path,
        name: str = "extension",
        reply: str = "hello",
    ):
        return build_release(
            extension(base / name, reply),
            base / f"{name}-dist",
            manifest=manifest(),
        )

    def test_valid_signature_is_deterministic_and_verifies(self):
        with tempfile.TemporaryDirectory() as directory:
            base = Path(directory)
            _, record_path, _ = self.build(base)
            private_key, allowed_signers = publisher_key(base)

            first_path, first = sign_release_record(
                record_path,
                private_key,
                PUBLISHER,
            )
            first_bytes = first_path.read_bytes()
            second_path, second = sign_release_record(
                record_path,
                private_key,
                PUBLISHER,
            )

            self.assertEqual(first, second)
            self.assertEqual(first_bytes, second_path.read_bytes())
            verified = verify_release_signature(
                record_path,
                first_path,
                allowed_signers,
                expected_publisher=PUBLISHER,
            )
            self.assertEqual(verified.publisher, PUBLISHER)

    def test_signature_rejects_release_record_tampering(self):
        with tempfile.TemporaryDirectory() as directory:
            base = Path(directory)
            _, record_path, _ = self.build(base)
            private_key, allowed_signers = publisher_key(base)
            signature_path, _ = sign_release_record(
                record_path,
                private_key,
                PUBLISHER,
            )
            value = json.loads(record_path.read_text(encoding="utf-8"))
            value["artifact_size"] += 1
            record_path.write_text(
                json.dumps(
                    value,
                    ensure_ascii=False,
                    sort_keys=True,
                    separators=(",", ":"),
                )
                + "\n",
                encoding="utf-8",
            )

            with self.assertRaisesRegex(
                SignatureError,
                "signed release digest mismatch",
            ):
                verify_release_signature(
                    record_path,
                    signature_path,
                    allowed_signers,
                )

    def test_signing_rejects_tampered_artifact(self):
        with tempfile.TemporaryDirectory() as directory:
            base = Path(directory)
            artifact, record_path, _ = self.build(base)
            private_key, _ = publisher_key(base)
            artifact.write_bytes(artifact.read_bytes() + b"tampered")

            with self.assertRaisesRegex(ReleaseError, "size mismatch"):
                sign_release_record(
                    record_path,
                    private_key,
                    PUBLISHER,
                )

    def test_signature_rejects_cryptographic_tampering(self):
        with tempfile.TemporaryDirectory() as directory:
            base = Path(directory)
            _, record_path, _ = self.build(base)
            private_key, allowed_signers = publisher_key(base)
            signature_path, _ = sign_release_record(
                record_path,
                private_key,
                PUBLISHER,
            )
            value = json.loads(signature_path.read_text(encoding="utf-8"))
            raw = bytearray(base64.b64decode(value["signature_base64"]))
            raw[len(raw) // 2] ^= 0x01
            value["signature_base64"] = base64.b64encode(raw).decode("ascii")
            signature_path.write_text(
                json.dumps(
                    value,
                    ensure_ascii=False,
                    sort_keys=True,
                    separators=(",", ":"),
                )
                + "\n",
                encoding="utf-8",
            )

            with self.assertRaisesRegex(
                SignatureError,
                "signature verification failed",
            ):
                verify_release_signature(
                    record_path,
                    signature_path,
                    allowed_signers,
                )

    def test_signature_rejects_untrusted_publisher_key(self):
        with tempfile.TemporaryDirectory() as directory:
            base = Path(directory)
            _, record_path, _ = self.build(base)
            private_key, _ = publisher_key(base / "trusted")
            untrusted_root = base / "untrusted"
            untrusted_root.mkdir()
            _, wrong_allowed_signers = publisher_key(untrusted_root)
            signature_path, _ = sign_release_record(
                record_path,
                private_key,
                PUBLISHER,
            )

            with self.assertRaisesRegex(
                SignatureError,
                "signature verification failed",
            ):
                verify_release_signature(
                    record_path,
                    signature_path,
                    wrong_allowed_signers,
                )

    def test_trust_file_rejects_non_ed25519_entries(self):
        with tempfile.TemporaryDirectory() as directory:
            base = Path(directory)
            _, record_path, _ = self.build(base)
            private_key, allowed_signers = publisher_key(base)
            signature_path, _ = sign_release_record(
                record_path,
                private_key,
                PUBLISHER,
            )
            allowed_signers.write_text(
                allowed_signers.read_text(encoding="utf-8").replace(
                    "ssh-ed25519",
                    "ssh-rsa",
                ),
                encoding="utf-8",
            )

            with self.assertRaisesRegex(SignatureError, "simple Ed25519"):
                verify_release_signature(
                    record_path,
                    signature_path,
                    allowed_signers,
                )

    def test_signature_binds_publisher_even_when_key_has_two_principals(self):
        with tempfile.TemporaryDirectory() as directory:
            base = Path(directory)
            _, record_path, _ = self.build(base)
            private_key, allowed_signers = publisher_key(base)
            public_key = private_key.with_suffix(".pub").read_text(
                encoding="utf-8"
            ).strip()
            with allowed_signers.open("a", encoding="utf-8") as stream:
                stream.write(f"other-publisher {public_key}\n")
            signature_path, _ = sign_release_record(
                record_path,
                private_key,
                PUBLISHER,
            )
            value = json.loads(signature_path.read_text(encoding="utf-8"))
            value["publisher"] = "other-publisher"
            signature_path.write_text(
                json.dumps(
                    value,
                    ensure_ascii=False,
                    sort_keys=True,
                    separators=(",", ":"),
                )
                + "\n",
                encoding="utf-8",
            )

            with self.assertRaisesRegex(
                SignatureError,
                "signature verification failed",
            ):
                verify_release_signature(
                    record_path,
                    signature_path,
                    allowed_signers,
                )

    def test_signing_rejects_noncanonical_release_record(self):
        with tempfile.TemporaryDirectory() as directory:
            base = Path(directory)
            _, record_path, _ = self.build(base)
            private_key, _ = publisher_key(base)
            record_path.write_bytes(b" " + record_path.read_bytes())

            with self.assertRaisesRegex(SignatureError, "not canonical"):
                sign_release_record(
                    record_path,
                    private_key,
                    PUBLISHER,
                )

    def test_signed_activation_persists_and_reverifies_signature(self):
        with tempfile.TemporaryDirectory() as directory:
            base = Path(directory)
            artifact, record_path, _ = self.build(base)
            private_key, allowed_signers = publisher_key(base)
            signature_path, _ = sign_release_record(
                record_path,
                private_key,
                PUBLISHER,
            )

            release_dir = install_release(
                artifact,
                record_path,
                base / "state",
                signature=signature_path,
                allowed_signers=allowed_signers,
                expected_publisher=PUBLISHER,
                health_check=lambda release: None,
            )

            self.assertTrue((release_dir / ".release.sig.json").is_file())
            record, signature = verify_installed_release_signature(
                release_dir,
                allowed_signers,
                expected_publisher=PUBLISHER,
            )
            self.assertEqual(record.plugin_id, "hello-command")
            self.assertEqual(signature.publisher, PUBLISHER)

    def test_signed_activation_requires_passing_health_check(self):
        with tempfile.TemporaryDirectory() as directory:
            base = Path(directory)
            artifact, record_path, _ = self.build(base)
            private_key, allowed_signers = publisher_key(base)
            signature_path, _ = sign_release_record(
                record_path,
                private_key,
                PUBLISHER,
            )
            state = base / "state"

            with self.assertRaisesRegex(
                ActivationError,
                "health check is required",
            ):
                install_release(
                    artifact,
                    record_path,
                    state,
                    signature=signature_path,
                    allowed_signers=allowed_signers,
                    expected_publisher=PUBLISHER,
                )
            self.assertFalse(state.exists())

    def test_failed_health_check_does_not_switch_current(self):
        with tempfile.TemporaryDirectory() as directory:
            base = Path(directory)
            artifact, record_path, record = self.build(base)
            private_key, allowed_signers = publisher_key(base)
            signature_path, _ = sign_release_record(
                record_path,
                private_key,
                PUBLISHER,
            )
            state = base / "state"

            with self.assertRaisesRegex(
                ActivationError,
                "health check failed",
            ):
                install_release(
                    artifact,
                    record_path,
                    state,
                    signature=signature_path,
                    allowed_signers=allowed_signers,
                    expected_publisher=PUBLISHER,
                    health_check=lambda release: (_ for _ in ()).throw(
                        RuntimeError("unhealthy")
                    ),
                )
            self.assertFalse((state / "current").exists())
            self.assertFalse(
                (state / "releases" / record.source_digest).exists()
            )

    def test_signed_activation_rejects_wrong_publisher_policy(self):
        with tempfile.TemporaryDirectory() as directory:
            base = Path(directory)
            artifact, record_path, _ = self.build(base)
            private_key, allowed_signers = publisher_key(base)
            signature_path, _ = sign_release_record(
                record_path,
                private_key,
                PUBLISHER,
            )
            state = base / "state"

            with self.assertRaisesRegex(
                SignatureError,
                "publisher does not match policy",
            ):
                install_release(
                    artifact,
                    record_path,
                    state,
                    signature=signature_path,
                    allowed_signers=allowed_signers,
                    expected_publisher="other-publisher",
                )
            self.assertFalse(state.exists())

    def test_signed_rollback_reverifies_previous_release(self):
        with tempfile.TemporaryDirectory() as directory:
            base = Path(directory)
            private_key, allowed_signers = publisher_key(base)
            first_artifact, first_record_path, first_record = self.build(
                base,
                "first",
                "one",
            )
            first_signature, _ = sign_release_record(
                first_record_path,
                private_key,
                PUBLISHER,
            )
            second_artifact, second_record_path, second_record = self.build(
                base,
                "second",
                "two",
            )
            second_signature, _ = sign_release_record(
                second_record_path,
                private_key,
                PUBLISHER,
            )
            state = base / "state"
            install_release(
                first_artifact,
                first_record_path,
                state,
                signature=first_signature,
                allowed_signers=allowed_signers,
                expected_publisher=PUBLISHER,
                health_check=lambda release: None,
            )
            install_release(
                second_artifact,
                second_record_path,
                state,
                signature=second_signature,
                allowed_signers=allowed_signers,
                expected_publisher=PUBLISHER,
                health_check=lambda release: None,
            )

            destination = rollback_release(
                state,
                allowed_signers=allowed_signers,
                expected_publisher=PUBLISHER,
            )

            self.assertEqual(destination.name, first_record.source_digest)
            self.assertEqual(
                (state / "current").resolve().name,
                first_record.source_digest,
            )
            self.assertEqual(
                (state / "previous").resolve().name,
                second_record.source_digest,
            )


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