"""Tests for ``muse auth keygen`` and ``muse auth register``.

Covers:
- keygen: key generation, file permissions, --force, duplicate key rejection
- keygen: public key / fingerprint format
- register: full challenge-response flow with a mocked hub
- register: token storage in identity.toml
- register: error paths (missing key, network errors, bad challenge token)
- keypair module: sign / verify round-trip, key loading
"""
from __future__ import annotations

import base64
import hashlib
import json
import pathlib
import stat
import unittest.mock
import urllib.error
import urllib.request
import types
from typing import TypedDict

import pytest
from tests.cli_test_helper import CliRunner

from muse.core import keypair as kp_module
from muse.core.store import JsonValue
from muse.core._types import Manifest

type _AuthPayload = dict[str, str | None]
type _JsonResponse = dict[str, JsonValue]


class _ChallengeResp(TypedDict, total=False):
    challengeToken: str
    isNewKey: bool
    algorithm: str


class _VerifyResp(TypedDict, total=False):
    token: str
    handle: str
    identityId: str
    isNewIdentity: bool
    authMethod: str

cli = None
runner = CliRunner()


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _env(tmp_home: pathlib.Path) -> Manifest:
    """Environment that redirects ~/.muse to a temp directory."""
    fake_home = tmp_home / "home"
    fake_home.mkdir(parents=True, exist_ok=True)
    return {"HOME": str(fake_home)}


def _patch_home(monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path) -> pathlib.Path:
    """Redirect pathlib.Path.home() to a temp dir for this test."""
    fake_home = tmp_path / "home"
    fake_home.mkdir(parents=True, exist_ok=True)
    monkeypatch.setattr(pathlib.Path, "home", staticmethod(lambda: fake_home))
    # Also redirect the module-level constants
    monkeypatch.setattr(kp_module, "_KEYS_DIR", fake_home / ".muse" / "keys")
    from muse.core import identity as id_module
    monkeypatch.setattr(id_module, "_IDENTITY_DIR", fake_home / ".muse")
    monkeypatch.setattr(id_module, "_IDENTITY_FILE", fake_home / ".muse" / "identity.toml")
    return fake_home


# ---------------------------------------------------------------------------
# keypair module unit tests
# ---------------------------------------------------------------------------


class TestKeypairModule:
    def test_generate_and_load_roundtrip(self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path) -> None:
        _patch_home(monkeypatch, tmp_path)
        pub_b64, fingerprint = kp_module.generate_keypair("example.com")
        # Fingerprint is 64 hex chars
        assert len(fingerprint) == 64
        assert all(c in "0123456789abcdef" for c in fingerprint)
        # Public key is base64url without padding
        assert "=" not in pub_b64
        # Load the key back
        private_key = kp_module.load_private_key("example.com")
        assert private_key is not None
        # Derived public key must match
        reloaded_pub = kp_module.public_key_to_b64url(private_key.public_key())
        assert reloaded_pub == pub_b64

    def test_key_file_permissions_0o600(self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path) -> None:
        _patch_home(monkeypatch, tmp_path)
        kp_module.generate_keypair("example.com")
        key_path = kp_module.key_path_for("example.com")
        mode = stat.S_IMODE(key_path.stat().st_mode)
        assert mode == 0o600

    def test_keys_dir_permissions_0o700(self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path) -> None:
        _patch_home(monkeypatch, tmp_path)
        kp_module.generate_keypair("example.com")
        keys_dir = kp_module._KEYS_DIR
        mode = stat.S_IMODE(keys_dir.stat().st_mode)
        assert mode == 0o700

    def test_load_nonexistent_key_returns_none(self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path) -> None:
        _patch_home(monkeypatch, tmp_path)
        assert kp_module.load_private_key("no.such.host") is None

    def test_sign_verify_roundtrip(self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path) -> None:
        from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
        from cryptography.hazmat.primitives.serialization import Encoding, PublicFormat

        _patch_home(monkeypatch, tmp_path)
        kp_module.generate_keypair("sign-test.example")
        private_key = kp_module.load_private_key("sign-test.example")
        assert private_key is not None

        nonce = b"\x01\x02\x03" * 10
        sig_b64 = kp_module.sign_bytes(private_key, nonce)
        # Verify the signature independently
        sig_bytes = base64.urlsafe_b64decode(sig_b64 + "==")
        public_key: Ed25519PublicKey = private_key.public_key()
        public_key.verify(sig_bytes, nonce)  # does not raise

    def test_fingerprint_matches_sha256(self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path) -> None:
        from cryptography.hazmat.primitives.serialization import Encoding, PublicFormat

        _patch_home(monkeypatch, tmp_path)
        kp_module.generate_keypair("fp-test.example")
        private_key = kp_module.load_private_key("fp-test.example")
        assert private_key is not None
        public_key = private_key.public_key()
        raw = public_key.public_bytes(Encoding.Raw, PublicFormat.Raw)
        expected = hashlib.sha256(raw).hexdigest()
        assert kp_module.public_key_fingerprint(public_key) == expected

    def test_hostname_sanitisation_colon_port(self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path) -> None:
        _patch_home(monkeypatch, tmp_path)
        kp_module.generate_keypair("localhost:10003")
        path = kp_module.key_path_for("localhost:10003")
        # No colon in filename
        assert ":" not in path.name
        assert path.exists()

    def test_generate_overwrites_existing_key(self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path) -> None:
        _patch_home(monkeypatch, tmp_path)
        pub1, _ = kp_module.generate_keypair("overwrite.test")
        pub2, _ = kp_module.generate_keypair("overwrite.test")
        # Keys should differ (new random key generated)
        assert pub1 != pub2


