"""
Integration tests for AI-enhanced database operations with real database.

Tier 2 Testing (NO MOCKING) - Uses real database infrastructure.

Verifies:
- End-to-end NL to SQL execution
- Data transformation pipelines with real database
- Quality assessment workflows
- Intelligent bulk operations
- Semantic database search

Prerequisites:
- PostgreSQL or SQLite database available
- DataFlow installed
- Real LLM provider configured (or skip with @pytest.mark.real_llm)
"""

import pytest
import os
from typing import Dict, List, Any


@pytest.fixture(scope="module")
def database_url():
    """
    Get database URL for integration tests.

    Uses SQLite for local testing, can be overridden with env var.
    """
    return os.getenv(
        'TEST_DATABASE_URL',
        'sqlite:///test_ai_db_operations.db'
    )


@pytest.fixture(scope="module")
def dataflow_instance(database_url):
    """
    Create real DataFlow instance for integration tests.

    This is a REAL database connection - no mocking.
    """
    try:
        from kailash_dataflow import DataFlow
    except ImportError:
        pytest.skip("DataFlow not available")

    # Create DataFlow instance
    db = DataFlow(database_url, auto_migrate=True)

    # Define test models
    @db.model
    class User:
        name: str
        email: str
        age: int = None
        department: str = None

    @db.model
    class Product:
        name: str
        inventory: int
        demand_score: float = None
        category: str = None

    yield db

    # Cleanup: Drop test database (SQLite)
    if database_url.startswith('sqlite:///'):
        import os
        db_file = database_url.replace('sqlite:///', '')
        if os.path.exists(db_file):
            os.remove(db_file)


@pytest.fixture
def agent_config():
    """Create agent configuration for integration tests."""
    from kaizen.core.config import BaseAgentConfig
    return BaseAgentConfig(
        llm_provider=os.getenv('TEST_LLM_PROVIDER', 'mock'),
        model=os.getenv('TEST_LLM_MODEL', 'gpt-4'),
        temperature=0.2
    )


class TestNLQueryToDataFlowExecution:
    """Integration tests for natural language query to execution."""

    @pytest.fixture
    def nl_agent_class(self):
        """Get NLToSQLAgent class."""
        try:
            from kaizen.integrations.dataflow.ai_enhanced_ops import NLToSQLAgent
            return NLToSQLAgent
        except ImportError:
            pytest.skip("AI enhanced operations not implemented yet")

    @pytest.fixture
    def setup_test_data(self, dataflow_instance):
        """Setup test data in real database."""
        from kailash.workflow.builder import WorkflowBuilder
        from kailash.runtime.local import LocalRuntime

        # Insert test users
        workflow = WorkflowBuilder()
        workflow.add_node("UserBulkCreateNode", "create_users", {
            "data": [
                {"name": "Alice Smith", "email": "alice@example.com", "age": 28, "department": "Engineering"},
                {"name": "Bob Jones", "email": "bob@example.com", "age": 35, "department": "Sales"},
                {"name": "Carol White", "email": "carol@example.com", "age": 22, "department": "Engineering"},
                {"name": "David Brown", "email": "david@example.com", "age": 45, "department": "Marketing"}
            ]
        })

        runtime = LocalRuntime()
        results, _ = runtime.execute(workflow.build())

        return results

    @pytest.mark.real_llm
    def test_nl_query_to_dataflow_execution(
        self,
        nl_agent_class,
        agent_config,
        dataflow_instance,
        setup_test_data
    ):
        """
        Test end-to-end NL query execution.

        Given: Real database with user data
        When: Natural language query executed
        Then: Returns correct results from database
        """
        agent = nl_agent_class(config=agent_config, db=dataflow_instance)

        result = agent.query(
            "Show me all users in the Engineering department"
        )

        # Should return results
        assert result is not None
        assert 'results' in result
        assert 'explanation' in result

        # Should have explanation
        assert len(result['explanation']) > 0

    @pytest.mark.real_llm
    def test_nl_query_with_age_filter(
        self,
        nl_agent_class,
        agent_config,
        dataflow_instance,
        setup_test_data
    ):
        """
        Test NL query with numeric comparison.

        Given: Database with users of different ages
        When: Query "users over 30"
        Then: Returns only users with age > 30
        """
        agent = nl_agent_class(config=agent_config, db=dataflow_instance)

        result = agent.query(
            "Find users over 30 years old"
        )

        # Should filter correctly
        assert result is not None

    @pytest.mark.real_llm
    def test_nl_query_aggregation(
        self,
        nl_agent_class,
        agent_config,
        dataflow_instance,
        setup_test_data
    ):
        """
        Test NL query with aggregation.

        Given: Database with users
        When: Query "count users by department"
        Then: Returns aggregated results
        """
        agent = nl_agent_class(config=agent_config, db=dataflow_instance)

        result = agent.query(
            "Count how many users are in each department"
        )

        # Should aggregate
        assert result is not None


