# 5x_aws_athena_auth_manager/auth.py
import os
from typing import Dict
import boto3
from botocore.exceptions import ClientError, NoCredentialsError

class AthenaConnectionManager:
    """
    Manages AWS Athena connections using environment variables.
    Required environment variables:
    - FIVEX_ATHENA_AWS_ACCESS_KEY: AWS access key ID
    - FIVEX_ATHENA_AWS_SECRET_KEY: AWS secret access key
    - FIVEX_ATHENA_AWS_REGION: AWS region (defaults to us-east-1)
    """
    
    def __init__(self):
        """Initialize the Athena connection manager using environment variables."""
        self._client = None
        self._validate_and_set_credentials()
    
    def _validate_and_set_credentials(self) -> None:
        """
        Validate and set AWS credentials from environment variables.
        
        Raises:
            ValueError: If required environment variables are missing
            ValueError: If credentials are empty or malformed
        """
        # Get credentials from environment
        access_key = os.getenv('FIVEX_ATHENA_AWS_ACCESS_KEY')
        secret_key = os.getenv('FIVEX_ATHENA_AWS_SECRET_KEY')
        region = os.getenv('FIVEX_ATHENA_AWS_REGION')

        # Check for missing credentials
        missing_vars = []
        if not access_key:
            missing_vars.append('FIVEX_ATHENA_AWS_ACCESS_KEY')
        if not secret_key:
            missing_vars.append('FIVEX_ATHENA_AWS_SECRET_KEY')
        if not region:
            missing_vars.append('FIVEX_ATHENA_AWS_REGION')    
            
        if missing_vars:
            raise ValueError(
                f"Missing required environment variables: {', '.join(missing_vars)}"
            )

        # Validate credential format
        if len(access_key.strip()) < 16:
            raise ValueError("AWS Access Key ID is invalid or malformed")
        if len(secret_key.strip()) < 16:
            raise ValueError("AWS Secret Access Key is invalid or malformed")

        self.credentials = {
            'aws_access_key_id': access_key,
            'aws_secret_access_key': secret_key,
            'region_name': region
        }

    def get_client(self) -> boto3.client:
        """
        Get or create an Athena client instance.
        
        Returns:
            boto3.client: Configured Athena client
            
        Raises:
            ClientError: If client creation fails due to invalid credentials
            ConnectionError: If connection to AWS fails
            Exception: For other unexpected errors
        """
        if self._client is None:
            try:
                self._client = boto3.client('athena', **self.credentials)
                # Test the connection immediately
                self._client.list_work_groups(MaxResults=1)
            except ClientError as e:
                error_code = e.response['Error']['Code']
                if error_code == 'InvalidAccessKeyId':
                    raise ValueError("Invalid AWS Access Key ID") from e
                elif error_code == 'SignatureDoesNotMatch':
                    raise ValueError("Invalid AWS Secret Access Key") from e
                elif error_code == 'UnrecognizedClientException':
                    raise ValueError("Invalid AWS credentials") from e
                else:
                    raise ClientError(
                        f"AWS Athena client error: {str(e)}",
                        e.operation_name
                    ) from e
            except NoCredentialsError as e:
                raise ValueError("No AWS credentials found") from e
            except Exception as e:
                raise ConnectionError(f"Failed to connect to AWS Athena: {str(e)}") from e
                
        return self._client

# Example usage
if __name__ == "__main__":
    try:
        # Create manager and get client
        manager = AthenaConnectionManager()
        client = manager.get_client()
        print("Successfully connected to AWS Athena")
        
    except ValueError as e:
        print(f"Configuration error: {str(e)}")
    except ClientError as e:
        print(f"AWS client error: {str(e)}")
    except ConnectionError as e:
        print(f"Connection error: {str(e)}")
    except Exception as e:
        print(f"Unexpected error: {str(e)}")