"""Tests for aablocks.auth module.

Auth uses OAuth2 authorization-code + PKCE flow against AWS Cognito.
The token lifecycle is:

    login() -> browser OAuth -> exchange code -> _save_token() -> disk + memory
    get_token() -> _load_token() -> memory cache or disk -> _refresh() if expired
    get_token_header() -> get_token() -> {"Authorization": "Bearer <id_token>", ...}
    logout() -> clear memory cache + delete token file

All tests that touch disk use the tmp_aab_dir fixture (from conftest.py)
which redirects ~/.aab to a temp directory. All HTTP calls are mocked so
nothing hits Cognito for real.
"""

from __future__ import annotations

import base64
import hashlib
import json
from time import time
from unittest.mock import MagicMock, patch
from urllib.parse import parse_qs, urlparse

import pytest
import requests

from aablocks.auth import (
    _decode_paste_token,
    _gated_login_url,
    _generate_code_challenge,
    _generate_code_verifier,
    _load_token,
    _manual_redirect_uri,
    _refresh,
    _save_token,
    get_token,
    get_token_header,
    login,
    logout,
)
from aablocks.cached_token import clear_cached_token, get_cached_token
from aablocks.config import set_config

DOMAIN = "auth.atlas.aalphabio.com"


# ---------------------------------------------------------------------------
# PKCE helpers
#
# PKCE (Proof Key for Code Exchange, RFC 7636) prevents authorization code
# interception attacks. The client generates a random code_verifier, hashes
# it into a code_challenge sent with the auth request, then proves possession
# of the verifier when exchanging the code for tokens.
# ---------------------------------------------------------------------------
class TestPKCE:
    def test_code_verifier_length(self):
        """Generated verifier is between 43 and 128 characters (RFC 7636)."""
        v = _generate_code_verifier(length=64)
        assert 43 <= len(v) <= 128

    def test_code_verifier_url_safe_chars(self):
        """Verifier contains only URL-safe base64 characters with no padding."""
        v = _generate_code_verifier()
        # RFC 7636 section 4.1: unreserved chars [A-Z] / [a-z] / [0-9] / "-" / "." / "_" / "~"
        # Our implementation uses base64url which is a subset of that.
        allowed = set(
            "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_"
        )
        assert set(v).issubset(allowed)

    def test_code_verifier_unique(self):
        """Two generated verifiers are different (randomness check)."""
        v1 = _generate_code_verifier()
        v2 = _generate_code_verifier()
        assert v1 != v2

    def test_code_challenge_is_s256(self):
        """Challenge matches the SHA-256 hash of the verifier (S256 method)."""
        verifier = "test_verifier_string"
        challenge = _generate_code_challenge(verifier)
        # Independently compute the expected S256 challenge to verify our implementation
        digest = hashlib.sha256(verifier.encode("ascii")).digest()
        expected = base64.urlsafe_b64encode(digest).decode("ascii").rstrip("=")
        assert challenge == expected


# ---------------------------------------------------------------------------
# Token loading / saving / refreshing
#
# Token resolution order in _load_token():
#   1. In-memory cache (fastest, avoids disk I/O)
#   2. Disk file (~/.aab/token.json)
#   3. If expired, attempt _refresh() using the refresh_token
#   4. If refresh fails or no token exists, return None (caller triggers login)
#
# _save_token() writes to both disk and memory cache, and converts the
# relative expires_in field to an absolute expires_at timestamp.
# ---------------------------------------------------------------------------
class TestLoadToken:
    def setup_method(self):
        """Start each test with a clean in-memory cache."""
        clear_cached_token(DOMAIN)

    def test_returns_none_when_no_file(self, tmp_aab_dir):
        """Returns None when no token file exists and cache is empty."""
        result = _load_token()
        assert result is None

    def test_loads_valid_token_from_file(self, write_token):
        """Reads a valid, non-expired token from disk and caches it in memory."""
        # expires_at is 1 hour in the future — well within the buffer window
        token = {
            "id_token": "abc123",
            "refresh_token": "refresh_abc",
            "expires_at": int(time()) + 3600,
        }
        write_token(token)
        with patch("aablocks.auth.log_auth"):
            result = _load_token()
        assert result is not None
        assert result["id_token"] == "abc123"
        # Verify the token was also stored in the in-memory cache
        assert get_cached_token(DOMAIN) is not None

    def test_refreshes_expired_token(self, write_token):
        """An expired token on disk triggers a refresh and returns the new token."""
        # expires_at in the past — _load_token should call _refresh()
        token = {
            "id_token": "old_token",
            "refresh_token": "refresh_abc",
            "expires_at": int(time()) - 100,
        }
        write_token(token)

        # Simulate Cognito returning a fresh token (without a refresh_token,
        # which is normal — Cognito sometimes omits it on refresh)
        new_token = {
            "id_token": "new_token",
            "access_token": "new_access",
            "expires_in": 3600,
        }

        mock_resp = MagicMock()
        mock_resp.ok = True
        mock_resp.json.return_value = new_token

        with (
            patch("aablocks.auth.requests.post", return_value=mock_resp),
            patch("aablocks.auth.log_auth"),
        ):
            result = _load_token()

        assert result is not None
        assert result["id_token"] == "new_token"
        # The old refresh_token should be carried forward since Cognito
        # didn't include one in the refresh response
        assert result["refresh_token"] == "refresh_abc"

    def test_returns_cached_if_not_expired(self, tmp_aab_dir):
        """Returns the in-memory cached token without reading disk when it's still valid."""
        from aablocks.cached_token import set_cached_token

        # Pre-populate the in-memory cache — _load_token should return this
        # immediately without touching the filesystem
        token = {
            "id_token": "cached_token",
            "expires_at": int(time()) + 3600,
        }
        set_cached_token(DOMAIN, token)
        result = _load_token()
        assert result == token


