"""Unit tests for DuckDBQueryTool (refactored)."""

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

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

from langchain_iceberg.tools.duckdb_tool_refactored import DuckDBQueryTool
from langchain_iceberg.exceptions import IcebergInvalidQueryError


class TestDuckDBQueryTool:
    """Test suite for DuckDBQueryTool class."""

    def test_init_with_defaults(self, mock_catalog, catalog_config):
        """Test initialization with default parameters."""
        with patch('langchain_iceberg.tools.duckdb_base.duckdb'):
            tool = DuckDBQueryTool(
                catalog=mock_catalog,
                catalog_config=catalog_config,
            )
            
            assert tool.name == "iceberg_sql_query"
            assert tool._query_timeout_seconds == 60
            assert tool._max_rows_per_query == 10000

    def test_init_with_custom_limits(self, mock_catalog, catalog_config):
        """Test initialization with custom limits."""
        with patch('langchain_iceberg.tools.duckdb_base.duckdb'):
            tool = DuckDBQueryTool(
                catalog=mock_catalog,
                catalog_config=catalog_config,
                query_timeout_seconds=120,
                max_rows_per_query=5000,
            )
            
            assert tool._query_timeout_seconds == 120
            assert tool._max_rows_per_query == 5000

    def test_run_simple_query(self, mock_catalog, catalog_config, mock_table):
        """Test running a simple SQL query."""
        with patch('langchain_iceberg.tools.duckdb_base.duckdb') as mock_duckdb:
            # Setup mocks
            mock_conn = Mock()
            mock_result = Mock()
            mock_df = pd.DataFrame({"total": [1000]})
            mock_result.fetchdf.return_value = mock_df
            mock_conn.execute.return_value = mock_result
            mock_duckdb.connect.return_value = mock_conn
            
            tool = DuckDBQueryTool(
                catalog=mock_catalog,
                catalog_config=catalog_config,
            )
            
            # Mock table loading
            mock_catalog.load_table.return_value = mock_table
            
            sql = "SELECT SUM(amount) as total FROM sales.orders"
            result = tool._run(sql_query=sql)
            
            assert result is not None
            assert isinstance(result, str)
            mock_conn.execute.assert_called()

    def test_run_join_query(self, mock_catalog, catalog_config, mock_table, sample_sql_query):
        """Test running a query with JOINs."""
        with patch('langchain_iceberg.tools.duckdb_base.duckdb') as mock_duckdb:
            # Setup mocks
            mock_conn = Mock()
            mock_result = Mock()
            mock_df = pd.DataFrame({
                "customer_segment": ["Premium", "Standard"],
                "order_count": [100, 200],
                "total_revenue": [50000, 30000]
            })
            mock_result.fetchdf.return_value = mock_df
            mock_conn.execute.return_value = mock_result
            mock_duckdb.connect.return_value = mock_conn
            
            tool = DuckDBQueryTool(
                catalog=mock_catalog,
                catalog_config=catalog_config,
            )
            
            # Mock table loading
            mock_catalog.load_table.return_value = mock_table
            
            result = tool._run(sql_query=sample_sql_query)
            
            assert result is not None
            assert isinstance(result, str)
            # Should have registered both tables
            assert "sales.orders" in tool._registered_tables
            assert "sales.customers" in tool._registered_tables

    def test_run_no_sql_query(self, mock_catalog, catalog_config):
        """Test error when no SQL query provided."""
        with patch('langchain_iceberg.tools.duckdb_base.duckdb'):
            tool = DuckDBQueryTool(
                catalog=mock_catalog,
                catalog_config=catalog_config,
            )
            
            with pytest.raises(IcebergInvalidQueryError, match="sql_query parameter is required"):
                tool._run(sql_query="")

    def test_run_timeout_exceeded(self, mock_catalog, catalog_config, mock_table):
        """Test timeout enforcement."""
        with patch('langchain_iceberg.tools.duckdb_base.duckdb') as mock_duckdb:
            # Setup mocks with slow execution
            mock_conn = Mock()
            mock_result = Mock()
            mock_df = pd.DataFrame({"total": [1000]})
            mock_result.fetchdf.return_value = mock_df
            mock_conn.execute.return_value = mock_result
            mock_duckdb.connect.return_value = mock_conn
            
            tool = DuckDBQueryTool(
                catalog=mock_catalog,
                catalog_config=catalog_config,
                query_timeout_seconds=0.001,  # Very short timeout
            )
            
            # Mock table loading
            mock_catalog.load_table.return_value = mock_table
            
            # Mock slow execution
            import time
            def slow_execute(*args, **kwargs):
                time.sleep(0.01)  # Longer than timeout
                return mock_result
            
            mock_conn.execute.side_effect = slow_execute
            
            sql = "SELECT * FROM sales.orders"
            
            with pytest.raises(IcebergInvalidQueryError, match="Query timeout exceeded"):
                tool._run(sql_query=sql)

    def test_run_row_limit_enforcement(self, mock_catalog, catalog_config, mock_table):
        """Test row limit enforcement."""
        with patch('langchain_iceberg.tools.duckdb_base.duckdb') as mock_duckdb:
            # Setup mocks with many rows
            mock_conn = Mock()
            mock_result = Mock()
            # Create DataFrame with more rows than limit
            mock_df = pd.DataFrame({"id": range(15000)})
            mock_result.fetchdf.return_value = mock_df
            mock_conn.execute.return_value = mock_result
            mock_duckdb.connect.return_value = mock_conn
            
            tool = DuckDBQueryTool(
                catalog=mock_catalog,
                catalog_config=catalog_config,
                max_rows_per_query=10000,
            )
            
            # Mock table loading
            mock_catalog.load_table.return_value = mock_table
            
            sql = "SELECT * FROM sales.orders"
            result = tool._run(sql_query=sql)
            
            assert "truncated" in result.lower()

    def test_run_json_input_format(self, mock_catalog, catalog_config, mock_table):
        """Test handling JSON string input."""
        with patch('langchain_iceberg.tools.duckdb_base.duckdb') as mock_duckdb:
            # Setup mocks
            mock_conn = Mock()
            mock_result = Mock()
            mock_df = pd.DataFrame({"total": [1000]})
            mock_result.fetchdf.return_value = mock_df
            mock_conn.execute.return_value = mock_result
            mock_duckdb.connect.return_value = mock_conn
            
            tool = DuckDBQueryTool(
                catalog=mock_catalog,
                catalog_config=catalog_config,
            )
            
            # Mock table loading
            mock_catalog.load_table.return_value = mock_table
            
            # JSON string input (some LLM frameworks send this)
            json_input = '{"sql_query": "SELECT SUM(amount) as total FROM sales.orders"}'
            result = tool._run(sql_query=json_input)
            
            assert result is not None
            mock_conn.execute.assert_called()

    def test_run_limit_injection(self, mock_catalog, catalog_config, mock_table):
        """Test automatic LIMIT injection for safety."""
        with patch('langchain_iceberg.tools.duckdb_base.duckdb') as mock_duckdb:
            # Setup mocks
            mock_conn = Mock()
            mock_result = Mock()
            mock_df = pd.DataFrame({"id": [1, 2, 3]})
            mock_result.fetchdf.return_value = mock_df
            mock_conn.execute.return_value = mock_result
            mock_duckdb.connect.return_value = mock_conn
            
            tool = DuckDBQueryTool(
                catalog=mock_catalog,
                catalog_config=catalog_config,
                max_rows_per_query=100,
            )
            
            # Mock table loading
            mock_catalog.load_table.return_value = mock_table
            
            # Query without LIMIT
            sql = "SELECT * FROM sales.orders WHERE status = 'completed'"
            tool._run(sql_query=sql)
            
            # Check that LIMIT was added
            executed_sql = str(mock_conn.execute.call_args)
            assert "LIMIT" in executed_sql.upper()


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

    def test_description_format(self, mock_catalog, catalog_config):
        """Test tool description is properly formatted."""
        with patch('langchain_iceberg.tools.duckdb_base.duckdb'):
            tool = DuckDBQueryTool(
                catalog=mock_catalog,
                catalog_config=catalog_config,
            )
            
            assert tool.description is not None
            assert "JOIN" in tool.description
            assert "sql_query" in tool.description
            assert len(tool.description) > 50

    def test_kwargs_handling(self, mock_catalog, catalog_config, mock_table):
        """Test handling of kwargs in _run method."""
        with patch('langchain_iceberg.tools.duckdb_base.duckdb') as mock_duckdb:
            mock_conn = Mock()
            mock_result = Mock()
            mock_df = pd.DataFrame({"total": [1000]})
            mock_result.fetchdf.return_value = mock_df
            mock_conn.execute.return_value = mock_result
            mock_duckdb.connect.return_value = mock_conn
            
            tool = DuckDBQueryTool(
                catalog=mock_catalog,
                catalog_config=catalog_config,
            )
            
            mock_catalog.load_table.return_value = mock_table
            
            # SQL in kwargs instead of direct parameter
            result = tool._run(sql_query="SELECT * FROM sales.orders")
            
            assert result is not None
