from unittest import mock

import pytest

from test_fixtures.models import ExecuteResult
from test_fixtures.snowflake import SnowflakeConnection

_PEM_KEY = '-----BEGIN PRIVATE KEY-----\nfakekey\n-----END PRIVATE KEY-----'
_CONNECT_KWARGS = {
    'account': 'sme-delphi',
    'user': 'TEST_USER',
    'warehouse': 'TEST_WH',
}


def _make_connection():
    # Patches exit after __init__, but conn._conn already holds the mock reference
    # so subsequent method calls on conn still hit the mock.
    with mock.patch('test_fixtures.snowflake.get_secret', return_value=_PEM_KEY):
        with mock.patch(
            'test_fixtures.snowflake.load_pem_private_key'
        ) as mock_load_key:
            mock_private_key = mock.Mock()
            mock_private_key.private_bytes.return_value = b'der-bytes'
            mock_load_key.return_value = mock_private_key
            with mock.patch('snowflake.connector.connect') as mock_connect:
                conn = SnowflakeConnection('qa/snowflake/key', **_CONNECT_KWARGS)
                return conn, mock_connect.return_value


def test_connects_with_der_key_and_kwargs():
    with mock.patch('test_fixtures.snowflake.get_secret', return_value=_PEM_KEY):
        with mock.patch(
            'test_fixtures.snowflake.load_pem_private_key'
        ) as mock_load_key:
            mock_private_key = mock.Mock()
            mock_private_key.private_bytes.return_value = b'der-bytes'
            mock_load_key.return_value = mock_private_key
            with mock.patch('snowflake.connector.connect') as mock_connect:
                SnowflakeConnection('qa/snowflake/key', **_CONNECT_KWARGS)
                mock_connect.assert_called_once_with(
                    **_CONNECT_KWARGS, private_key=b'der-bytes'
                )


def test_wraps_bare_key_in_pem_headers():
    bare_key = 'fakekey'
    with mock.patch('test_fixtures.snowflake.get_secret', return_value=bare_key):
        with mock.patch(
            'test_fixtures.snowflake.load_pem_private_key'
        ) as mock_load_key:
            mock_private_key = mock.Mock()
            mock_private_key.private_bytes.return_value = b'der-bytes'
            mock_load_key.return_value = mock_private_key
            with mock.patch('snowflake.connector.connect'):
                SnowflakeConnection('qa/snowflake/key')
            mock_load_key.assert_called_once_with(
                b'-----BEGIN PRIVATE KEY-----\nfakekey\n-----END PRIVATE KEY-----',
                password=None,
            )


def test_raises_if_secret_is_not_string():
    with mock.patch(
        'test_fixtures.snowflake.get_secret', return_value={'key': 'value'}
    ):
        with pytest.raises(ValueError, match='must be a plain string'):
            SnowflakeConnection('qa/snowflake/key')


def test_fetchall():
    conn, mock_sf_conn = _make_connection()
    mock_cursor = mock_sf_conn.cursor.return_value.__enter__.return_value
    mock_cursor.fetchall.return_value = [{'col': 'a'}, {'col': 'b'}]
    assert conn.fetchall('SELECT * FROM t') == [{'col': 'a'}, {'col': 'b'}]


def test_fetchone():
    conn, mock_sf_conn = _make_connection()
    mock_sf_conn.cursor.return_value.__enter__.return_value.fetchone.return_value = {
        'col': 'a'
    }
    assert conn.fetchone('SELECT * FROM t WHERE id = %s', (1,)) == {'col': 'a'}


def test_execute_returns_rowcount_and_rows_when_description_present():
    conn, mock_sf_conn = _make_connection()
    mock_cursor = mock_sf_conn.cursor.return_value.__enter__.return_value
    mock_cursor.rowcount = 1
    mock_cursor.description = [('id',)]
    mock_cursor.fetchall.return_value = [{'id': 42}]
    result = conn.execute('INSERT INTO t (x) VALUES (%s) RETURNING id', (1,))
    assert isinstance(result, ExecuteResult)
    assert result.rowcount == 1
    assert result.rows == [{'id': 42}]


def test_execute_rows_empty_when_no_result_set():
    conn, mock_sf_conn = _make_connection()
    mock_cursor = mock_sf_conn.cursor.return_value.__enter__.return_value
    mock_cursor.rowcount = 2
    mock_cursor.description = None
    result = conn.execute('DELETE FROM t WHERE id = %s', (1,))
    assert result.rowcount == 2
    assert result.rows == []
    mock_cursor.fetchall.assert_not_called()


def test_close():
    conn, mock_sf_conn = _make_connection()
    conn.close()
    mock_sf_conn.close.assert_called_once()


def test_context_manager_closes_on_exit():
    with mock.patch('test_fixtures.snowflake.get_secret', return_value=_PEM_KEY):
        with mock.patch(
            'test_fixtures.snowflake.load_pem_private_key'
        ) as mock_load_key:
            mock_private_key = mock.Mock()
            mock_private_key.private_bytes.return_value = b'der-bytes'
            mock_load_key.return_value = mock_private_key
            with mock.patch('snowflake.connector.connect') as mock_connect:
                with SnowflakeConnection('qa/snowflake/key', **_CONNECT_KWARGS):
                    pass
                mock_connect.return_value.close.assert_called_once()