"""Unit tests for AuthSession helper methods."""

import base64
import unittest
from datetime import timedelta
from unittest.mock import MagicMock, patch

from keylime.models.base.types import Timestamp
from keylime.models.verifier.auth_session import AuthSession
from keylime.shared_data import cleanup_global_shared_memory, get_shared_memory


class TestAuthSessionHelpers(unittest.TestCase):
    """Test cases for AuthSession helper methods."""

    def setUp(self):
        """Set up test fixtures."""
        # Clean up any existing shared memory
        cleanup_global_shared_memory()
        self.shared_memory = get_shared_memory()
        self.sessions_cache = self.shared_memory.get_or_create_dict("auth_sessions")
        self.test_agent_id = "test-agent-123"

    def tearDown(self):
        """Clean up after tests."""
        cleanup_global_shared_memory()

    def test_delete_stale_from_memory_removes_expired_nonce(self):
        """Test that delete_stale_from_memory removes sessions with expired nonces."""
        now = Timestamp.now()
        past_time = now - timedelta(seconds=10)

        # Create a session with expired nonce
        self.sessions_cache[1] = {  # type: ignore[index]
            "session_id": 1,
            "agent_id": self.test_agent_id,
            "nonce_expires_at": past_time,
            "token_expires_at": None,
        }

        # Call delete_stale_from_memory
        AuthSession.delete_stale_from_memory(self.test_agent_id)

        # Session should be removed
        self.assertNotIn(1, self.sessions_cache)

    def test_delete_stale_from_memory_removes_expired_token(self):
        """Test that delete_stale_from_memory removes sessions with expired tokens."""
        now = Timestamp.now()
        past_time = now - timedelta(seconds=10)

        # Create a session with expired token
        self.sessions_cache[2] = {  # type: ignore[index]
            "session_id": 2,
            "agent_id": self.test_agent_id,
            "nonce_expires_at": None,
            "token_expires_at": past_time,
        }

        # Call delete_stale_from_memory
        AuthSession.delete_stale_from_memory(self.test_agent_id)

        # Session should be removed
        self.assertNotIn(2, self.sessions_cache)

    def test_delete_stale_from_memory_keeps_valid_sessions(self):
        """Test that delete_stale_from_memory keeps valid sessions."""
        now = Timestamp.now()
        future_time = now + timedelta(seconds=60)

        # Create a valid session
        self.sessions_cache[3] = {  # type: ignore[index]
            "session_id": 3,
            "agent_id": self.test_agent_id,
            "nonce_expires_at": future_time,
            "token_expires_at": future_time,
        }

        # Call delete_stale_from_memory
        AuthSession.delete_stale_from_memory(self.test_agent_id)

        # Session should still be present
        self.assertIn(3, self.sessions_cache)

    def test_delete_stale_from_memory_only_affects_target_agent(self):
        """Test that delete_stale_from_memory only affects the specified agent."""
        now = Timestamp.now()
        past_time = now - timedelta(seconds=10)

        # Create sessions for different agents
        self.sessions_cache[4] = {  # type: ignore[index]
            "session_id": 4,
            "agent_id": self.test_agent_id,
            "nonce_expires_at": past_time,
        }
        self.sessions_cache[5] = {  # type: ignore[index]
            "session_id": 5,
            "agent_id": "other-agent-456",
            "nonce_expires_at": past_time,
        }

        # Call delete_stale_from_memory for test_agent_id
        AuthSession.delete_stale_from_memory(self.test_agent_id)

        # Only test_agent_id session should be removed
        self.assertNotIn(4, self.sessions_cache)
        self.assertIn(5, self.sessions_cache)

    def test_get_active_session_for_agent_from_memory(self):
        """Test that get_active_session_for_agent retrieves from shared memory."""
        now = Timestamp.now()
        future_time = now + timedelta(seconds=60)

        # Create an active session in shared memory
        session_data = {
            "session_id": 6,
            "agent_id": self.test_agent_id,
            "active": True,
            "token": "test-token-123",
            "token_expires_at": future_time,
        }
        self.sessions_cache[6] = session_data  # type: ignore[index]

        # Retrieve the session
        result = AuthSession.get_active_session_for_agent(self.test_agent_id)

        # Should return the session data
        self.assertIsNotNone(result)
        self.assertEqual(result["agent_id"], self.test_agent_id)  # type: ignore[index]
        self.assertEqual(result["token"], "test-token-123")  # type: ignore[index]

    @patch("keylime.models.verifier.auth_session.AuthSession.all")
    def test_get_active_session_for_agent_ignores_inactive(self, mock_all):
        """Test that get_active_session_for_agent ignores inactive sessions."""
        now = Timestamp.now()
        future_time = now + timedelta(seconds=60)

        # Mock empty database result
        mock_all.return_value = []

        # Create an inactive session
        self.sessions_cache[7] = {  # type: ignore[index]
            "session_id": 7,
            "agent_id": self.test_agent_id,
            "active": False,
            "token_expires_at": future_time,
        }

        # Retrieve the session
        result = AuthSession.get_active_session_for_agent(self.test_agent_id)

        # Should return None
        self.assertIsNone(result)

    @patch("keylime.models.verifier.auth_session.AuthSession.all")
    def test_get_active_session_for_agent_ignores_expired(self, mock_all):
        """Test that get_active_session_for_agent ignores expired sessions."""
        now = Timestamp.now()
        past_time = now - timedelta(seconds=10)

        # Mock empty database result
        mock_all.return_value = []

        # Create an expired session
        self.sessions_cache[8] = {  # type: ignore[index]
            "session_id": 8,
            "agent_id": self.test_agent_id,
            "active": True,
            "token_expires_at": past_time,
        }

        # Retrieve the session
        result = AuthSession.get_active_session_for_agent(self.test_agent_id)

        # Should return None
        self.assertIsNone(result)

    @patch("keylime.models.verifier.auth_session.AuthSession.all")
    def test_get_active_session_for_agent_from_database(self, mock_all):
        """Test that get_active_session_for_agent falls back to database."""
        now = Timestamp.now()
        future_time = now + timedelta(seconds=60)

        # Mock database session
        mock_session = MagicMock()
        mock_session.token = "db-token-456"
        mock_session.agent_id = self.test_agent_id
        mock_session.active = True
        mock_session.token_expires_at = future_time

        mock_all.return_value = [mock_session]

        # Retrieve the session (not in shared memory, should check DB)
        result = AuthSession.get_active_session_for_agent(self.test_agent_id)

        # Should return session data from database
        self.assertIsNotNone(result)
        self.assertEqual(result["agent_id"], self.test_agent_id)  # type: ignore[index]
        self.assertEqual(result["token"], "db-token-456")  # type: ignore[index]

        # Should also populate shared memory cache
        session_id = hash("db-token-456")
        self.assertIn(session_id, self.sessions_cache)

    def test_delete_active_session_for_agent(self):
        """Test that delete_active_session_for_agent removes the session."""
        # Create an active session
        self.sessions_cache[9] = {  # type: ignore[index]
            "session_id": 9,
            "agent_id": self.test_agent_id,
            "active": True,
        }

        # Delete the session
        AuthSession.delete_active_session_for_agent(self.test_agent_id)

        # Session should be removed
        self.assertNotIn(9, self.sessions_cache)

    def test_delete_active_session_for_agent_only_active(self):
        """Test that delete_active_session_for_agent only deletes active sessions."""
        # Create inactive and active sessions
        self.sessions_cache[10] = {  # type: ignore[index]
            "session_id": 10,
            "agent_id": self.test_agent_id,
            "active": False,
        }
        self.sessions_cache[11] = {  # type: ignore[index]
            "session_id": 11,
            "agent_id": self.test_agent_id,
            "active": True,
        }

        # Delete active session
        AuthSession.delete_active_session_for_agent(self.test_agent_id)

        # Only active session should be removed
        self.assertIn(10, self.sessions_cache)
        self.assertNotIn(11, self.sessions_cache)

    def test_delete_active_session_for_agent_only_target_agent(self):
        """Test that delete_active_session_for_agent only affects target agent."""
        # Create sessions for different agents
        self.sessions_cache[12] = {  # type: ignore[index]
            "session_id": 12,
            "agent_id": self.test_agent_id,
            "active": True,
        }
        self.sessions_cache[13] = {  # type: ignore[index]
            "session_id": 13,
            "agent_id": "other-agent-789",
            "active": True,
        }

        # Delete session for test_agent_id
        AuthSession.delete_active_session_for_agent(self.test_agent_id)

        # Only test_agent_id session should be removed
        self.assertNotIn(12, self.sessions_cache)
        self.assertIn(13, self.sessions_cache)


