"""Tests for src/connectors/snowflake/utils.py - Snowflake utility functions."""

import tempfile
from pathlib import Path
from unittest.mock import mock_open, patch

import pytest
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa

from src.connectors.snowflake.utils import (
    get_max_query_bytes,
    get_private_key,
    get_private_key_from_file,
    get_private_key_from_str,
)


class TestGetMaxQueryBytes:
    """Test get_max_query_bytes function."""

    def test_get_max_query_bytes_returns_16mb(self):
        """Test returns 16 MB (16777216 bytes)."""
        result = get_max_query_bytes()
        assert result == 16 * 1024 * 1024
        assert result == 16777216


class TestGetPrivateKey:
    """Test get_private_key function."""

    @patch('src.connectors.snowflake.utils.get_private_key_from_str')
    def test_get_private_key_from_string(self, mock_from_str):
        """Test loading private key from string when provided."""
        mock_from_str.return_value = b'key_bytes'

        result = get_private_key(
            private_key='key_content',
            private_key_path=None,
            private_key_passphrase=None,
        )

        assert result == b'key_bytes'
        mock_from_str.assert_called_once_with('key_content', None)

    @patch('src.connectors.snowflake.utils.get_private_key_from_str')
    @patch('src.connectors.snowflake.utils.get_private_key_from_file')
    def test_get_private_key_prefers_string_over_file(
        self, mock_from_file, mock_from_str
    ):
        """Test string key is preferred when both string and file are provided."""
        mock_from_str.return_value = b'string_key'

        result = get_private_key(
            private_key='key_content',
            private_key_path='/path/to/key',
            private_key_passphrase='pass',
        )

        assert result == b'string_key'
        mock_from_str.assert_called_once_with('key_content', 'pass')
        mock_from_file.assert_not_called()

    @patch('src.connectors.snowflake.utils.get_private_key_from_file')
    def test_get_private_key_from_file_path(self, mock_from_file):
        """Test loading private key from file when string not provided."""
        mock_from_file.return_value = b'file_key'

        result = get_private_key(
            private_key=None,
            private_key_path='/path/to/key.pem',
            private_key_passphrase='secret',
        )

        assert result == b'file_key'
        mock_from_file.assert_called_once_with('/path/to/key.pem', 'secret')

    def test_get_private_key_returns_none_when_no_source(self):
        """Test returns None when neither string nor file provided."""
        result = get_private_key(
            private_key=None, private_key_path=None, private_key_passphrase=None
        )

        assert result is None


