#!/usr/bin/env python3
"""
Unit tests for Snowflake object management functionality
"""

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

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

from snowflake_object_manager import SnowflakeObjectManager
from user_management import UserManager, PermissionManager, PermissionType
from sigma_api_client import SigmaAPIClient

class TestSnowflakeObjectManager(unittest.TestCase):
    """Test cases for SnowflakeObjectManager"""
    
    def setUp(self):
        """Set up test fixtures"""
        self.mock_client = Mock(spec=SigmaAPIClient)
        self.snowflake_manager = SnowflakeObjectManager(self.mock_client)
    
    def test_init(self):
        """Test SnowflakeObjectManager initialization"""
        self.assertEqual(self.snowflake_manager.api, self.mock_client)
        self.assertIsInstance(self.snowflake_manager.user_manager, UserManager)
        self.assertIsInstance(self.snowflake_manager.permission_manager, PermissionManager)
    
    def test_list_connections_success(self):
        """Test successful connections listing"""
        mock_response = {
            'entries': [
                {'connectionId': '1', 'name': 'Snowflake Conn 1', 'type': 'snowflake'},
                {'connectionId': '2', 'name': 'Postgres Conn', 'type': 'postgres'},
                {'connectionId': '3', 'name': 'Snowflake Conn 2', 'type': 'snowflake'}
            ]
        }
        self.mock_client.get.return_value = mock_response
        
        result = self.snowflake_manager.list_connections()
        
        self.assertEqual(len(result), 3)
        self.mock_client.get.assert_called_once_with('/v2/connections')
    
    def test_list_connections_error(self):
        """Test connections listing error handling"""
        self.mock_client.get.side_effect = Exception("API Error")
        
        result = self.snowflake_manager.list_connections()
        
        self.assertEqual(result, [])
    
    def test_get_snowflake_connection_by_name(self):
        """Test finding Snowflake connection by name"""
        mock_connections = [
            {'connectionId': '1', 'name': 'Snowflake Conn 1', 'type': 'snowflake'},
            {'connectionId': '2', 'name': 'Postgres Conn', 'type': 'postgres'},
            {'connectionId': '3', 'name': 'Target Connection', 'type': 'snowflake'}
        ]
        
        with patch.object(self.snowflake_manager, 'list_connections', return_value=mock_connections):
            result = self.snowflake_manager.get_snowflake_connection('Target Connection')
        
        self.assertIsNotNone(result)
        self.assertEqual(result['connectionId'], '3')
        self.assertEqual(result['name'], 'Target Connection')
    
    def test_get_snowflake_connection_first_snowflake(self):
        """Test getting first Snowflake connection when no name specified"""
        mock_connections = [
            {'connectionId': '1', 'name': 'Postgres Conn', 'type': 'postgres'},
            {'connectionId': '2', 'name': 'First Snowflake', 'type': 'snowflake'},
            {'connectionId': '3', 'name': 'Second Snowflake', 'type': 'snowflake'}
        ]
        
        with patch.object(self.snowflake_manager, 'list_connections', return_value=mock_connections):
            result = self.snowflake_manager.get_snowflake_connection()
        
        self.assertIsNotNone(result)
        self.assertEqual(result['connectionId'], '2')  # First Snowflake connection
    
    def test_get_snowflake_connection_not_found(self):
        """Test Snowflake connection not found"""
        mock_connections = [
            {'connectionId': '1', 'name': 'Postgres Conn', 'type': 'postgres'}
        ]
        
        with patch.object(self.snowflake_manager, 'list_connections', return_value=mock_connections):
            result = self.snowflake_manager.get_snowflake_connection('Nonexistent')
        
        self.assertIsNone(result)
    
    def test_list_datasets_success(self):
        """Test successful datasets listing"""
        mock_response = {
            'entries': [
                {'datasetId': '1', 'path': 'DB1.SCHEMA1.TABLE1'},
                {'datasetId': '2', 'path': 'DB1.SCHEMA2.TABLE2'}
            ]
        }
        self.mock_client.get.return_value = mock_response
        
        result = self.snowflake_manager.list_datasets('connection_id')
        
        self.assertEqual(len(result), 2)
        self.mock_client.get.assert_called_once_with('/v2/datasets', params={'connectionId': 'connection_id'})
    
    def test_list_datasets_no_connection_filter(self):
        """Test datasets listing without connection filter"""
        mock_response = {'entries': []}
        self.mock_client.get.return_value = mock_response
        
        self.snowflake_manager.list_datasets()
        
        self.mock_client.get.assert_called_once_with('/v2/datasets', params={})
    
    def test_get_snowflake_objects_success(self):
        """Test successful Snowflake objects organization"""
        mock_connection = {'connectionId': 'conn_id', 'name': 'Test Connection'}
        mock_datasets = [
            {'datasetId': 'ds1', 'path': 'PROD.FACTS.TABLE1', 'type': 'table'},
            {'datasetId': 'ds2', 'path': 'PROD.FACTS.TABLE2', 'type': 'table'},
            {'datasetId': 'ds3', 'path': 'PROD.DIM.TABLE3', 'type': 'table'},
            {'datasetId': 'ds4', 'path': 'DEV.FACTS.TABLE4', 'type': 'table'}
        ]
        
        with patch.object(self.snowflake_manager, 'get_snowflake_connection', return_value=mock_connection):
            with patch.object(self.snowflake_manager, 'list_datasets', return_value=mock_datasets):
                result = self.snowflake_manager.get_snowflake_objects()
        
        # Verify structure
        self.assertIn('connection', result)
        self.assertIn('databases', result)
        self.assertEqual(result['connection']['name'], 'Test Connection')
        
        # Check PROD database
        self.assertIn('PROD', result['databases'])
        prod_db = result['databases']['PROD']
        self.assertIn('schemas', prod_db)
        self.assertIn('FACTS', prod_db['schemas'])
        self.assertIn('DIM', prod_db['schemas'])
        
        # Check FACTS schema under PROD
        facts_schema = prod_db['schemas']['FACTS']
        self.assertIn('tables', facts_schema)
        self.assertEqual(len(facts_schema['tables']), 2)
        
        # Check table details
        table1 = facts_schema['tables'][0]
        self.assertEqual(table1['name'], 'TABLE1')
        self.assertEqual(table1['dataset_id'], 'ds1')
        self.assertEqual(table1['full_path'], 'PROD.FACTS.TABLE1')
    
    def test_get_snowflake_objects_connection_not_found(self):
        """Test Snowflake objects when connection not found"""
        with patch.object(self.snowflake_manager, 'get_snowflake_connection', return_value=None):
            result = self.snowflake_manager.get_snowflake_objects('Nonexistent')
        
        self.assertIn('error', result)
        self.assertEqual(result['error'], 'Snowflake connection not found')
    
    def test_grant_database_access_success(self):
        """Test successful database access grant"""
        # Mock user lookup
        mock_user = {'memberId': 'user123', 'email': 'test@test.com'}
        
        # Mock Snowflake objects
        mock_objects = {
            'databases': {
                'PROD': {
                    'schemas': {
                        'FACTS': {
                            'tables': [
                                {'name': 'TABLE1', 'dataset_id': 'ds1'},
                                {'name': 'TABLE2', 'dataset_id': 'ds2'}
                            ]
                        }
                    }
                }
            }
        }
        
        with patch.object(self.snowflake_manager.user_manager, 'get_member_by_email', return_value=mock_user):
            with patch.object(self.snowflake_manager, 'get_snowflake_objects', return_value=mock_objects):
                with patch.object(self.snowflake_manager, '_grant_dataset_access', return_value={'success': True}):
                    result = self.snowflake_manager.grant_database_access('test@test.com', 'PROD')
        
        # Should have 2 successful grants (one per table)
        successful_grants = [r for r in result if r.get('success')]
        self.assertEqual(len(successful_grants), 2)
    
    def test_grant_database_access_user_not_found(self):
        """Test database access grant with user not found"""
        with patch.object(self.snowflake_manager.user_manager, 'get_member_by_email', return_value=None):
            with self.assertRaises(ValueError) as context:
                self.snowflake_manager.grant_database_access('nonexistent@test.com', 'PROD')
        
        self.assertIn('not found', str(context.exception))
    
    def test_grant_table_access_success(self):
        """Test successful table access grant"""
        mock_user = {'memberId': 'user123', 'email': 'test@test.com'}
        mock_objects = {
            'databases': {
                'PROD': {
                    'schemas': {
                        'FACTS': {
                            'tables': [
                                {'name': 'TARGET_TABLE', 'dataset_id': 'target_ds'}
                            ]
                        }
                    }
                }
            }
        }
        
        with patch.object(self.snowflake_manager.user_manager, 'get_member_by_email', return_value=mock_user):
            with patch.object(self.snowflake_manager, 'get_snowflake_objects', return_value=mock_objects):
                with patch.object(self.snowflake_manager, '_grant_dataset_access', return_value={'success': True}):
                    result = self.snowflake_manager.grant_table_access(
                        'test@test.com', 'PROD', 'FACTS', 'TARGET_TABLE'
                    )
        
        self.assertTrue(result['success'])
        self.assertEqual(result['dataset_id'], 'target_ds')
    
    def test_grant_table_access_table_not_found(self):
        """Test table access grant with table not found"""
        mock_user = {'memberId': 'user123', 'email': 'test@test.com'}
        mock_objects = {
            'databases': {
                'PROD': {
                    'schemas': {
                        'FACTS': {
                            'tables': []
                        }
                    }
                }
            }
        }
        
        with patch.object(self.snowflake_manager.user_manager, 'get_member_by_email', return_value=mock_user):
            with patch.object(self.snowflake_manager, 'get_snowflake_objects', return_value=mock_objects):
                result = self.snowflake_manager.grant_table_access(
                    'test@test.com', 'PROD', 'FACTS', 'NONEXISTENT_TABLE'
                )
        
        self.assertFalse(result['success'])
        self.assertIn('not found', result['error'])
    
    def test_bulk_grant_access_success(self):
        """Test successful bulk access grant"""
        users = ['user1@test.com', 'user2@test.com']
        objects = [
            {'database': 'PROD', 'schema': 'FACTS', 'table': 'TABLE1'},
            {'database': 'PROD', 'schema': 'DIM'}
        ]
        
        # Mock grant methods to return success
        with patch.object(self.snowflake_manager, 'grant_table_access', return_value={'success': True}):
            with patch.object(self.snowflake_manager, 'grant_schema_access', return_value=[{'success': True}]):
                result = self.snowflake_manager.bulk_grant_access(users, objects)
        
        # Should have 4 results (2 users × 2 objects)
        self.assertEqual(len(result), 4)

if __name__ == '__main__':
    unittest.main()