"""Unit tests for DuckDBExecutorBase."""

import pytest
from unittest.mock import Mock, patch, MagicMock
import sys
import os

# Add parent directory to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from langchain_iceberg.tools.duckdb_base import DuckDBExecutorBase
from langchain_iceberg.exceptions import IcebergConnectionError


class TestDuckDBExecutorBase:
    """Test suite for DuckDBExecutorBase class."""

    def test_init_requires_duckdb(self, mock_catalog, catalog_config):
        """Test that initialization fails without DuckDB."""
        with patch('langchain_iceberg.tools.duckdb_base.HAS_DUCKDB', False):
            with pytest.raises(ImportError, match="DuckDB is required"):
                DuckDBExecutorBase(
                    catalog=mock_catalog,
                    catalog_config=catalog_config,
                )

    def test_init_success(self, mock_catalog, catalog_config):
        """Test successful initialization."""
        with patch('langchain_iceberg.tools.duckdb_base.duckdb') as mock_duckdb:
            mock_conn = Mock()
            mock_duckdb.connect.return_value = mock_conn
            
            base = DuckDBExecutorBase(
                catalog=mock_catalog,
                catalog_config=catalog_config,
            )
            
            assert base.catalog == mock_catalog
            assert base._catalog_config == catalog_config
            assert base._duckdb_conn == mock_conn

    def test_extract_table_references_simple(self, mock_catalog, catalog_config):
        """Test extracting table references from simple SQL."""
        with patch('langchain_iceberg.tools.duckdb_base.duckdb'):
            base = DuckDBExecutorBase(
                catalog=mock_catalog,
                catalog_config=catalog_config,
            )
            
            sql = "SELECT * FROM sales.orders WHERE status = 'completed'"
            tables = base._extract_table_references(sql)
            
            assert "sales.orders" in tables

    def test_extract_table_references_join(self, mock_catalog, catalog_config, sample_sql_query):
        """Test extracting table references from JOIN query."""
        with patch('langchain_iceberg.tools.duckdb_base.duckdb'):
            base = DuckDBExecutorBase(
                catalog=mock_catalog,
                catalog_config=catalog_config,
            )
            
            tables = base._extract_table_references(sample_sql_query)
            
            assert "sales.orders" in tables
            assert "sales.customers" in tables

    def test_get_table_location_success(self, mock_catalog, catalog_config, mock_table):
        """Test getting table location."""
        with patch('langchain_iceberg.tools.duckdb_base.duckdb'):
            base = DuckDBExecutorBase(
                catalog=mock_catalog,
                catalog_config=catalog_config,
            )
            
            # Mock catalog.load_table
            mock_catalog.load_table.return_value = mock_table
            
            table_root, metadata_file = base._get_table_location("sales.orders")
            
            assert "s3://test-bucket/warehouse/sales/orders" in table_root
            mock_catalog.load_table.assert_called_once()

    def test_get_table_location_invalid_format(self, mock_catalog, catalog_config):
        """Test error on invalid table_id format."""
        with patch('langchain_iceberg.tools.duckdb_base.duckdb'):
            base = DuckDBExecutorBase(
                catalog=mock_catalog,
                catalog_config=catalog_config,
            )
            
            with pytest.raises(ValueError, match="Invalid table_id format"):
                base._get_table_location("invalid_table")

    def test_register_table_success(self, mock_catalog, catalog_config, mock_table):
        """Test successful table registration."""
        with patch('langchain_iceberg.tools.duckdb_base.duckdb') as mock_duckdb:
            mock_conn = Mock()
            mock_duckdb.connect.return_value = mock_conn
            
            base = DuckDBExecutorBase(
                catalog=mock_catalog,
                catalog_config=catalog_config,
            )
            
            # Mock catalog.load_table
            mock_catalog.load_table.return_value = mock_table
            
            view_name = base._register_table("sales.orders")
            
            assert view_name == "sales_orders"
            assert "sales.orders" in base._registered_tables
            mock_conn.execute.assert_called()

    def test_register_table_idempotent(self, mock_catalog, catalog_config, mock_table):
        """Test that registering same table twice is idempotent."""
        with patch('langchain_iceberg.tools.duckdb_base.duckdb') as mock_duckdb:
            mock_conn = Mock()
            mock_duckdb.connect.return_value = mock_conn
            
            base = DuckDBExecutorBase(
                catalog=mock_catalog,
                catalog_config=catalog_config,
            )
            
            # Mock catalog.load_table
            mock_catalog.load_table.return_value = mock_table
            
            # Register first time
            view_name1 = base._register_table("sales.orders")
            call_count_1 = mock_conn.execute.call_count
            
            # Register second time (should skip)
            view_name2 = base._register_table("sales.orders")
            call_count_2 = mock_conn.execute.call_count
            
            assert view_name1 == view_name2
            assert call_count_2 == call_count_1  # No additional call

    def test_format_results(self, mock_catalog, catalog_config):
        """Test result formatting."""
        with patch('langchain_iceberg.tools.duckdb_base.duckdb'):
            base = DuckDBExecutorBase(
                catalog=mock_catalog,
                catalog_config=catalog_config,
            )
            
            # Mock DataFrame
            mock_df = Mock()
            mock_df.to_string = Mock(return_value="test results")
            
            with patch('langchain_iceberg.utils.formatters.ResultFormatter.format_table', 
                      return_value="formatted results"):
                result = base._format_results(
                    result_df=mock_df,
                    limit=100,
                    include_metadata=True,
                    execution_time_ms=150.5,
                    was_truncated=False
                )
                
                assert "formatted results" in result
                assert "150ms" in result

    def test_configure_cloud_storage_s3(self, mock_catalog, catalog_config):
        """Test S3 storage configuration."""
        with patch('langchain_iceberg.tools.duckdb_base.duckdb') as mock_duckdb:
            mock_conn = Mock()
            mock_duckdb.connect.return_value = mock_conn
            
            base = DuckDBExecutorBase(
                catalog=mock_catalog,
                catalog_config=catalog_config,
            )
            
            # Check that S3 secret creation was attempted
            calls = [str(call) for call in mock_conn.execute.call_args_list]
            s3_calls = [c for c in calls if 'CREATE SECRET' in c or 'iceberg_s3' in c]
            assert len(s3_calls) > 0