class TestGetPrivateKeyFromFile:
    """Test get_private_key_from_file function."""

    def _generate_test_key(self, passphrase: str | None = None) -> bytes:
        """Generate a test RSA private key in PEM format."""
        private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)

        encryption = serialization.NoEncryption()
        if passphrase:
            encryption = serialization.BestAvailableEncryption(
                passphrase.encode('utf-8')
            )

        return private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=encryption,
        )

    def test_get_private_key_from_file_unencrypted(self):
        """Test loading unencrypted private key from file."""
        key_pem = self._generate_test_key()

        with tempfile.NamedTemporaryFile(mode='wb', delete=False, suffix='.pem') as f:
            f.write(key_pem)
            temp_path = f.name

        try:
            result = get_private_key_from_file(temp_path)

            # Verify it returns DER format bytes
            assert isinstance(result, bytes)
            assert len(result) > 0

            # Verify it can be loaded as a private key in DER format
            serialization.load_der_private_key(result, password=None)
        finally:
            Path(temp_path).unlink()

    def test_get_private_key_from_file_with_passphrase(self):
        """Test loading encrypted private key from file with passphrase."""
        passphrase = 'test_password_123'
        key_pem = self._generate_test_key(passphrase)

        with tempfile.NamedTemporaryFile(mode='wb', delete=False, suffix='.pem') as f:
            f.write(key_pem)
            temp_path = f.name

        try:
            result = get_private_key_from_file(temp_path, passphrase)

            assert isinstance(result, bytes)
            assert len(result) > 0

            # Verify the DER key can be loaded
            serialization.load_der_private_key(result, password=None)
        finally:
            Path(temp_path).unlink()

    def test_get_private_key_from_file_expands_tilde(self):
        """Test that tilde in path is expanded."""
        with patch('os.path.expanduser') as mock_expand:
            mock_expand.return_value = '/home/user/.ssh/key.pem'

            with patch('builtins.open', mock_open(read_data=self._generate_test_key())):
                get_private_key_from_file('~/.ssh/key.pem')

                mock_expand.assert_called_once_with('~/.ssh/key.pem')

    def test_get_private_key_from_file_wrong_passphrase(self):
        """Test loading encrypted key with wrong passphrase raises error."""
        passphrase = 'correct_password'
        key_pem = self._generate_test_key(passphrase)

        with tempfile.NamedTemporaryFile(mode='wb', delete=False, suffix='.pem') as f:
            f.write(key_pem)
            temp_path = f.name

        try:
            with pytest.raises(ValueError):
                get_private_key_from_file(temp_path, 'wrong_password')
        finally:
            Path(temp_path).unlink()

    def test_get_private_key_from_file_missing_passphrase(self):
        """Test loading encrypted key without passphrase raises error."""
        passphrase = 'required_password'
        key_pem = self._generate_test_key(passphrase)

        with tempfile.NamedTemporaryFile(mode='wb', delete=False, suffix='.pem') as f:
            f.write(key_pem)
            temp_path = f.name

        try:
            with pytest.raises(TypeError):
                get_private_key_from_file(temp_path, None)
        finally:
            Path(temp_path).unlink()


class TestGetPrivateKeyFromStr:
    """Test get_private_key_from_str function."""

    def _generate_test_key(self, passphrase: str | None = None) -> str:
        """Generate a test RSA private key in PEM format as string."""
        private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)

        encryption = serialization.NoEncryption()
        if passphrase:
            encryption = serialization.BestAvailableEncryption(
                passphrase.encode('utf-8')
            )

        key_bytes = private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=encryption,
        )
        return key_bytes.decode('utf-8')

    def test_get_private_key_from_str_unencrypted(self):
        """Test loading unencrypted private key from string."""
        key_pem = self._generate_test_key()

        result = get_private_key_from_str(key_pem)

        # Verify it returns DER format bytes
        assert isinstance(result, bytes)
        assert len(result) > 0

        # Verify it can be loaded as a private key in DER format
        serialization.load_der_private_key(result, password=None)

    def test_get_private_key_from_str_with_passphrase(self):
        """Test loading encrypted private key from string with passphrase."""
        passphrase = 'secure_pass_456'
        key_pem = self._generate_test_key(passphrase)

        result = get_private_key_from_str(key_pem, passphrase)

        assert isinstance(result, bytes)
        assert len(result) > 0

        # Verify the DER key can be loaded
        serialization.load_der_private_key(result, password=None)

    def test_get_private_key_from_str_wrong_passphrase(self):
        """Test loading encrypted key with wrong passphrase raises error."""
        passphrase = 'correct_pass'
        key_pem = self._generate_test_key(passphrase)

        with pytest.raises(ValueError):
            get_private_key_from_str(key_pem, 'wrong_pass')

    def test_get_private_key_from_str_missing_passphrase(self):
        """Test loading encrypted key without passphrase raises error."""
        passphrase = 'needed_pass'
        key_pem = self._generate_test_key(passphrase)

        with pytest.raises(TypeError):
            get_private_key_from_str(key_pem, None)

    def test_get_private_key_from_str_invalid_key_format(self):
        """Test loading invalid key format raises error."""
        invalid_key = 'not a valid private key'

        with pytest.raises(ValueError):
            get_private_key_from_str(invalid_key)