class TestAuthSessionCore(unittest.TestCase):
    """Test cases for core AuthSession functionality."""

    def setUp(self):
        """Set up test fixtures."""
        cleanup_global_shared_memory()
        self.test_agent_id = "test-agent-123"

    def tearDown(self):
        """Clean up after tests."""
        cleanup_global_shared_memory()

    def test_create_in_memory_success(self):
        """Test successful creation of in-memory session."""
        request_data = {
            "data": {
                "type": "session",
                "attributes": {
                    "agent_id": self.test_agent_id,
                    "authentication_supported": [
                        {
                            "authentication_class": "pop",
                            "authentication_type": "tpm_pop",
                            "capabilities": {
                                "supported_hash_algorithms": ["sha256", "sha384"],
                                "supported_signing_schemes": ["rsassa", "ecdsa"],
                            },
                        }
                    ],
                },
            }
        }

        result = AuthSession.create_in_memory(self.test_agent_id, request_data)

        # Should return session data dictionary (not errors)
        self.assertNotIn("errors", result)
        self.assertIn("session_id", result)
        self.assertIn("nonce", result)
        self.assertIn("response", result)
        self.assertEqual(result["agent_id"], self.test_agent_id)

        # Response should be JSON:API compliant
        response = result["response"]
        self.assertIn("data", response)
        self.assertEqual(response["data"]["type"], "session")

    def test_create_in_memory_validation_error(self):
        """Test that create_in_memory validates input."""
        # Missing tpm_pop in authentication_supported
        bad_request_data = {
            "data": {
                "type": "session",
                "attributes": {
                    "agent_id": self.test_agent_id,
                    "authentication_supported": [
                        {
                            "authentication_class": "pop",
                            "authentication_type": "other_type",  # Not tpm_pop
                        }
                    ],
                },
            }
        }

        result = AuthSession.create_in_memory(self.test_agent_id, bad_request_data)

        # Should return errors
        self.assertIn("errors", result)
        self.assertIn("authentication_supported", result["errors"])

    @patch("keylime.models.verifier.auth_session.get_session")
    @patch.object(AuthSession, "all")
    def test_authenticate_agent_success(self, mock_all, mock_get_session):
        """Test successful agent authentication with valid token."""
        # Create a mock agent
        mock_agent = MagicMock()
        mock_agent.agent_id = self.test_agent_id

        # Mock session query
        mock_db_session = MagicMock()
        mock_db_session.query.return_value.filter.return_value.one_or_none.return_value = mock_agent
        mock_get_session.return_value = mock_db_session

        # Mock AuthSession.all to return an active session
        mock_auth_session = MagicMock()
        mock_auth_session.token = "test-token"
        mock_auth_session.active = True
        mock_auth_session.agent_id = self.test_agent_id
        mock_all.return_value = [mock_auth_session]

        result = AuthSession.authenticate_agent("test-token")

        # Should return the agent
        self.assertIsNotNone(result)
        self.assertEqual(result.agent_id, self.test_agent_id)  # type: ignore[union-attr]

    @patch.object(AuthSession, "all")
    def test_authenticate_agent_inactive_session(self, mock_all):
        """Test that inactive sessions cannot authenticate."""
        # Mock AuthSession.all to return empty list (inactive sessions are filtered out)
        mock_all.return_value = []

        result = AuthSession.authenticate_agent("test-token")

        # Should return False
        self.assertFalse(result)

    @patch.object(AuthSession, "all")
    def test_authenticate_agent_no_session(self, mock_all):
        """Test that authentication fails when session doesn't exist."""
        # Mock AuthSession.all to return empty list (no sessions exist)
        mock_all.return_value = []

        result = AuthSession.authenticate_agent("test-token")

        # Should return False
        self.assertFalse(result)

    @patch.object(AuthSession, "empty")
    def test_create_with_agent(self, mock_empty):
        """Test AuthSession.create() with an enrolled agent."""
        # Mock agent
        mock_agent = MagicMock()
        mock_agent.agent_id = self.test_agent_id

        # Mock empty session
        mock_session = MagicMock()
        mock_session.initialise = MagicMock()
        mock_session.receive_capabilities = MagicMock()
        mock_empty.return_value = mock_session

        # Call create
        data = {"data": {"attributes": {"authentication_supported": [{"authentication_type": "tpm_pop"}]}}}
        AuthSession.create(mock_agent, data)

        # Verify initialization
        mock_session.initialise.assert_called_once_with(self.test_agent_id)
        mock_session.receive_capabilities.assert_called_once_with(data, mock_agent)

    @patch.object(AuthSession, "empty")
    def test_create_without_agent(self, mock_empty):
        """Test AuthSession.create() with unenrolled agent."""
        # Mock empty session
        mock_session = MagicMock()
        mock_session.initialise = MagicMock()
        mock_session.receive_capabilities = MagicMock()
        mock_empty.return_value = mock_session

        # Call create without agent (using agent_id)
        data = {"data": {"attributes": {"authentication_supported": [{"authentication_type": "tpm_pop"}]}}}
        AuthSession.create(None, data, agent_id=self.test_agent_id)

        # Verify initialization
        mock_session.initialise.assert_called_once_with(self.test_agent_id)
        mock_session.receive_capabilities.assert_called_once_with(data, None)

    @patch.object(AuthSession, "empty")
    def test_create_from_memory(self, mock_empty):
        """Test AuthSession.create_from_memory()."""
        # Mock agent
        mock_agent = MagicMock()
        mock_agent.agent_id = self.test_agent_id

        # Create session data
        now = Timestamp.now()
        session_data = {
            "token": "test-token",
            "agent_id": self.test_agent_id,
            "nonce": b"test-nonce",
            "nonce_created_at": now,
            "nonce_expires_at": now + timedelta(seconds=60),
            "hash_algorithm": "sha256",
            "signing_scheme": "rsassa",
        }

        # Mock empty session
        mock_session = MagicMock()
        mock_session.receive_pop = MagicMock()
        mock_empty.return_value = mock_session

        # Call create_from_memory
        pop_request = {"data": {"attributes": {"authentication_provided": []}}}
        AuthSession.create_from_memory(session_data, mock_agent, pop_request)

        # Verify session attributes were set
        self.assertEqual(mock_session.token, "test-token")
        self.assertEqual(mock_session.agent_id, self.test_agent_id)
        self.assertEqual(mock_session.nonce, b"test-nonce")
        self.assertFalse(mock_session.active)
        mock_session.receive_pop.assert_called_once_with(mock_agent, pop_request)


