"""Tests for lockstock_core.rotate — Account Key rotation."""

import base64
import pytest
import respx
import httpx

from lockstock_core.keyring import (
    generate_account_key,
    generate_dek,
    derive_kek,
    wrap_dek,
    unwrap_dek,
)
from lockstock_core.rotate import (
    fetch_operator_keys,
    rotate_account_key,
    main,
)

SERVER = "https://test.lockstock.example.com"
API_KEY = "lsk_test_key_rotate"


def _make_agent_entry(agent_id, account_key, key_version=1):
    """Create a realistic agent entry with properly wrapped DEK."""
    kek = derive_kek(account_key, agent_id)
    dek = generate_dek()
    wrapped = wrap_dek(kek, dek)
    return {
        "agent_id": agent_id,
        "key_version": key_version,
        "wrapped_dek_b64": base64.b64encode(wrapped).decode("ascii"),
        "created_at": "2026-02-15T23:00:00Z",
    }, dek


class TestFetchOperatorKeys:
    """Test operator key listing."""

    @respx.mock
    def test_fetch_operator_keys(self):
        response_data = {
            "operator_id": "user-123",
            "agents": [{"agent_id": "agent_a", "key_version": 1}],
            "total": 1,
        }
        respx.get(f"{SERVER}/v1/operators/keys").mock(
            return_value=httpx.Response(200, json=response_data)
        )
        result = fetch_operator_keys(SERVER, API_KEY)
        assert result["operator_id"] == "user-123"
        assert result["total"] == 1

    @respx.mock
    def test_fetch_operator_keys_unauthorized(self):
        respx.get(f"{SERVER}/v1/operators/keys").mock(
            return_value=httpx.Response(401, json={"error": "unauthorized"})
        )
        with pytest.raises(httpx.HTTPStatusError):
            fetch_operator_keys(SERVER, API_KEY)


class TestRotateAccountKey:
    """Test the full rotation flow."""

    @respx.mock
    def test_rotate_single_agent(self):
        old_key = generate_account_key()
        new_key = generate_account_key()
        agent_id = "agent_test_rotate_1"

        entry, original_dek = _make_agent_entry(agent_id, old_key)

        # Mock operator keys listing
        respx.get(f"{SERVER}/v1/operators/keys").mock(
            return_value=httpx.Response(200, json={
                "operator_id": "user-123",
                "agents": [entry],
                "total": 1,
            })
        )

        # Mock PUT for re-wrapped DEK
        put_calls = []
        def capture_put(request):
            put_calls.append(request)
            return httpx.Response(200, json={"stored": True})

        respx.put(f"{SERVER}/v1/agents/{agent_id}/keys/1").mock(side_effect=capture_put)

        results = rotate_account_key(
            old_account_key=old_key,
            new_account_key=new_key,
            api_key=API_KEY,
            server_url=SERVER,
        )

        assert len(results) == 1
        assert results[0]["status"] == "rotated"
        assert results[0]["agent_id"] == agent_id

        # Verify the PUT was called
        assert len(put_calls) == 1

        # Verify the new wrapped DEK can be unwrapped with new key
        put_body = put_calls[0].content
        import json
        new_wrapped_b64 = json.loads(put_body)["wrapped_dek_b64"]
        new_wrapped = base64.b64decode(new_wrapped_b64)
        new_kek = derive_kek(new_key, agent_id)
        recovered_dek = unwrap_dek(new_kek, new_wrapped)
        assert recovered_dek == original_dek

    @respx.mock
    def test_rotate_multiple_agents(self):
        old_key = generate_account_key()
        new_key = generate_account_key()

        entry_a, dek_a = _make_agent_entry("agent_a", old_key)
        entry_b, dek_b = _make_agent_entry("agent_b", old_key)

        respx.get(f"{SERVER}/v1/operators/keys").mock(
            return_value=httpx.Response(200, json={
                "operator_id": "user-123",
                "agents": [entry_a, entry_b],
                "total": 2,
            })
        )
        respx.put(f"{SERVER}/v1/agents/agent_a/keys/1").mock(
            return_value=httpx.Response(200, json={"stored": True})
        )
        respx.put(f"{SERVER}/v1/agents/agent_b/keys/1").mock(
            return_value=httpx.Response(200, json={"stored": True})
        )

        results = rotate_account_key(
            old_account_key=old_key,
            new_account_key=new_key,
            api_key=API_KEY,
            server_url=SERVER,
        )

        assert len(results) == 2
        assert all(r["status"] == "rotated" for r in results)

    @respx.mock
    def test_dry_run_does_not_write(self):
        old_key = generate_account_key()
        new_key = generate_account_key()

        entry, _ = _make_agent_entry("agent_dry", old_key)

        respx.get(f"{SERVER}/v1/operators/keys").mock(
            return_value=httpx.Response(200, json={
                "operator_id": "user-123",
                "agents": [entry],
                "total": 1,
            })
        )

        # No PUT mock — if rotation tries to PUT, it will fail
        results = rotate_account_key(
            old_account_key=old_key,
            new_account_key=new_key,
            api_key=API_KEY,
            server_url=SERVER,
            dry_run=True,
        )

        assert len(results) == 1
        assert results[0]["status"] == "dry_run"
        assert "old_wrapped_prefix" in results[0]
        assert "new_wrapped_prefix" in results[0]

    @respx.mock
    def test_wrong_old_key_fails_gracefully(self):
        real_key = generate_account_key()
        wrong_key = generate_account_key()
        new_key = generate_account_key()

        entry, _ = _make_agent_entry("agent_wrong", real_key)

        respx.get(f"{SERVER}/v1/operators/keys").mock(
            return_value=httpx.Response(200, json={
                "operator_id": "user-123",
                "agents": [entry],
                "total": 1,
            })
        )

        results = rotate_account_key(
            old_account_key=wrong_key,  # wrong!
            new_account_key=new_key,
            api_key=API_KEY,
            server_url=SERVER,
        )

        assert len(results) == 1
        assert results[0]["status"] == "failed"
        assert "error" in results[0]

    @respx.mock
    def test_empty_operator_no_agents(self):
        old_key = generate_account_key()
        new_key = generate_account_key()

        respx.get(f"{SERVER}/v1/operators/keys").mock(
            return_value=httpx.Response(200, json={
                "operator_id": "user-123",
                "agents": [],
                "total": 0,
            })
        )

        results = rotate_account_key(
            old_account_key=old_key,
            new_account_key=new_key,
            api_key=API_KEY,
            server_url=SERVER,
        )

        assert len(results) == 0


class TestCLI:
    """Test CLI argument parsing."""

    def test_generate_key(self, capsys):
        code = main(["--generate-key"])
        assert code == 0
        output = capsys.readouterr().out.strip()
        assert len(output) == 64  # 32 bytes = 64 hex chars
        # Verify it's valid hex
        bytes.fromhex(output)

    def test_same_keys_rejected(self, capsys):
        key = "aa" * 32
        code = main(["--old-key", key, "--new-key", key, "--api-key", "test"])
        assert code == 1
        assert "identical" in capsys.readouterr().err

    def test_invalid_key_length(self, capsys):
        code = main(["--old-key", "abcd", "--new-key", "bb" * 32, "--api-key", "test"])
        assert code == 1
        assert "32 bytes" in capsys.readouterr().err

    def test_missing_required_args(self, capsys):
        with pytest.raises(SystemExit):
            main(["--old-key", "aa" * 32])
