import os
import pytest
import pandas as pd
import networkx as nx
from unittest.mock import MagicMock, patch, call
from syda.schemas import ModelConfig
from syda.generate import SyntheticDataGenerator


class TestSyntheticDataGenerator:
    """Test the SyntheticDataGenerator class."""

    @pytest.fixture
    def sample_schema(self):
        """Sample schema for testing."""
        return {
            "id": {"type": "number", "description": "Unique identifier"},
            "name": {"type": "text", "description": "Customer name"},
            "email": {"type": "email", "description": "Customer email"}
        }

    @pytest.fixture
    def sample_schema_with_fk(self):
        """Sample schema with foreign key for testing."""
        return {
            "id": {"type": "number", "description": "Unique identifier"},
            "customer_id": {
                "type": "foreign_key",
                "description": "Reference to customer",
                "references": {"schema": "Customer", "field": "id"}
            }
        }

    def test_initialization(self):
        """Test generator initialization."""
        # Create with OpenAI configuration
        generator1 = SyntheticDataGenerator(
            model_config=ModelConfig(provider="openai", model_name="gpt-4"),
            openai_api_key="test_key"
        )
        assert generator1.model_config.provider == "openai"
        assert generator1.model_config.model_name == "gpt-4"
        
        # Create with Anthropic configuration
        generator2 = SyntheticDataGenerator(
            model_config=ModelConfig(provider="anthropic", model_name="claude-3"),
            anthropic_api_key="test_key"
        )
        assert generator2.model_config.provider == "anthropic"
        assert generator2.model_config.model_name == "claude-3"

        # Create with Gemini configuration
        generator3 = SyntheticDataGenerator(
            model_config=ModelConfig(provider="gemini", model_name="gemini-2.5-flash"),
            gemini_api_key="test_key"
        )
        assert generator3.model_config.provider == "gemini"
        assert generator3.model_config.model_name == "gemini-2.5-flash"

    @patch("syda.generate.SchemaLoader")
    def test_generate_for_schemas(self, mock_schema_loader, sample_schema):
        """Test generating data for schemas."""
        # Configure the mock
        mock_loader_instance = MagicMock()
        mock_schema_loader.return_value = mock_loader_instance
        # Configure the load_schema method to return the expected 6 values
        mock_loader_instance.load_schema.return_value = (sample_schema, {}, "Test schema", {}, {}, [])
        
        # Create the generator with a mock LLM client to avoid API calls
        generator = SyntheticDataGenerator(model_config=ModelConfig(provider="openai", model_name="gpt-4"), openai_api_key="test_key")
        
        # Mock the _generate_data method to avoid actual data generation
        generator._generate_data = MagicMock(return_value=pd.DataFrame())
        
        # Generate data using the schema
        result = generator.generate_for_schemas(
            schemas={"customers": sample_schema},
            default_sample_size=5
        )
        
        # Check that the result is a dictionary containing the customers DataFrame
        assert isinstance(result, dict)
        assert "customers" in result
        assert isinstance(result["customers"], pd.DataFrame)
        
    @patch("syda.generate.SchemaLoader")
    def test_generate_with_dependencies(self, mock_schema_loader):
        """Test generating data with foreign key dependencies."""
        # Create schemas with foreign key relationships
        schemas = {
            "customers": {"id": {"type": "number"}, "name": {"type": "text"}},
            "orders": {
                "id": {"type": "number"},
                "customer_id": {
                    "type": "foreign_key",
                    "references": {"schema": "customers", "field": "id"}
                },
                "amount": {"type": "number"}
            }
        }

        # Create mock instances and preset return values
        customers_df = pd.DataFrame({"id": [101, 102], "name": ["Customer A", "Customer B"]})
        orders_df = pd.DataFrame({"id": [1, 2], "customer_id": [101, 102], "amount": [100.0, 200.0]})
        
        # Create a generator with proper model configuration
        with patch("syda.dependency_handler.DependencyHandler") as mock_dependency_handler, \
             patch("syda.generate.SyntheticDataGenerator._generate_data") as mock_generate_data:
            
            # Set up dependency handler mock
            mock_dep_instance = mock_dependency_handler.return_value
            mock_dep_instance.extract_dependencies.return_value = {"orders": ["customers"]}
            mock_dep_instance.determine_generation_order.return_value = ["customers", "orders"]
            mock_dep_instance.has_cycle.return_value = False
            
            # Set up schema loader mock to return expected values
            mock_loader_instance = mock_schema_loader.return_value
            
            def mock_load_schema(schema_source):
                # Debug: Print what's being loaded
                print(f"Loading schema: {schema_source}")
                if isinstance(schema_source, dict):
                    if schema_source == schemas["customers"]:
                        return (schema_source, {}, "Customers Schema", {}, {}, [])
                    elif schema_source == schemas["orders"]:
                        return (schema_source, {}, "Orders Schema", 
                              {"customer_id": ("customers", "id")}, {}, ["customers"])
                return (schema_source, {}, "Unknown Schema", {}, {}, [])
                
            mock_loader_instance.load_schema.side_effect = mock_load_schema
            
            # Configure _generate_data mock to return our predetermined DataFrames
            def side_effect_generate_data(table_schema, metadata=None, table_description=None, prompt=None, sample_size=None, schema_name=None, **kwargs):
                print(f"Generating data for schema_name: {schema_name}")
                if schema_name == "customers":
                    return customers_df
                elif schema_name == "orders":
                    return orders_df
                # Fallback based on the schema itself
                if isinstance(table_schema, dict) and "name" in table_schema:
                    return customers_df
                if isinstance(table_schema, dict) and "customer_id" in table_schema:
                    return orders_df
                return pd.DataFrame()
                
            mock_generate_data.side_effect = side_effect_generate_data
            
            # Create the generator and run the test
            generator = SyntheticDataGenerator(
                model_config=ModelConfig(provider="openai", model_name="gpt-4"),
                openai_api_key="test_key"
            )
            
            # Generate data using the schemas
            print("Calling generate_for_schemas with schemas:")
            for name, schema in schemas.items():
                print(f"  - {name}: {schema}")
                
            result = generator.generate_for_schemas(
                schemas=schemas,
                default_sample_size=5
            )
            
            print(f"Result keys: {result.keys() if result else 'Empty result'}")
            
            # Verify the results contain both DataFrames
            assert isinstance(result, dict)
            assert "customers" in result, "Customers DataFrame not found in result"
            assert "orders" in result, "Orders DataFrame not found in result"
            assert isinstance(result["customers"], pd.DataFrame)
            assert isinstance(result["orders"], pd.DataFrame)
            
            # Verify the mock calls were made correctly
            assert mock_generate_data.call_count >= 2, f"_generate_data was called {mock_generate_data.call_count} times, expected at least 2"
    
    def test_generate_data(self, sample_schema):
        """Test generating data."""
        # Create the generator
        generator = SyntheticDataGenerator(model_config=ModelConfig(provider="openai", model_name="gpt-4"), openai_api_key="test_key")
        
        # Create a mock for _generate_data
        with patch.object(generator, "_generate_data") as mock_generate_data:
            # Mock response from the generate_data method
            mock_df = pd.DataFrame({
                "id": [1, 2],
                "name": ["Alice", "Bob"],
                "email": ["alice@example.com", "bob@example.com"]
            })
            mock_generate_data.return_value = mock_df
            
            # Configure schema loader mock to handle our test schema
            generator.schema_loader = MagicMock()
            generator.schema_loader.load_schema.return_value = (sample_schema, {}, "Test schema", {}, {}, [])
            
            # Generate data using a simple schema
            result = generator.generate_for_schemas({
                "customers": sample_schema
            }, default_sample_size=2)
            
            # Check that mock_generate_data was called
            assert mock_generate_data.called, "_generate_data was not called"
            
            # Check that the result contains the customers DataFrame
            assert isinstance(result, dict)
            assert "customers" in result
            assert isinstance(result["customers"], pd.DataFrame)
        
    def test_schema_loading_error_handling(self):
        """Test error handling when loading an invalid schema."""
        # Create a generator with proper model configuration
        generator = SyntheticDataGenerator(model_config=ModelConfig(provider="openai", model_name="gpt-4"), openai_api_key="test_key")
        
        # Create a schema loader that raises an exception
        generator.schema_loader = MagicMock()
        generator.schema_loader.load_schema.side_effect = ValueError("Invalid schema format")
        
        # Try to generate data with an invalid schema source
        with pytest.raises(ValueError, match="Invalid schema format"):
            generator.generate_for_schemas({"invalid_schema": "invalid_source"})
    
    @patch("syda.dependency_handler.DependencyHandler")
    def test_generate_data_with_foreign_keys(self, mock_dependency_handler, sample_schema_with_fk):
        """Test generating data with foreign key dependencies."""
        # Configure the dependency handler mock
        mock_dep_instance = MagicMock()
        mock_dependency_handler.return_value = mock_dep_instance
        mock_dep_instance.extract_dependencies.return_value = {"orders": ["customers"]}
        mock_dep_instance.build_dependency_graph.return_value = MagicMock()
        mock_dep_instance.determine_generation_order.return_value = ["customers", "orders"]
        
        # Create a generator with proper model configuration
        generator = SyntheticDataGenerator(model_config=ModelConfig(provider="openai", model_name="gpt-4"), openai_api_key="test_key")
        
        # Mock the _generate_data method to return known DataFrames
        def mock_generate_data(table_schema, metadata=None, table_description=None, prompt=None, sample_size=None, schema_name=None, **kwargs):
            if schema_name == "customers" or (isinstance(table_schema, dict) and "name" in table_schema):
                return pd.DataFrame({"id": [101, 102], "name": ["Test1", "Test2"]})
            else:
                return pd.DataFrame({"id": [1, 2], "customer_id": [101, 102]})
                
        generator._generate_data = MagicMock(side_effect=mock_generate_data)
        
        # Create schemas with foreign key relationships
        schemas = {
            "customers": {"id": {"type": "number"}, "name": {"type": "text"}},
            "orders": {
                "id": {"type": "number"},
                "customer_id": {"type": "foreign_key", "references": {"schema": "customers", "field": "id"}}
            }
        }
        
        # Mock the schema loading to handle foreign keys correctly
        def mock_load_schema(schema_source):
            if "id" in schema_source and "name" in schema_source:
                # Customer schema
                return (schema_source, {}, "Customers Schema", {}, {}, [])
            else:
                # Order schema with foreign key
                return (schema_source, {}, "Orders Schema",
                       {"customer_id": ("customers", "id")}, {}, ["customers"])
                
        generator.schema_loader = MagicMock()
        generator.schema_loader.load_schema.side_effect = mock_load_schema
        
        # Generate data with foreign keys
        results = generator.generate_for_schemas(schemas, default_sample_size=2)
        
        # Verify results
        assert "customers" in results
        assert "orders" in results
        assert isinstance(results["customers"], pd.DataFrame)
        assert isinstance(results["orders"], pd.DataFrame)
        
        # Check that customer IDs in orders match IDs in customers
        orders_df = results["orders"]
        customers_df = results["customers"]
        assert set(orders_df["customer_id"]).issubset(set(customers_df["id"]))
    
    @patch("syda.generate.create_empty_dataframe")
    @patch("syda.generate.GeneratorManager")
    def test_generate_random_data(self, mock_generator_manager, mock_create_empty_df, sample_schema):
        """Test fallback to random data generation when LLM generation fails."""
        # Configure the empty DataFrame mock
        mock_df = pd.DataFrame(columns=["id", "name", "email"])
        mock_create_empty_df.return_value = mock_df
        
        # Configure the generator manager mock
        mock_manager_instance = MagicMock()
        mock_generator_manager.return_value = mock_manager_instance
        mock_manager_instance.generate.side_effect = [1, "Alice", "alice@example.com"]
        mock_manager_instance.apply_custom_generators.return_value = pd.DataFrame({
            "id": [1],
            "name": ["Alice"],
            "email": ["alice@example.com"]
        })
        
        # Create the generator with proper model configuration
        generator = SyntheticDataGenerator(model_config=ModelConfig(provider="openai", model_name="gpt-4"), openai_api_key="test_key")
        
        # Mock schema loader to provide the test schema
        generator.schema_loader = MagicMock()
        generator.schema_loader.load_schema.return_value = (sample_schema, {}, "Test schema", {}, {}, [])
        
        # Mock _generate_data to raise an exception to trigger fallback
        with patch.object(generator, "_generate_data", side_effect=Exception("Simulated LLM failure")):
            # Set up schemas for the test
            schemas = {"customers": sample_schema}
            
            # When _generate_data fails, the code should fall back to random generation
            # Pass a small default_sample_size to make the test faster
            with pytest.raises(Exception, match="Simulated LLM failure"):
                generator.generate_for_schemas(schemas, default_sample_size=1)
    
    @patch("syda.generate.save_dataframes")
    def test_save_dataframes(self, mock_save_dataframes):
        """Test saving DataFrames."""
        # Create DataFrames
        df1 = pd.DataFrame({"id": [1, 2]})
        df2 = pd.DataFrame({"id": [3, 4]})
        
        # Create the generator with proper model configuration
        generator = SyntheticDataGenerator(model_config=ModelConfig(provider="openai", model_name="gpt-4"), openai_api_key="test_key")
        
        # Set up dataframes for the test
        dataframes = {
            "table1": df1,
            "table2": df2
        }
        
        # Directly use the save_dataframes function that's being used in generate_for_schemas
        from syda.generate import save_dataframes
        save_dataframes(dataframes, "output_dir", format="csv")
        
        # Check that save_dataframes was called with the correct parameters
        # Adjust assertion to match the actual call (without filenames parameter)
        mock_save_dataframes.assert_called_once_with(
            dataframes, "output_dir", format="csv"
        )
    
    def test_generate_all_schemas(self):
        """Test generating data for all schemas."""
        # Create a generator with proper model configuration
        generator = SyntheticDataGenerator(model_config=ModelConfig(provider="openai", model_name="gpt-4"), openai_api_key="test_key")
        
        # Create test schemas
        schemas = {
            "table1": {"id": {"type": "number"}},
            "table2": {"id": {"type": "number"}}
        }
        
        # Mock the schema loader
        generator.schema_loader = MagicMock()
        generator.schema_loader.load_schema.side_effect = [
            (schemas["table1"], {}, "Table 1", {}, {}, []),
            (schemas["table2"], {}, "Table 2", {}, {}, [])
        ]
        
        # Mock _generate_data to avoid actual LLM calls
        generator._generate_data = MagicMock(return_value=pd.DataFrame({"id": [1, 2]}))
        
        # Generate data for all schemas
        result = generator.generate_for_schemas(
            schemas=schemas,
            default_sample_size=10
        )
        
        # Verify results contain both tables
        assert "table1" in result
        assert "table2" in result
        assert isinstance(result["table1"], pd.DataFrame)
        assert isinstance(result["table2"], pd.DataFrame)
        
    def test_detect_dependency_cycle(self):
        """Test that a circular dependency in schemas is handled gracefully."""
        # Create schemas with a circular dependency
        schemas = {
            'table1': {'id': {'type': 'number'}, 'table2_id': {'type': 'foreign_key', 'references': {'schema': 'table2', 'field': 'id'}}},
            'table2': {'id': {'type': 'number'}, 'table1_id': {'type': 'foreign_key', 'references': {'schema': 'table1', 'field': 'id'}}}
        }
        
        # Mock the dependency handler's methods to simulate a cycle
        with patch('syda.dependency_handler.DependencyHandler.determine_generation_order') as mock_determine:
            # Simulate a circular dependency error
            mock_determine.side_effect = ValueError("Circular dependencies detected in schemas")
            
            # Mock the foreign key handler to prevent foreign key errors
            with patch('syda.dependency_handler.ForeignKeyHandler.apply_foreign_keys'):
                # Mock _generate_data to avoid actual API calls and return dataframes with proper columns
                with patch('syda.generate.SyntheticDataGenerator._generate_data') as mock_generate_data:
                    # Return a DataFrame with the 'id' column for both tables
                    mock_generate_data.side_effect = lambda **kwargs: pd.DataFrame({'id': [1, 2], 
                                                                                 'table1_id' if 'table2' in kwargs['table_description'] else 'table2_id': [1, 2]})
                    
                    # Create the generator
                    generator = SyntheticDataGenerator(
                        model_config=ModelConfig(provider="openai", model_name="gpt-4"),
                        openai_api_key="test_key"
                    )
                    
                    # The method should handle the cycle and not raise an exception
                    result = generator.generate_for_schemas(schemas)
                    
                    # Verify that determine_generation_order was called
                    mock_determine.assert_called_once()
                    
                    # Verify that we still got results for both schemas
                    assert set(result.keys()) == set(['table1', 'table2'])
                    assert isinstance(result['table1'], pd.DataFrame)
                    assert isinstance(result['table2'], pd.DataFrame)
            
    def test_client_error_handling(self):
        """Test handling of client errors when generating data."""
        # Create a schema
        test_schema = {'id': {'type': 'number'}}
        
        # Create the generator
        generator = SyntheticDataGenerator(
            model_config=ModelConfig(provider="openai", model_name="gpt-4"),
            openai_api_key="test_key"
        )
        
        # Mock the schema loading to avoid external dependencies
        generator.schema_loader = MagicMock()
        generator.schema_loader.load_schema.return_value = (test_schema, {}, "Test schema", {}, {}, [])
        
        # Patch the instance method _generate_data to raise an exception
        with patch.object(SyntheticDataGenerator, "_generate_data") as mock_generate_data:
            # Mock the generate_data method to raise a client error
            mock_generate_data.side_effect = Exception("Simulated client error")
            
            # Test that the exception is properly propagated
            with pytest.raises(Exception, match="Simulated client error"):
                generator.generate_for_schemas({"customers": test_schema}, default_sample_size=10)