class TestAuthSessionReceiveCapabilities(unittest.TestCase):
    """Test cases for AuthSession.receive_capabilities()."""

    def setUp(self):
        """Set up test fixtures."""
        cleanup_global_shared_memory()
        self.test_agent_id = "test-agent-123"

    def tearDown(self):
        """Clean up after tests."""
        cleanup_global_shared_memory()

    def test_receive_capabilities_missing_tpm_pop(self):
        """Test receive_capabilities with missing tpm_pop."""
        # Create mock agent and session
        mock_agent = MagicMock()
        mock_agent.agent_id = self.test_agent_id

        mock_session = MagicMock()
        mock_session.nonce = None
        mock_session._add_error = MagicMock()  # pylint: disable=protected-access

        # Call receive_capabilities without tpm_pop
        data = {
            "data": {
                "attributes": {
                    "authentication_supported": [
                        {
                            "authentication_type": "other_type",
                        }
                    ]
                }
            }
        }

        # Call the actual method directly
        AuthSession.receive_capabilities(mock_session, data, None)

        # Should add error
        mock_session._add_error.assert_called_with(  # pylint: disable=protected-access
            "authentication_supported", "must include tpm_pop authentication type"
        )

    def test_receive_capabilities_already_received(self):
        """Test that receive_capabilities fails if nonce already set."""
        # Create mock session with nonce already set
        mock_session = MagicMock()
        mock_session.nonce = b"existing-nonce"

        # Call receive_capabilities should raise ValueError
        data = {"data": {"attributes": {"authentication_supported": [{"authentication_type": "tpm_pop"}]}}}

        with self.assertRaises(ValueError):
            AuthSession.receive_capabilities(mock_session, data, None)


