from typing import List, Dict, Any, Optional
from sigma_api_client import SigmaAPIClient
from user_management import UserManager, PermissionManager, GranteeType, PermissionType

class SnowflakeObjectManager:
    """
    Manages Snowflake objects accessible through Sigma and user access grants
    """
    
    def __init__(self, api_client: SigmaAPIClient):
        self.api = api_client
        self.user_manager = UserManager(api_client)
        self.permission_manager = PermissionManager(api_client)
    
    def list_connections(self) -> List[Dict[str, Any]]:
        """
        List all data connections in Sigma
        """
        try:
            response = self.api.get('/v2/connections')
            return response.get('entries', [])
        except Exception as e:
            print(f"Error listing connections: {e}")
            return []
    
    def get_snowflake_connection(self, connection_name: str = None) -> Optional[Dict[str, Any]]:
        """
        Get Snowflake connection details
        
        Args:
            connection_name: Optional name of specific connection to find
        """
        connections = self.list_connections()
        
        for conn in connections:
            if conn.get('type', '').lower() == 'snowflake':
                if connection_name is None or conn.get('name') == connection_name:
                    return conn
        
        return None
    
    def list_datasets(self, connection_id: str = None) -> List[Dict[str, Any]]:
        """
        List datasets (tables/views) available through connections
        
        Args:
            connection_id: Optional connection ID to filter by
        """
        try:
            params = {}
            if connection_id:
                params['connectionId'] = connection_id
            
            response = self.api.get('/v2/datasets', params=params)
            return response.get('entries', [])
        except Exception as e:
            print(f"Error listing datasets: {e}")
            return []
    
    def get_snowflake_objects(self, connection_name: str = None) -> Dict[str, Any]:
        """
        Get organized list of Snowflake objects (databases, schemas, tables)
        
        Args:
            connection_name: Name of Snowflake connection
        """
        # Get Snowflake connection
        snowflake_conn = self.get_snowflake_connection(connection_name)
        if not snowflake_conn:
            return {'error': 'Snowflake connection not found'}
        
        connection_id = snowflake_conn['connectionId']
        
        # Get datasets for this connection
        datasets = self.list_datasets(connection_id)
        
        # Organize by database and schema
        organized_objects = {
            'connection': snowflake_conn,
            'databases': {}
        }
        
        for dataset in datasets:
            # Parse dataset path (typically database.schema.table format)
            path_parts = dataset.get('path', '').split('.')
            
            if len(path_parts) >= 3:
                database = path_parts[0]
                schema = path_parts[1]
                table = path_parts[2]
                
                if database not in organized_objects['databases']:
                    organized_objects['databases'][database] = {'schemas': {}}
                
                if schema not in organized_objects['databases'][database]['schemas']:
                    organized_objects['databases'][database]['schemas'][schema] = {'tables': []}
                
                organized_objects['databases'][database]['schemas'][schema]['tables'].append({
                    'name': table,
                    'dataset_id': dataset.get('datasetId'),
                    'full_path': dataset.get('path'),
                    'type': dataset.get('type', 'table')
                })
        
        return organized_objects
    
    def grant_database_access(self, user_email: str, database_name: str, 
                             permission: PermissionType = PermissionType.VIEW) -> List[Dict[str, Any]]:
        """
        Grant user access to all objects in a Snowflake database
        
        Args:
            user_email: Email of the user to grant access
            database_name: Name of the database
            permission: Permission level to grant
        """
        user = self.user_manager.get_member_by_email(user_email)
        if not user:
            raise ValueError(f"User with email {user_email} not found")
        
        user_id = user['memberId']
        
        # Get all Snowflake objects
        snowflake_objects = self.get_snowflake_objects()
        
        if 'error' in snowflake_objects:
            return [{'error': snowflake_objects['error']}]
        
        results = []
        
        # Find the database and grant access to all its tables
        if database_name in snowflake_objects['databases']:
            db_info = snowflake_objects['databases'][database_name]
            
            for schema_name, schema_info in db_info['schemas'].items():
                for table in schema_info['tables']:
                    try:
                        grant_result = self._grant_dataset_access(
                            dataset_id=table['dataset_id'],
                            user_id=user_id,
                            permission=permission
                        )
                        
                        results.append({
                            'success': True,
                            'object': f"{database_name}.{schema_name}.{table['name']}",
                            'dataset_id': table['dataset_id'],
                            'grant': grant_result
                        })
                        
                    except Exception as e:
                        results.append({
                            'success': False,
                            'object': f"{database_name}.{schema_name}.{table['name']}",
                            'error': str(e)
                        })
        else:
            results.append({
                'success': False,
                'error': f"Database {database_name} not found"
            })
        
        return results
    
    def grant_schema_access(self, user_email: str, database_name: str, schema_name: str,
                           permission: PermissionType = PermissionType.VIEW) -> List[Dict[str, Any]]:
        """
        Grant user access to all objects in a Snowflake schema
        
        Args:
            user_email: Email of the user to grant access
            database_name: Name of the database
            schema_name: Name of the schema
            permission: Permission level to grant
        """
        user = self.user_manager.get_member_by_email(user_email)
        if not user:
            raise ValueError(f"User with email {user_email} not found")
        
        user_id = user['memberId']
        
        # Get all Snowflake objects
        snowflake_objects = self.get_snowflake_objects()
        
        if 'error' in snowflake_objects:
            return [{'error': snowflake_objects['error']}]
        
        results = []
        
        # Find the schema and grant access to all its tables
        if (database_name in snowflake_objects['databases'] and 
            schema_name in snowflake_objects['databases'][database_name]['schemas']):
            
            schema_info = snowflake_objects['databases'][database_name]['schemas'][schema_name]
            
            for table in schema_info['tables']:
                try:
                    grant_result = self._grant_dataset_access(
                        dataset_id=table['dataset_id'],
                        user_id=user_id,
                        permission=permission
                    )
                    
                    results.append({
                        'success': True,
                        'object': f"{database_name}.{schema_name}.{table['name']}",
                        'dataset_id': table['dataset_id'],
                        'grant': grant_result
                    })
                    
                except Exception as e:
                    results.append({
                        'success': False,
                        'object': f"{database_name}.{schema_name}.{table['name']}",
                        'error': str(e)
                    })
        else:
            results.append({
                'success': False,
                'error': f"Schema {database_name}.{schema_name} not found"
            })
        
        return results
    
    def grant_table_access(self, user_email: str, database_name: str, schema_name: str, 
                          table_name: str, permission: PermissionType = PermissionType.VIEW) -> Dict[str, Any]:
        """
        Grant user access to a specific Snowflake table
        
        Args:
            user_email: Email of the user to grant access
            database_name: Name of the database
            schema_name: Name of the schema
            table_name: Name of the table
            permission: Permission level to grant
        """
        user = self.user_manager.get_member_by_email(user_email)
        if not user:
            raise ValueError(f"User with email {user_email} not found")
        
        user_id = user['memberId']
        
        # Get all Snowflake objects
        snowflake_objects = self.get_snowflake_objects()
        
        if 'error' in snowflake_objects:
            return {'error': snowflake_objects['error']}
        
        # Find the specific table
        if (database_name in snowflake_objects['databases'] and 
            schema_name in snowflake_objects['databases'][database_name]['schemas']):
            
            schema_info = snowflake_objects['databases'][database_name]['schemas'][schema_name]
            
            for table in schema_info['tables']:
                if table['name'] == table_name:
                    try:
                        grant_result = self._grant_dataset_access(
                            dataset_id=table['dataset_id'],
                            user_id=user_id,
                            permission=permission
                        )
                        
                        return {
                            'success': True,
                            'object': f"{database_name}.{schema_name}.{table_name}",
                            'dataset_id': table['dataset_id'],
                            'grant': grant_result
                        }
                        
                    except Exception as e:
                        return {
                            'success': False,
                            'object': f"{database_name}.{schema_name}.{table_name}",
                            'error': str(e)
                        }
        
        return {
            'success': False,
            'error': f"Table {database_name}.{schema_name}.{table_name} not found"
        }
    
    def bulk_grant_access(self, user_emails: List[str], objects: List[Dict[str, str]], 
                         permission: PermissionType = PermissionType.VIEW) -> List[Dict[str, Any]]:
        """
        Grant multiple users access to multiple Snowflake objects
        
        Args:
            user_emails: List of user email addresses
            objects: List of objects with 'database', 'schema', and optionally 'table' keys
            permission: Permission level to grant
        """
        results = []
        
        for user_email in user_emails:
            for obj in objects:
                database = obj['database']
                schema = obj.get('schema')
                table = obj.get('table')
                
                try:
                    if table:
                        # Grant table access
                        result = self.grant_table_access(user_email, database, schema, table, permission)
                        results.append({
                            'user_email': user_email,
                            'object_type': 'table',
                            'object': f"{database}.{schema}.{table}",
                            'result': result
                        })
                    elif schema:
                        # Grant schema access
                        schema_results = self.grant_schema_access(user_email, database, schema, permission)
                        results.append({
                            'user_email': user_email,
                            'object_type': 'schema',
                            'object': f"{database}.{schema}",
                            'result': schema_results
                        })
                    else:
                        # Grant database access
                        db_results = self.grant_database_access(user_email, database, permission)
                        results.append({
                            'user_email': user_email,
                            'object_type': 'database',
                            'object': database,
                            'result': db_results
                        })
                        
                except Exception as e:
                    results.append({
                        'user_email': user_email,
                        'object': obj,
                        'success': False,
                        'error': str(e)
                    })
        
        return results
    
    def _grant_dataset_access(self, dataset_id: str, user_id: str, 
                             permission: PermissionType) -> Dict[str, Any]:
        """
        Grant access to a specific dataset using the grants API
        
        Args:
            dataset_id: ID of the dataset
            user_id: ID of the user
            permission: Permission level to grant
        """
        # Use the grants API to grant dataset access
        return self.permission_manager.api.post('/v2/grants', {
            'resource': dataset_id,
            'grantee': user_id,
            'granteeType': GranteeType.USER.value,
            'permission': permission.value
        })