class TestSaveToken:
    def setup_method(self):
        """Start each test with a clean in-memory cache."""
        clear_cached_token(DOMAIN)

    def test_computes_expires_at_from_expires_in(self, tmp_aab_dir):
        """Converts relative expires_in to an absolute expires_at timestamp."""
        # Cognito returns expires_in (seconds until expiry). _save_token
        # converts this to an absolute expires_at for easier comparison later.
        token = {"id_token": "abc", "expires_in": 3600}
        before = int(time())
        with patch("aablocks.auth.log_auth"):
            _save_token(token)
        after = int(time())
        assert "expires_at" in token
        assert before + 3600 <= token["expires_at"] <= after + 3600

    def test_persists_to_disk(self, tmp_aab_dir):
        """Writes the token to token.json on disk, keyed by auth domain."""
        token = {"id_token": "abc", "expires_at": 9999999999}
        _save_token(token)
        file_path = tmp_aab_dir / "token.json"
        assert file_path.exists()
        data = json.loads(file_path.read_text())
        assert DOMAIN in data
        assert data[DOMAIN]["id_token"] == "abc"

    def test_caches_in_memory(self, tmp_aab_dir):
        """Stores the token in the in-memory cache after saving."""
        token = {"id_token": "abc", "expires_at": 9999999999}
        _save_token(token)
        assert get_cached_token(DOMAIN)["id_token"] == "abc"


# ---------------------------------------------------------------------------
# _refresh() sends a grant_type=refresh_token POST to the Cognito token
# endpoint. It returns the new token on success, or None on any failure
# (missing refresh_token, HTTP error, network error). Returning None signals
# the caller to trigger a fresh login flow.
# ---------------------------------------------------------------------------
class TestRefresh:
    def setup_method(self):
        """Start each test with a clean in-memory cache."""
        clear_cached_token(DOMAIN)

    def test_no_refresh_token_returns_none(self, tmp_aab_dir):
        """Returns None when the old token has no refresh_token field."""
        # Without a refresh_token there's nothing to send to Cognito
        with patch("aablocks.auth.log_auth"):
            result = _refresh({"id_token": "old"})
        assert result is None

    def test_successful_refresh(self, tmp_aab_dir):
        """A successful refresh returns a new token with the refresh_token carried forward."""
        old_token = {"id_token": "old", "refresh_token": "rt123"}
        new_token_data = {
            "id_token": "new",
            "access_token": "new_access",
            "expires_in": 3600,
        }

        mock_resp = MagicMock()
        mock_resp.ok = True
        mock_resp.json.return_value = new_token_data

        with (
            patch("aablocks.auth.requests.post", return_value=mock_resp),
            patch("aablocks.auth.log_auth"),
        ):
            result = _refresh(old_token)

        assert result is not None
        assert result["id_token"] == "new"
        # Cognito may omit refresh_token on refresh — _refresh carries the old one forward
        assert result["refresh_token"] == "rt123"
        # expires_at should have been computed from expires_in by _save_token
        assert "expires_at" in result

    def test_failed_refresh_returns_none(self, tmp_aab_dir):
        """Returns None when Cognito rejects the refresh request (e.g. 400)."""
        old_token = {"id_token": "old", "refresh_token": "rt123"}

        # Cognito returns 400 when the refresh token is expired or revoked
        mock_resp = MagicMock()
        mock_resp.ok = False
        mock_resp.status_code = 400

        with (
            patch("aablocks.auth.requests.post", return_value=mock_resp),
            patch("aablocks.auth.log_auth"),
        ):
            result = _refresh(old_token)

        assert result is None

    def test_network_error_returns_none(self, tmp_aab_dir):
        """Returns None when a network error prevents the refresh request."""
        old_token = {"id_token": "old", "refresh_token": "rt123"}

        # Simulate being offline — the caller should fall back to a fresh login
        with (
            patch(
                "aablocks.auth.requests.post",
                side_effect=requests.ConnectionError("offline"),
            ),
            patch("aablocks.auth.log_auth"),
        ):
            result = _refresh(old_token)

        assert result is None