class TestAuthSessionReceivePop(unittest.TestCase):
    """Test cases for AuthSession.receive_pop()."""

    def setUp(self):
        """Set up test fixtures."""
        cleanup_global_shared_memory()
        self.test_agent_id = "test-agent-123"

    def tearDown(self):
        """Clean up after tests."""
        cleanup_global_shared_memory()

    @patch("keylime.models.verifier.auth_session.Tpm")
    @patch("keylime.models.verifier.auth_session.config")
    def test_receive_pop_success(self, mock_config, mock_tpm):
        """Test successful receive_pop."""
        # Mock config
        mock_config.getint.return_value = 3600  # session lifetime

        # Mock agent with AK
        mock_agent = MagicMock()
        mock_agent.agent_id = self.test_agent_id
        mock_agent.ak_tpm = base64.b64encode(b"ak-tpm-data").decode("utf-8")

        # Mock session
        mock_session = MagicMock()
        mock_session.agent_id = self.test_agent_id
        mock_session.nonce = b"test-nonce"
        mock_session.hash_algorithm = "sha256"
        mock_session.signing_scheme = "rsassa"
        mock_session._add_error = MagicMock()  # pylint: disable=protected-access

        # Mock TPM verification to succeed
        mock_tpm.verify_tpm_object.return_value = None

        # Call receive_pop
        message = base64.b64encode(b"attest-data").decode("utf-8")
        signature = base64.b64encode(b"signature-data").decode("utf-8")
        data = {
            "data": {
                "attributes": {
                    "authentication_provided": [
                        {
                            "authentication_type": "tpm_pop",
                            "data": {
                                "message": message,
                                "signature": signature,
                            },
                        }
                    ]
                }
            }
        }

        AuthSession.receive_pop(mock_session, mock_agent, data)

        # Verify TPM verification was called
        mock_tpm.verify_tpm_object.assert_called_once()

    def test_receive_pop_wrong_agent(self):
        """Test receive_pop with wrong agent."""
        # Mock agent with different agent_id
        mock_agent = MagicMock()
        mock_agent.agent_id = "different-agent"
        mock_agent.ak_tpm = base64.b64encode(b"ak-tpm-data").decode("utf-8")

        # Mock session
        mock_session = MagicMock()
        mock_session.agent_id = self.test_agent_id
        mock_session._add_error = MagicMock()  # pylint: disable=protected-access

        # Call receive_pop
        data = {"data": {"attributes": {"authentication_provided": []}}}
        AuthSession.receive_pop(mock_session, mock_agent, data)

        # Should return early without adding errors (agent mismatch)
        # The function returns early if agent_id doesn't match

    def test_receive_pop_missing_proof(self):
        """Test receive_pop with missing proof data."""
        # Mock agent
        mock_agent = MagicMock()
        mock_agent.agent_id = self.test_agent_id
        mock_agent.ak_tpm = base64.b64encode(b"ak-tpm-data").decode("utf-8")

        # Mock session
        mock_session = MagicMock()
        mock_session.agent_id = self.test_agent_id
        mock_session._add_error = MagicMock()  # pylint: disable=protected-access

        # Call receive_pop without proof
        data = {"data": {"attributes": {"authentication_provided": []}}}
        AuthSession.receive_pop(mock_session, mock_agent, data)

        # Should add error
        mock_session._add_error.assert_called_with(  # pylint: disable=protected-access
            "authentication_provided", "must include at least one authentication method"
        )