class TestDuckDBExecutorBaseEdgeCases:
    """Test edge cases and error handling."""

    def test_extract_tables_no_tables(self, mock_catalog, catalog_config):
        """Test extracting tables from query with no table references."""
        with patch('langchain_iceberg.tools.duckdb_base.duckdb'):
            base = DuckDBExecutorBase(
                catalog=mock_catalog,
                catalog_config=catalog_config,
            )
            
            sql = "SELECT 1 + 1"
            tables = base._extract_table_references(sql)
            
            assert len(tables) == 0

    def test_execute_sql_query_no_tables(self, mock_catalog, catalog_config):
        """Test error when no tables found in query."""
        with patch('langchain_iceberg.tools.duckdb_base.duckdb'):
            base = DuckDBExecutorBase(
                catalog=mock_catalog,
                catalog_config=catalog_config,
            )
            
            from langchain_iceberg.exceptions import IcebergInvalidQueryError
            
            with pytest.raises(IcebergInvalidQueryError, match="No table references found"):
                base._execute_sql_query("SELECT 1 + 1")

    def test_table_location_no_metadata(self, mock_catalog, catalog_config):
        """Test error when table has no location metadata."""
        with patch('langchain_iceberg.tools.duckdb_base.duckdb'):
            base = DuckDBExecutorBase(
                catalog=mock_catalog,
                catalog_config=catalog_config,
            )
            
            # Mock table with no metadata
            bad_table = Mock()
            bad_table.metadata = None
            bad_table.metadata_location = None
            bad_table.location = None
            
            mock_catalog.load_table.return_value = bad_table
            
            with pytest.raises(IcebergConnectionError, 
                             match="Cannot determine storage location"):
                base._get_table_location("sales.orders")
