"""Unit tests for DuckDB connector utilities."""

import csv
import os
import tempfile
from unittest.mock import MagicMock, Mock, patch

import duckdb
import pyarrow as pa
import pyarrow.parquet as pq
import pytest

from src.connectors.duckdb import duckdb_connection
from src.connectors.duckdb.utils import get_file_columns
from src.enums import FileType
from src.errors import InvalidFileTypeError
from src.schemas import FileMetadata


class TestDuckDBUtilsGetFileColumns:
    """Tests for get_file_columns utility functions."""

    @pytest.fixture
    def temp_csv(self):
        """Create temporary CSV file."""
        with tempfile.NamedTemporaryFile(
            mode='w', suffix='.csv', delete=False, encoding='utf-8'
        ) as f:
            writer = csv.writer(f)
            writer.writerow(['col1', 'col2', 'col3'])
            writer.writerow(['value1', 'value2', 'value3'])
            path = f.name
        yield path
        if os.path.exists(path):
            os.remove(path)

    @pytest.fixture
    def temp_parquet(self):
        """Create temporary Parquet file."""
        with tempfile.NamedTemporaryFile(
            mode='wb', suffix='.parquet', delete=False
        ) as f:
            path = f.name

        table = pa.table({'col1': ['value1', 'value2'], 'col2': ['value3', 'value4']})
        pq.write_table(table, path)
        yield path
        if os.path.exists(path):
            os.remove(path)

    def test_get_file_columns_csv(self, temp_csv):
        """Test get_file_columns with CSV file."""
        conn = duckdb.connect(':memory:')
        metadata = FileMetadata(
            file_path=temp_csv,
            file_type=FileType.CSV,
            encoding='utf-8',
            gzipped=False,
        )

        columns = get_file_columns(conn, metadata)

        assert columns == {'col1', 'col2', 'col3'}

    def test_get_file_columns_parquet(self, temp_parquet):
        """Test get_file_columns with Parquet file."""
        conn = duckdb.connect(':memory:')
        metadata = FileMetadata(
            file_path=temp_parquet,
            file_type=FileType.PQT,
            encoding='utf-8',
            gzipped=False,
        )

        columns = get_file_columns(conn, metadata)

        assert columns == {'col1', 'col2'}

    def test_get_file_columns_invalid_type_raises_error(self):
        """Test get_file_columns with invalid file type raises error."""
        conn = duckdb.connect(':memory:')
        # Create metadata with an invalid file type using Mock
        metadata = Mock(spec=FileMetadata)
        metadata.file_path = '/fake/path.txt'
        metadata.file_type = 'invalid'
        metadata.encoding = 'utf-8'
        metadata.gzipped = False

        with pytest.raises(InvalidFileTypeError, match='Unsupported file type'):
            get_file_columns(conn, metadata)


class TestDuckDBConnection:
    """Tests for duckdb_connection context manager."""

    @patch('src.connectors.duckdb.connection.os.makedirs')
    @patch('src.connectors.duckdb.connection.duckdb.connect')
    @patch('src.connectors.duckdb.connection.get_free_disk_space')
    def test_duckdb_connection_yields_connection(
        self, mock_get_free_disk_space, mock_connect, mock_makedirs
    ):
        """Test duckdb_connection yields a connection object."""
        from src.connectors.duckdb import DuckDBConfig

        mock_conn = MagicMock(spec=duckdb.DuckDBPyConnection)
        mock_connect.return_value = mock_conn
        mock_get_free_disk_space.return_value = 1024 * 1024 * 1024  # 1GB

        config = DuckDBConfig()
        with duckdb_connection(config) as conn:
            assert conn is not None
            assert conn == mock_conn

        mock_conn.close.assert_called_once()

    @patch('src.connectors.duckdb.connection.os.makedirs')
    @patch('src.connectors.duckdb.connection.duckdb.connect')
    @patch('src.connectors.duckdb.connection.get_free_disk_space')
    def test_duckdb_connection_with_memory_limit(
        self, mock_get_free_disk_space, mock_connect, mock_makedirs
    ):
        """Test duckdb_connection with memory limits."""
        from src.connectors.duckdb import DuckDBConfig

        mock_conn = MagicMock(spec=duckdb.DuckDBPyConnection)
        mock_connect.return_value = mock_conn
        mock_get_free_disk_space.return_value = 1024 * 1024 * 1024  # 1GB

        config = DuckDBConfig(max_memory_mb=512)
        with duckdb_connection(config) as conn:
            assert conn is not None
            assert conn == mock_conn

        # Verify connect was called with memory limit
        call_args = mock_connect.call_args
        assert 'config' in call_args[1]
        assert '512MB' in call_args[1]['config']['max_memory']
        mock_conn.close.assert_called_once()

    @patch('src.connectors.duckdb.connection.os.makedirs')
    @patch('src.connectors.duckdb.connection.duckdb.connect')
    @patch('src.connectors.duckdb.connection.get_free_disk_space')
    def test_duckdb_connection_closes_on_error(
        self, mock_get_free_disk_space, mock_connect, mock_makedirs
    ):
        """Test duckdb_connection closes connection on error."""
        from src.connectors.duckdb import DuckDBConfig

        mock_conn = MagicMock(spec=duckdb.DuckDBPyConnection)
        mock_conn.execute.return_value = None  # Mock the extension loading
        mock_connect.return_value = mock_conn
        mock_get_free_disk_space.return_value = 1024 * 1024 * 1024  # 1GB

        config = DuckDBConfig()
        # Test that close is called even when an error occurs
        with pytest.raises(RuntimeError, match='Test error'):
            with duckdb_connection(config):
                # Simulate an error occurring during usage
                raise RuntimeError('Test error')

        # Verify close was called in the finally block
        mock_conn.close.assert_called_once()