# ---------------------------------------------------------------------------
# muse auth keygen CLI tests
# ---------------------------------------------------------------------------


class TestAuthKeygenCLI:
    HUB = "http://localhost:10003"

    def test_generates_key_successfully(self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path) -> None:
        _patch_home(monkeypatch, tmp_path)
        result = runner.invoke(cli, ["auth", "keygen", "--hub", self.HUB], catch_exceptions=False)
        assert result.exit_code == 0
        assert "Ed25519 keypair generated" in result.output
        assert "Private key:" in result.output
        assert "Fingerprint" in result.output

    def test_key_file_created(self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path) -> None:
        _patch_home(monkeypatch, tmp_path)
        runner.invoke(cli, ["auth", "keygen", "--hub", self.HUB], catch_exceptions=False)
        key_path = kp_module.key_path_for("localhost:10003")
        assert key_path.exists()

    def test_force_flag_overwrites(self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path) -> None:
        _patch_home(monkeypatch, tmp_path)
        runner.invoke(cli, ["auth", "keygen", "--hub", self.HUB], catch_exceptions=False)
        r1 = runner.invoke(cli, ["auth", "keygen", "--hub", self.HUB], catch_exceptions=False)
        # Without --force, second keygen should fail
        assert r1.exit_code != 0
        r2 = runner.invoke(cli, ["auth", "keygen", "--hub", self.HUB, "--force"], catch_exceptions=False)
        assert r2.exit_code == 0

    def test_no_hub_fails(self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path) -> None:
        _patch_home(monkeypatch, tmp_path)
        # chdir to a directory with no hub config so get_hub_url(None) returns None.
        # MUSE_REPO_ROOT only affects find_repo_root(), not get_hub_url().
        monkeypatch.chdir(tmp_path)
        result = runner.invoke(cli, ["auth", "keygen"])
        assert result.exit_code != 0

    def test_label_shown_in_output(self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path) -> None:
        _patch_home(monkeypatch, tmp_path)
        result = runner.invoke(cli, ["auth", "keygen", "--hub", self.HUB, "--label", "My Laptop"],
                               catch_exceptions=False)
        assert result.exit_code == 0
        assert "My Laptop" in result.output


# ---------------------------------------------------------------------------
# muse auth register CLI tests (hub is mocked)
# ---------------------------------------------------------------------------