# ---------------------------------------------------------------------------
# get_token / get_token_header
#
# get_token() is the main entry point for obtaining a valid token. It tries
# _load_token() first (cache -> disk -> refresh), and falls back to login()
# if nothing is available. get_token_header() wraps the result into the
# headers dict that API requests need.
# ---------------------------------------------------------------------------
class TestGetToken:
    def setup_method(self):
        """Start each test with a clean in-memory cache."""
        clear_cached_token(DOMAIN)

    def test_returns_loaded_token(self, tmp_aab_dir):
        """Returns the token directly when _load_token finds a valid one."""
        token = {"id_token": "abc", "expires_at": int(time()) + 3600}
        with patch("aablocks.auth._load_token", return_value=token):
            result = get_token()
        assert result["id_token"] == "abc"

    def test_triggers_login_when_no_token(self, tmp_aab_dir):
        """Triggers the login flow when no existing token is found."""
        token = {"id_token": "new_from_login", "expires_at": int(time()) + 3600}

        def fake_login(**kwargs):
            from aablocks.cached_token import set_cached_token

            set_cached_token(DOMAIN, token)  # Simulate _exchange_code_for_token

        with (
            patch("aablocks.auth._load_token", return_value=None),
            patch("aablocks.auth.login", side_effect=fake_login),
        ):
            result = get_token()
        assert result["id_token"] == "new_from_login"

    def test_raises_if_login_fails_to_cache(self, tmp_aab_dir):
        """Raises RuntimeError if login completes but doesn't cache a token."""
        # This guards against a bug in the login flow where the browser
        # callback succeeds but the token exchange silently fails
        with (
            patch("aablocks.auth._load_token", return_value=None),
            patch("aablocks.auth.login"),  # no-op login that doesn't cache anything
        ):
            with pytest.raises(RuntimeError, match="Login failed"):
                get_token()


class TestGetTokenHeader:
    def setup_method(self):
        """Start each test with a clean in-memory cache."""
        clear_cached_token(DOMAIN)

    def test_returns_bearer_header(self, tmp_aab_dir):
        """Returns Authorization, X-Client-Type, and X-Client-Version headers."""
        token = {"id_token": "my_id_token", "expires_at": int(time()) + 3600}
        with patch("aablocks.auth._load_token", return_value=token):
            headers = get_token_header()
        # The API gateway validates the Bearer token against the Cognito user pool
        assert headers["Authorization"] == "Bearer my_id_token"
        assert "X-Client-Type" in headers
        assert "X-Client-Version" in headers

    def test_raises_if_no_id_token(self, tmp_aab_dir):
        """Raises RuntimeError when the token dict has no id_token field."""
        # This could happen if the token response was malformed
        token = {"access_token": "only_access", "expires_at": int(time()) + 3600}
        with patch("aablocks.auth._load_token", return_value=token):
            with pytest.raises(RuntimeError, match="No id_token"):
                get_token_header()


# ---------------------------------------------------------------------------
# login / logout
#
# login() is the full OAuth flow: PKCE generation -> browser -> callback ->
# code exchange -> save. We only test the short-circuit path here (already
# logged in) because the full browser flow would require mocking the HTTP
# server, webbrowser.open, and the callback handler.
#
# logout() clears both the in-memory cache and the disk token file.
# ---------------------------------------------------------------------------
class TestLogout:
    def setup_method(self):
        """Start each test with a clean in-memory cache."""
        clear_cached_token(DOMAIN)

    def test_clears_cache_and_file(self, write_token, tmp_aab_dir):
        """Clears the in-memory cache and deletes the token file from disk."""
        from aablocks.cached_token import set_cached_token

        # Set up state: token in both memory and on disk
        set_cached_token(DOMAIN, {"id_token": "abc"})
        write_token({"id_token": "abc"})
        token_path = tmp_aab_dir / "token.json"
        assert token_path.exists()

        with patch("aablocks.auth.log_auth"):
            logout()

        # Both storage locations should be cleared
        assert get_cached_token(DOMAIN) is None
        # File should be deleted since no other domains remain
        assert not token_path.exists()

    def test_no_file_prints_no_session(self, tmp_aab_dir, capsys):
        """Prints 'No active session' when there is no token file to delete."""
        logout()
        captured = capsys.readouterr()
        assert "No active session" in captured.out