class TestAuthSessionDeleteStale(unittest.TestCase):
    """Test cases for AuthSession.delete_stale()."""

    def setUp(self):
        """Set up test fixtures."""
        cleanup_global_shared_memory()
        self.test_agent_id = "test-agent-123"

    def tearDown(self):
        """Clean up after tests."""
        cleanup_global_shared_memory()

    @patch.object(AuthSession, "all")
    def test_delete_stale_removes_expired_sessions(self, mock_all):
        """Test that delete_stale removes expired sessions."""
        now = Timestamp.now()
        past_time = now - timedelta(seconds=10)

        # Mock expired sessions
        mock_session1 = MagicMock()
        mock_session1.nonce_expires_at = now + timedelta(seconds=60)  # Not expired
        mock_session1.token_expires_at = now + timedelta(seconds=3600)  # Not expired
        mock_session1.delete = MagicMock()

        mock_session2 = MagicMock()
        mock_session2.nonce_expires_at = past_time  # Expired
        mock_session2.token_expires_at = now + timedelta(seconds=3600)
        mock_session2.delete = MagicMock()

        mock_all.return_value = [mock_session1, mock_session2]

        # Call delete_stale
        AuthSession.delete_stale(self.test_agent_id)

        # Note: Current implementation has a bug - it deletes if NOT expired
        # This test documents the current behavior


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