class TestAuthRegisterCLI:
    HUB = "http://localhost:10003"

    def _setup_key(self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path) -> tuple[str, str]:
        _patch_home(monkeypatch, tmp_path)
        pub_b64, fingerprint = kp_module.generate_keypair("localhost:10003")
        return pub_b64, fingerprint

    def _mock_hub(
        self,
        monkeypatch: pytest.MonkeyPatch,
        nonce_hex: str,
        challenge_resp: _ChallengeResp | None = None,
        verify_resp: _VerifyResp | None = None,
    ) -> None:
        """Patch urllib.request.urlopen to simulate hub challenge-response."""
        _challenge_resp = challenge_resp or {
            "challengeToken": nonce_hex,
            "isNewKey": True,
            "algorithm": "ed25519",
        }
        _verify_resp = verify_resp or {
            "handle": "alice",
            "identityId": "id-123",
            "isNewIdentity": True,
            "authMethod": "ed25519",
        }
        call_count = 0

        class FakeResponse:
            def __init__(self, data: bytes) -> None:
                self._data = data

            def read(self, n: int = -1) -> bytes:
                return self._data[:n] if n >= 0 else self._data

            def __enter__(self) -> "FakeResponse":
                return self

            def __exit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: types.TracebackType | None) -> None:
                pass

        def fake_urlopen(req: urllib.request.Request, timeout: int = 30) -> FakeResponse:
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                return FakeResponse(json.dumps(_challenge_resp).encode())
            return FakeResponse(json.dumps(_verify_resp).encode())

        import muse.cli.commands.auth as auth_mod
        monkeypatch.setattr(auth_mod, "_json_post_raw", lambda base, path, payload: _challenge_resp if "challenge" in path else _verify_resp)

    def test_full_registration_stores_identity(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        self._setup_key(monkeypatch, tmp_path)
        import secrets
        nonce_hex = secrets.token_hex(32)
        self._mock_hub(monkeypatch, nonce_hex)

        result = runner.invoke(
            cli,
            ["auth", "register", "--hub", self.HUB, "--handle", "alice"],
            catch_exceptions=False,
        )
        assert result.exit_code == 0
        assert "alice" in result.output

        # Ed25519 identity must be persisted
        from muse.core.identity import load_identity
        entry = load_identity(self.HUB)
        assert entry is not None
        assert entry.get("handle") == "alice"
        assert entry.get("key_path") is not None
        assert "token" not in entry

    def test_no_key_fails_gracefully(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        _patch_home(monkeypatch, tmp_path)
        result = runner.invoke(
            cli,
            ["auth", "register", "--hub", self.HUB, "--handle", "alice"],
        )
        assert result.exit_code != 0
        assert "keygen" in result.output.lower() or "no ed25519 key" in result.output.lower()

    def test_new_key_without_handle_fails(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        self._setup_key(monkeypatch, tmp_path)
        import muse.cli.commands.auth as auth_mod

        def fake_json_post(base: str, path: str, payload: _AuthPayload) -> _JsonResponse:
            if "challenge" in path:
                return {
                    "challengeToken": "ab" * 32,
                    "isNewKey": True,
                    "algorithm": "ed25519",
                }
            return {}  # should not reach here

        monkeypatch.setattr(auth_mod, "_json_post_raw", fake_json_post)
        result = runner.invoke(cli, ["auth", "register", "--hub", self.HUB])
        assert result.exit_code != 0
        assert "--handle" in result.output

    def test_agent_flag_marks_identity_as_agent(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        self._setup_key(monkeypatch, tmp_path)
        import muse.cli.commands.auth as auth_mod
        import secrets

        nonce_hex = secrets.token_hex(32)

        def fake_json_post(base: str, path: str, payload: _AuthPayload) -> _JsonResponse:
            if "challenge" in path:
                return {"challengeToken": nonce_hex, "isNewKey": False, "algorithm": "ed25519"}
            return {"handle": "bot", "identityId": "id-bot", "isNewIdentity": False, "authMethod": "ed25519"}

        monkeypatch.setattr(auth_mod, "_json_post_raw", fake_json_post)
        runner.invoke(cli, ["auth", "register", "--hub", self.HUB, "--handle", "bot", "--agent"], catch_exceptions=False)

        from muse.core.identity import load_identity
        entry = load_identity(self.HUB)
        assert entry is not None
        assert entry.get("type") == "agent"

    def test_bad_challenge_token_fails(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        self._setup_key(monkeypatch, tmp_path)
        import muse.cli.commands.auth as auth_mod

        def fake_json_post(base: str, path: str, payload: _AuthPayload) -> _JsonResponse:
            return {"challengeToken": "not-valid-hex!", "isNewKey": False, "algorithm": "ed25519"}

        monkeypatch.setattr(auth_mod, "_json_post_raw", fake_json_post)
        result = runner.invoke(cli, ["auth", "register", "--hub", self.HUB, "--handle", "alice"])
        assert result.exit_code != 0

    def test_empty_verify_response_fails(
        self, monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
    ) -> None:
        """Verify response with no handle and no --handle fallback fails."""
        self._setup_key(monkeypatch, tmp_path)
        import muse.cli.commands.auth as auth_mod
        import secrets

        nonce_hex = secrets.token_hex(32)

        def fake_json_post(base: str, path: str, payload: _AuthPayload) -> _JsonResponse:
            if "challenge" in path:
                return {"challengeToken": nonce_hex, "isNewKey": False, "algorithm": "ed25519"}
            return {}  # No handle field

        monkeypatch.setattr(auth_mod, "_json_post_raw", fake_json_post)
        # No --handle flag, and hub returns no handle → must fail
        result = runner.invoke(cli, ["auth", "register", "--hub", self.HUB])
        assert result.exit_code != 0