from __future__ import annotations

import io
import json
import tempfile
import unittest
from contextlib import redirect_stdout
from unittest.mock import patch

from a2a_sdl.cli import _load_handler_spec, _load_key_registry, _load_payload, _load_trust_governance_policy, main
from a2a_sdl.envelope import build_envelope
from a2a_sdl.schema import get_builtin_descriptor


class CLITests(unittest.TestCase):
    def test_load_payload_requires_one_input(self) -> None:
        with self.assertRaises(ValueError):
            _load_payload(None, None)

    def test_load_trust_governance_policy(self) -> None:
        with tempfile.NamedTemporaryFile("w+", suffix=".json") as handle:
            handle.write(
                '{"approver_keys":{"did:key:ops-a":"pub-a","did:key:ops-b":"pub-b"},"threshold":2}'
            )
            handle.flush()
            policy = _load_trust_governance_policy(handle.name)
        self.assertEqual(policy.threshold, 2)
        self.assertIn("did:key:ops-a", policy.approver_keys)

    def test_send_passes_timeout_and_retry_options(self) -> None:
        captured: dict[str, object] = {}

        def fake_send_http(url: str, envelope: dict, **kwargs):
            captured["url"] = url
            captured["kwargs"] = kwargs
            return build_envelope(
                msg_type="res",
                from_identity=envelope["to"],
                to_identity=envelope["from"],
                content_type="error.v1",
                payload={
                    "code": "BAD_REQUEST",
                    "message": "stub",
                    "details": {},
                    "retryable": False,
                },
                schema=get_builtin_descriptor("error.v1"),
            )

        payload_json = (
            '{"kind":"task.v1","goal":"x","inputs":{},"constraints":{"time_budget_s":1,'
            '"compute_budget":"low","safety":{}},"deliverables":[{"type":"text",'
            '"description":"d"}],"acceptance":["ok"],"context":{}}'
        )

        argv = [
            "send",
            "--url",
            "http://127.0.0.1:9999/a2a",
            "--ct",
            "task.v1",
            "--payload-json",
            payload_json,
            "--timeout",
            "33",
            "--retry-attempts",
            "2",
            "--retry-backoff-s",
            "0.25",
        ]

        with patch("a2a_sdl.cli.send_http", side_effect=fake_send_http):
            with redirect_stdout(io.StringIO()):
                code = main(argv)

        self.assertEqual(code, 0)
        self.assertEqual(captured["url"], "http://127.0.0.1:9999/a2a")
        kwargs = captured["kwargs"]
        assert isinstance(kwargs, dict)
        self.assertEqual(kwargs["timeout"], 33.0)
        self.assertEqual(kwargs["retry_attempts"], 2)
        self.assertEqual(kwargs["retry_backoff_s"], 0.25)

    def test_send_uses_auto_negotiate_sender(self) -> None:
        calls: dict[str, int] = {"send_http": 0, "send_http_with_auto_downgrade": 0}

        def fake_send_http(url: str, envelope: dict, **kwargs):
            calls["send_http"] += 1
            return build_envelope(
                msg_type="res",
                from_identity=envelope["to"],
                to_identity=envelope["from"],
                content_type="error.v1",
                payload={
                    "code": "BAD_REQUEST",
                    "message": "stub",
                    "details": {},
                    "retryable": False,
                },
                schema=get_builtin_descriptor("error.v1"),
            )

        def fake_send_http_with_auto_downgrade(url: str, envelope: dict, **kwargs):
            calls["send_http_with_auto_downgrade"] += 1
            return build_envelope(
                msg_type="res",
                from_identity=envelope["to"],
                to_identity=envelope["from"],
                content_type="negotiation.v1",
                payload={"need": {}, "have": {}, "ask": [], "supported_ct": ["negotiation.v1"]},
                schema=get_builtin_descriptor("negotiation.v1"),
            )

        payload_json = (
            '{"kind":"task.v1","goal":"x","inputs":{},"constraints":{"time_budget_s":1,'
            '"compute_budget":"low","safety":{}},"deliverables":[{"type":"text",'
            '"description":"d"}],"acceptance":["ok"],"context":{}}'
        )

        argv = [
            "send",
            "--url",
            "http://127.0.0.1:9999/a2a",
            "--ct",
            "task.v1",
            "--payload-json",
            payload_json,
            "--auto-negotiate",
        ]

        with patch("a2a_sdl.cli.send_http", side_effect=fake_send_http):
            with patch("a2a_sdl.cli.send_http_with_auto_downgrade", side_effect=fake_send_http_with_auto_downgrade):
                with redirect_stdout(io.StringIO()):
                    code = main(argv)

        self.assertEqual(code, 0)
        self.assertEqual(calls["send_http"], 0)
        self.assertEqual(calls["send_http_with_auto_downgrade"], 1)

    def test_send_secure_adds_enc_sig_and_replay(self) -> None:
        captured: dict[str, object] = {}

        def fake_send_http(url: str, envelope: dict, **kwargs):
            captured["envelope"] = envelope
            return build_envelope(
                msg_type="res",
                from_identity=envelope["to"],
                to_identity=envelope["from"],
                content_type="error.v1",
                payload={
                    "code": "BAD_REQUEST",
                    "message": "stub",
                    "details": {},
                    "retryable": False,
                },
                schema=get_builtin_descriptor("error.v1"),
            )

        payload_json = (
            '{"kind":"task.v1","goal":"x","inputs":{},"constraints":{"time_budget_s":1,'
            '"compute_budget":"low","safety":{}},"deliverables":[{"type":"text",'
            '"description":"d"}],"acceptance":["ok"],"context":{}}'
        )

        with tempfile.TemporaryDirectory() as tmpdir:
            from a2a_sdl.security import generate_signing_keypair, generate_x25519_keypair

            signing = generate_signing_keypair()
            x25519 = generate_x25519_keypair()
            sign_path = f"{tmpdir}/sign.pem"
            with open(sign_path, "w", encoding="utf-8") as f:
                f.write(signing["private_key_pem"])

            argv = [
                "send",
                "--url",
                "http://127.0.0.1:9999/a2a",
                "--ct",
                "task.v1",
                "--payload-json",
                payload_json,
                "--secure",
                "--sign-key",
                sign_path,
                "--sign-kid",
                "did:key:sender#sig1",
                "--encrypt-kid",
                "did:key:receiver#enc1",
                "--encrypt-pub",
                x25519["public_key_b64"],
            ]

            with patch("a2a_sdl.cli.send_http", side_effect=fake_send_http):
                with redirect_stdout(io.StringIO()):
                    code = main(argv)

        self.assertEqual(code, 0)
        envelope = captured["envelope"]
        assert isinstance(envelope, dict)
        sec = envelope["sec"]
        self.assertEqual(sec["mode"], "enc+sig")
        self.assertTrue(isinstance(sec.get("replay"), dict))
        self.assertEqual(sec["sig"]["alg"], "ed25519")

    def test_send_attaches_session_binding(self) -> None:
        captured: dict[str, object] = {}

        def fake_send_http(url: str, envelope: dict, **kwargs):
            _ = url, kwargs
            captured["envelope"] = envelope
            return build_envelope(
                msg_type="res",
                from_identity=envelope["to"],
                to_identity=envelope["from"],
                content_type="error.v1",
                payload={
                    "code": "BAD_REQUEST",
                    "message": "stub",
                    "details": {},
                    "retryable": False,
                },
                schema=get_builtin_descriptor("error.v1"),
            )

        payload_json = (
            '{"kind":"task.v1","goal":"x","inputs":{},"constraints":{"time_budget_s":1,'
            '"compute_budget":"low","safety":{}},"deliverables":[{"type":"text",'
            '"description":"d"}],"acceptance":["ok"],"context":{}}'
        )

        argv = [
            "send",
            "--url",
            "http://127.0.0.1:9999/a2a",
            "--ct",
            "task.v1",
            "--payload-json",
            payload_json,
            "--session-binding-id",
            "sha256:binding-123",
            "--session-binding-exp",
            "2099-01-01T00:00:00Z",
        ]

        with patch("a2a_sdl.cli.send_http", side_effect=fake_send_http):
            with redirect_stdout(io.StringIO()):
                code = main(argv)

        self.assertEqual(code, 0)
        envelope = captured["envelope"]
        assert isinstance(envelope, dict)
        sec = envelope["sec"]
        self.assertEqual(sec["session"]["binding_id"], "sha256:binding-123")
        self.assertEqual(sec["session"]["exp"], "2099-01-01T00:00:00Z")

    def test_send_passes_tls_options(self) -> None:
        captured: dict[str, object] = {}

        def fake_send_http(url: str, envelope: dict, **kwargs):
            _ = envelope
            captured["url"] = url
            captured["kwargs"] = kwargs
            return build_envelope(
                msg_type="res",
                from_identity={"agent_id": "did:key:server", "name": "s", "instance": "i", "role": "executor"},
                to_identity={"agent_id": "did:key:sender", "name": "c", "instance": "i", "role": "planner"},
                content_type="error.v1",
                payload={
                    "code": "BAD_REQUEST",
                    "message": "stub",
                    "details": {},
                    "retryable": False,
                },
                schema=get_builtin_descriptor("error.v1"),
            )

        payload_json = (
            '{"kind":"task.v1","goal":"x","inputs":{},"constraints":{"time_budget_s":1,'
            '"compute_budget":"low","safety":{}},"deliverables":[{"type":"text",'
            '"description":"d"}],"acceptance":["ok"],"context":{}}'
        )
        argv = [
            "send",
            "--url",
            "https://127.0.0.1:9999/a2a",
            "--ct",
            "task.v1",
            "--payload-json",
            payload_json,
            "--tls-ca-file",
            "/tmp/ca.pem",
            "--tls-client-cert-file",
            "/tmp/client.crt",
            "--tls-client-key-file",
            "/tmp/client.key",
            "--tls-insecure-skip-verify",
        ]
        with patch("a2a_sdl.cli.send_http", side_effect=fake_send_http):
            with redirect_stdout(io.StringIO()):
                code = main(argv)

        self.assertEqual(code, 0)
        kwargs = captured["kwargs"]
        assert isinstance(kwargs, dict)
        self.assertEqual(kwargs["tls_ca_file"], "/tmp/ca.pem")
        self.assertEqual(kwargs["tls_client_cert_file"], "/tmp/client.crt")
        self.assertEqual(kwargs["tls_client_key_file"], "/tmp/client.key")
        self.assertTrue(kwargs["tls_insecure_skip_verify"])

    def test_serve_secure_required_needs_key_files(self) -> None:
        with redirect_stdout(io.StringIO()):
            with patch("sys.stderr", new_callable=io.StringIO) as stderr:
                code = main(["serve", "--secure-required"])
        self.assertEqual(code, 2)
        self.assertIn("--secure-required needs", stderr.getvalue())

    def test_serve_rejects_mixed_trustsync_modes(self) -> None:
        with tempfile.NamedTemporaryFile("w+", suffix=".json") as policy_file:
            policy_file.write('{"approver_keys":{"did:key:ops-a":"pub-a"},"threshold":1}')
            policy_file.flush()
            with redirect_stdout(io.StringIO()):
                with patch("sys.stderr", new_callable=io.StringIO) as stderr:
                    code = main(
                        [
                            "serve",
                            "--deployment-mode",
                            "dev",
                            "--trust-governance-file",
                            policy_file.name,
                            "--trust-sync-verify-key-file",
                            "pub-key",
                        ]
                    )
        self.assertEqual(code, 2)
        self.assertIn("use either --trust-governance-file or --trust-sync-verify-key-file", stderr.getvalue())

    def test_serve_session_binding_required_needs_trusted_keys(self) -> None:
        with redirect_stdout(io.StringIO()):
            with patch("sys.stderr", new_callable=io.StringIO) as stderr:
                code = main(["serve", "--session-binding-required"])
        self.assertEqual(code, 2)
        self.assertIn("--session-binding-required needs trusted signing keys", stderr.getvalue())

    def test_serve_prod_requires_tls_by_default(self) -> None:
        with tempfile.NamedTemporaryFile("w+", suffix=".json") as trusted, tempfile.NamedTemporaryFile(
            "w+", suffix=".json"
        ) as decrypt:
            trusted.write('{"kid-1":"pub"}')
            trusted.flush()
            decrypt.write('{"enc-1":"priv"}')
            decrypt.flush()
            with redirect_stdout(io.StringIO()):
                with patch("sys.stderr", new_callable=io.StringIO) as stderr:
                    code = main(
                        [
                            "serve",
                            "--deployment-mode",
                            "prod",
                            "--trusted-signing-keys-file",
                            trusted.name,
                            "--decrypt-keys-file",
                            decrypt.name,
                        ]
                    )
        self.assertEqual(code, 2)
        self.assertIn("requires TLS", stderr.getvalue())

    def test_serve_prod_requires_durable_replay_storage(self) -> None:
        with tempfile.NamedTemporaryFile("w+", suffix=".json") as trusted, tempfile.NamedTemporaryFile(
            "w+", suffix=".json"
        ) as decrypt:
            trusted.write('{"kid-1":"pub"}')
            trusted.flush()
            decrypt.write('{"enc-1":"priv"}')
            decrypt.flush()
            with redirect_stdout(io.StringIO()):
                with patch("sys.stderr", new_callable=io.StringIO) as stderr:
                    code = main(
                        [
                            "serve",
                            "--deployment-mode",
                            "prod",
                            "--trusted-signing-keys-file",
                            trusted.name,
                            "--decrypt-keys-file",
                            decrypt.name,
                            "--tls-cert-file",
                            "cert.pem",
                            "--tls-key-file",
                            "key.pem",
                        ]
                    )
        self.assertEqual(code, 2)
        self.assertIn("requires durable replay storage", stderr.getvalue())

    def test_serve_reports_tls_initialization_error(self) -> None:
        with redirect_stdout(io.StringIO()):
            with patch("sys.stderr", new_callable=io.StringIO) as stderr:
                code = main(
                    [
                        "serve",
                        "--deployment-mode",
                        "dev",
                        "--tls-cert-file",
                        "missing-cert.pem",
                        "--tls-key-file",
                        "missing-key.pem",
                    ]
                )
        self.assertEqual(code, 2)
        self.assertIn("failed to initialize server", stderr.getvalue())

    def test_swarm_requires_minimum_two_ports(self) -> None:
        with redirect_stdout(io.StringIO()):
            with patch("sys.stderr", new_callable=io.StringIO) as stderr:
                code = main(["swarm", "--ports", "9001"])
        self.assertEqual(code, 2)
        self.assertIn("at least 2 ports", stderr.getvalue())

    def test_conformance_command_outputs_json_report(self) -> None:
        with redirect_stdout(io.StringIO()) as stdout:
            code = main(
                [
                    "conformance",
                    "--transport",
                    "core",
                    "--mode",
                    "dev",
                    "--skip-load",
                    "--format",
                    "json",
                ]
            )
        self.assertEqual(code, 0)
        report = json.loads(stdout.getvalue())
        self.assertTrue(report["passed"])
        self.assertEqual(report["summary"]["failed"], 0)

    def test_load_handler_spec_valid(self) -> None:
        ct, handler = _load_handler_spec("artifact.v1=json:loads")
        self.assertEqual(ct, "artifact.v1")
        self.assertTrue(callable(handler))

    def test_load_handler_spec_invalid_format(self) -> None:
        with self.assertRaises(ValueError):
            _load_handler_spec("artifact.v1-json:loads")

    def test_load_handler_spec_rejects_untrusted_module_prefix(self) -> None:
        with self.assertRaises(ValueError):
            _load_handler_spec(
                "artifact.v1=json:loads",
                allowed_module_prefixes=("trusted_pkg",),
            )

    def test_load_handler_spec_requires_manifest_in_strict_mode(self) -> None:
        with self.assertRaises(ValueError):
            _load_handler_spec(
                "artifact.v1=json:loads",
                require_manifest=True,
            )

    def test_load_key_registry_parses_rotation_and_revocation(self) -> None:
        registry_json = """
        {
          "trusted_signing_keys": {"kid-1": "pub1"},
          "required_kid_by_agent": {"agent-a": "kid-1"},
          "allowed_kids_by_agent": {"agent-a": ["kid-1", "kid-2"]},
          "revoked_kids": ["kid-old"],
          "kid_not_after": {"kid-1": "2099-01-01T00:00:00Z"}
        }
        """
        with tempfile.NamedTemporaryFile("w+", suffix=".json") as tmp:
            tmp.write(registry_json)
            tmp.flush()
            registry = _load_key_registry(tmp.name)
        self.assertEqual(registry["trusted_signing_keys"]["kid-1"], "pub1")
        self.assertIn("kid-2", registry["allowed_kids_by_agent"]["agent-a"])
        self.assertIn("kid-old", registry["revoked_kids"])


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