class TestLogin:
    def setup_method(self):
        """Start each test with a clean in-memory cache."""
        clear_cached_token(DOMAIN)

    def test_noop_when_already_logged_in(self, tmp_aab_dir):
        """Skips the login flow entirely when a valid token already exists."""
        # login() calls _load_token() first — if it returns a token, the
        # browser flow is skipped entirely. This test would fail if login()
        # tried to open a browser or start the callback server.
        token = {"id_token": "valid", "expires_at": int(time()) + 3600}
        with patch("aablocks.auth._load_token", return_value=token):
            login()


class TestManualRedirectUri:
    def test_defaults_to_prod_web_app(self, tmp_aab_dir):
        """The manual flow redirects to the web app's /cli-login page."""
        assert _manual_redirect_uri() == "https://atlas.aalphabio.com/cli-login"

    def test_uses_configured_web_url(self, tmp_aab_dir):
        """A custom web_url config points the redirect at that host."""
        set_config("web_url", "https://data-test.aalpha-internal.com/")
        # Trailing slash is normalized so we don't get a doubled "//cli-login".
        assert (
            _manual_redirect_uri() == "https://data-test.aalpha-internal.com/cli-login"
        )


class TestGatedLoginUrl:
    def test_wraps_authorize_url_in_cli_login_gate(self, tmp_aab_dir):
        """The authorize URL is forwarded as a single percent-encoded ``auth`` param."""
        authorize_url = (
            "https://auth.atlas.aalphabio.com/login?response_type=code"
            "&client_id=abc&redirect_uri=http%3A%2F%2Flocalhost%3A6565%2Fcallback"
            "&state=xyz&code_challenge=cc&code_challenge_method=S256"
        )
        gated = _gated_login_url(authorize_url)

        parsed = urlparse(gated)
        assert parsed.scheme == "https"
        assert parsed.netloc == "atlas.aalphabio.com"
        assert parsed.path == "/cli-login"
        # The whole authorize URL (including its own query) round-trips intact.
        assert parse_qs(parsed.query)["auth"] == [authorize_url]

    def test_uses_configured_web_url(self, tmp_aab_dir):
        """A custom web_url config points the gate at that host."""
        set_config("web_url", "https://data-test.aalpha-internal.com/")
        gated = _gated_login_url("https://auth.example.com/login?state=s")
        assert gated.startswith("https://data-test.aalpha-internal.com/cli-login?auth=")


class TestDecodePasteToken:
    def test_decodes_base64url_token(self):
        """The base64url JSON blob from /cli-login round-trips to (code, state)."""
        blob = json.dumps({"code": "the-code", "state": "the-state"})
        token = base64.urlsafe_b64encode(blob.encode()).decode().rstrip("=")
        assert _decode_paste_token(token) == ("the-code", "the-state")

    def test_decodes_token_with_url_unsafe_bytes(self):
        """Padding-free urlsafe decoding handles values that need - and _."""
        # Force bytes that base64-encode with + and / in standard alphabet so
        # the urlsafe round-trip (which uses - and _) is actually exercised.
        blob = json.dumps({"code": "a?b>c~code", "state": "s/t+ate"})
        token = base64.urlsafe_b64encode(blob.encode()).decode().rstrip("=")
        assert _decode_paste_token(token) == ("a?b>c~code", "s/t+ate")

    def test_accepts_full_redirect_url(self):
        """A pasted redirect URL with a query string is also accepted."""
        url = "https://atlas.aalphabio.com/cli-login?code=xyz&state=abc"
        assert _decode_paste_token(url) == ("xyz", "abc")

    def test_rejects_garbage(self):
        """Undecodable input raises a clear error rather than crashing."""
        with pytest.raises(RuntimeError, match="Could not decode the login code"):
            _decode_paste_token("not a real token")

    def test_rejects_url_without_code(self):
        """A redirect URL missing the code is rejected."""
        with pytest.raises(RuntimeError, match="Could not find an authorization code"):
            _decode_paste_token("https://atlas.aalphabio.com/cli-login?state=abc")