class TestDataTransformationPipeline:
    """Integration tests for data transformation with real database."""

    @pytest.fixture
    def transform_agent_class(self):
        """Get DataTransformAgent class."""
        try:
            from kaizen.integrations.dataflow.ai_enhanced_ops import DataTransformAgent
            return DataTransformAgent
        except ImportError:
            pytest.skip("Data transform agent not implemented yet")

    @pytest.mark.real_llm
    def test_transform_data_between_schemas(
        self,
        transform_agent_class,
        agent_config,
        dataflow_instance
    ):
        """
        Test data transformation pipeline.

        Given: Source data with different schema
        When: Agent transforms to target schema
        Then: Data inserted into real database
        """
        agent = transform_agent_class(config=agent_config, db=dataflow_instance)

        # Source data with different field names
        source_data = [
            {
                'full_name': 'Test User 1',
                'email_address': 'test1@example.com',
                'years_old': 30
            },
            {
                'full_name': 'Test User 2',
                'email_address': 'test2@example.com',
                'years_old': 25
            }
        ]

        result = agent.transform_data(
            source_data=source_data,
            target_table='User'
        )

        # Should insert data
        assert result is not None
        assert 'inserted_count' in result or 'transformed_data' in result

    @pytest.mark.real_llm
    def test_transformation_with_quality_check(
        self,
        transform_agent_class,
        agent_config,
        dataflow_instance
    ):
        """
        Test transformation includes quality assessment.

        Given: Source data with quality issues
        When: Agent transforms data
        Then: Reports quality issues and handles gracefully
        """
        agent = transform_agent_class(config=agent_config, db=dataflow_instance)

        source_data = [
            {'name': '', 'email': 'invalid-email', 'age': -5},  # Bad data
            {'name': 'Valid User', 'email': 'valid@example.com', 'age': 30}
        ]

        result = agent.transform_data(
            source_data=source_data,
            target_table='User'
        )

        # Should identify quality issues
        assert result is not None
        if 'quality_issues' in result:
            assert len(result['quality_issues']) > 0


class TestQualityAssessmentWorkflow:
    """Integration tests for data quality assessment."""

    @pytest.fixture
    def quality_agent_class(self):
        """Get DataQualityAgent class."""
        try:
            from kaizen.integrations.dataflow.ai_enhanced_ops import DataQualityAgent
            return DataQualityAgent
        except ImportError:
            pytest.skip("Data quality agent not implemented yet")

    @pytest.fixture
    def setup_quality_test_data(self, dataflow_instance):
        """Setup data with quality issues."""
        from kailash.workflow.builder import WorkflowBuilder
        from kailash.runtime.local import LocalRuntime

        workflow = WorkflowBuilder()
        workflow.add_node("UserBulkCreateNode", "create_users", {
            "data": [
                {"name": "Good User", "email": "good@example.com", "age": 25},
                {"name": "", "email": "bad", "age": 999},  # Quality issues
                {"name": "Another Good User", "email": "another@example.com", "age": 30}
            ]
        })

        runtime = LocalRuntime()
        results, _ = runtime.execute(workflow.build())

        return results

    @pytest.mark.real_llm
    def test_assess_and_fix_data_quality(
        self,
        quality_agent_class,
        agent_config,
        dataflow_instance,
        setup_quality_test_data
    ):
        """
        Test quality assessment and fixing workflow.

        Given: Real database with quality issues
        When: Agent assesses and fixes
        Then: Quality improves
        """
        agent = quality_agent_class(config=agent_config)

        # Get sample data from database
        from kailash.workflow.builder import WorkflowBuilder
        from kailash.runtime.local import LocalRuntime

        workflow = WorkflowBuilder()
        workflow.add_node("UserListNode", "list_users", {"limit": 10})

        runtime = LocalRuntime()
        results, _ = runtime.execute(workflow.build())

        users = results.get('list_users', [])

        # Assess quality
        schema = {
            'columns': {
                'name': {'type': 'str', 'nullable': False},
                'email': {'type': 'str', 'nullable': False},
                'age': {'type': 'int', 'nullable': True}
            }
        }

        result = agent.assess_quality(
            data_sample=users,
            schema=schema
        )

        # Should identify issues
        assert result is not None
        assert 'quality_score' in result or 'issues_found' in result


class TestIntelligentBulkInsert:
    """Integration tests for AI-optimized bulk operations."""

    @pytest.fixture
    def bulk_optimizer_class(self):
        """Get bulk operation optimizer class."""
        try:
            from kaizen.integrations.dataflow.query_optimizer import BulkOperationOptimizer
            return BulkOperationOptimizer
        except ImportError:
            pytest.skip("Bulk optimizer not implemented yet")

    @pytest.mark.real_llm
    def test_intelligent_bulk_insert_optimization(
        self,
        bulk_optimizer_class,
        agent_config,
        dataflow_instance
    ):
        """
        Test AI-optimized bulk insert.

        Given: Large dataset for insertion
        When: Agent optimizes bulk operation
        Then: Inserts efficiently with optimal batch size
        """
        optimizer = bulk_optimizer_class(config=agent_config)

        # Generate large dataset
        large_dataset = [
            {
                'name': f'User {i}',
                'email': f'user{i}@example.com',
                'age': 20 + (i % 50)
            }
            for i in range(1000)
        ]

        result = optimizer.optimize_bulk_insert(
            data=large_dataset,
            target_table='User'
        )

        # Should optimize
        assert result is not None
        assert 'batch_size' in result or 'strategy' in result


class TestSemanticDatabaseSearch:
    """Integration tests for semantic search across tables."""

    @pytest.fixture
    def semantic_agent_class(self):
        """Get SemanticSearchAgent class."""
        try:
            from kaizen.integrations.dataflow.ai_enhanced_ops import SemanticSearchAgent
            return SemanticSearchAgent
        except ImportError:
            pytest.skip("Semantic search agent not implemented yet")

    @pytest.fixture
    def setup_multi_table_data(self, dataflow_instance):
        """Setup data across multiple tables."""
        from kailash.workflow.builder import WorkflowBuilder
        from kailash.runtime.local import LocalRuntime

        # Insert users
        workflow = WorkflowBuilder()
        workflow.add_node("UserBulkCreateNode", "create_users", {
            "data": [
                {"name": "Alice", "email": "alice@example.com", "age": 25},
                {"name": "Bob", "email": "bob@example.com", "age": 30}
            ]
        })

        # Insert products
        workflow.add_node("ProductBulkCreateNode", "create_products", {
            "data": [
                {"name": "Laptop", "inventory": 5, "demand_score": 0.9, "category": "Electronics"},
                {"name": "Phone", "inventory": 15, "demand_score": 0.95, "category": "Electronics"},
                {"name": "Desk", "inventory": 20, "demand_score": 0.3, "category": "Furniture"}
            ]
        })

        runtime = LocalRuntime()
        results, _ = runtime.execute(workflow.build())

        return results

    @pytest.mark.real_llm
    def test_semantic_search_across_tables(
        self,
        semantic_agent_class,
        agent_config,
        dataflow_instance,
        setup_multi_table_data
    ):
        """
        Test semantic search across multiple tables.

        Given: Real database with users and products
        When: Semantic query executed
        Then: Identifies relevant tables and returns results
        """
        agent = semantic_agent_class(config=agent_config, db=dataflow_instance)

        result = agent.search(
            query="Find popular electronics with low inventory"
        )

        # Should search semantically
        assert result is not None
        assert 'relevant_tables' in result or 'results' in result

    @pytest.mark.real_llm
    def test_semantic_search_with_context(
        self,
        semantic_agent_class,
        agent_config,
        dataflow_instance,
        setup_multi_table_data
    ):
        """
        Test semantic search with additional context.

        Given: Database with multiple tables
        When: Query with business context
        Then: Uses context to improve search
        """
        agent = semantic_agent_class(config=agent_config, db=dataflow_instance)

        result = agent.search(
            query="Items we should restock soon",
            context="Focus on high demand products"
        )

        # Should use context
        assert result is not None


class TestEndToEndAIPipeline:
    """Integration tests for complete AI-enhanced workflows."""

    @pytest.fixture
    def all_agent_classes(self):
        """Get all AI agent classes."""
        try:
            from kaizen.integrations.dataflow.ai_enhanced_ops import (
                NLToSQLAgent,
                DataTransformAgent,
                DataQualityAgent
            )
            return {
                'nl_to_sql': NLToSQLAgent,
                'transform': DataTransformAgent,
                'quality': DataQualityAgent
            }
        except ImportError:
            pytest.skip("AI enhanced operations not implemented yet")

    @pytest.mark.real_llm
    def test_complete_ai_pipeline(
        self,
        all_agent_classes,
        agent_config,
        dataflow_instance
    ):
        """
        Test complete AI-enhanced data pipeline.

        Given: Raw data needing transformation
        When: Pipeline runs (quality check → transform → query)
        Then: Data processed and queryable via NL
        """
        # Step 1: Assess quality
        quality_agent = all_agent_classes['quality'](config=agent_config)

        raw_data = [
            {'full_name': '  Alice Smith  ', 'email': 'ALICE@EXAMPLE.COM', 'age': 28},
            {'full_name': 'Bob Jones', 'email': 'bob@example.com', 'age': 35}
        ]

        quality_result = quality_agent.assess_quality(
            data_sample=raw_data,
            schema={'columns': {'name': {'type': 'str'}, 'email': {'type': 'str'}, 'age': {'type': 'int'}}}
        )

        # Step 2: Transform and load
        transform_agent = all_agent_classes['transform'](config=agent_config, db=dataflow_instance)

        transform_result = transform_agent.transform_data(
            source_data=raw_data,
            target_table='User'
        )

        # Step 3: Query via natural language
        nl_agent = all_agent_classes['nl_to_sql'](config=agent_config, db=dataflow_instance)

        query_result = nl_agent.query("Show me all users")

        # All steps should succeed
        assert quality_result is not None
        assert transform_result is not None
        assert query_result is not None


if __name__ == '__main__':
    pytest.main([__file__, '-v'])
