# SPDX-License-Identifier: Apache-2.0
# Copyright 2025 S3 Fetch Project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import boto3
import os
from botocore.exceptions import ClientError, NoCredentialsError
import logging

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class S3Fetcher:
    """
    A class to fetch objects from Amazon S3 using secure credential management.
    
    Credentials are loaded from environment variables or AWS credential files.
    """
    
    def __init__(self, aws_access_key_id=None, aws_secret_access_key=None, aws_region=None):
        """
        Initialize S3 client with secure credential management.
        
        Args:
            aws_access_key_id (str, optional): AWS Access Key ID. If not provided, 
                will use environment variable AWS_ACCESS_KEY_ID or AWS credential files.
            aws_secret_access_key (str, optional): AWS Secret Access Key. If not provided,
                will use environment variable AWS_SECRET_ACCESS_KEY or AWS credential files.
            aws_region (str, optional): AWS region. If not provided, will use 
                environment variable AWS_DEFAULT_REGION or default to 'us-east-1'.
        """
        # Get credentials from environment variables or AWS credential files
        self.aws_access_key_id = os.getenv('AWS_ACCESS_KEY_ID')
        self.aws_secret_access_key = os.getenv('AWS_SECRET_ACCESS_KEY')
        self.aws_region = os.getenv('AWS_DEFAULT_REGION', 'us-east-1')
        
        # Initialize S3 client with secure credential management
        try:
            # If explicit credentials are provided, use them
            if self.aws_access_key_id and self.aws_secret_access_key:
                self.s3_client = boto3.client(
                    's3',
                    aws_access_key_id=self.aws_access_key_id,
                    aws_secret_access_key=self.aws_secret_access_key,
                    region_name=self.aws_region
                )
                logger.info("S3 client initialized successfully with environment credentials")
            else:
                # Use default credential chain (IAM roles, credential files, etc.)
                self.s3_client = boto3.client('s3', region_name=self.aws_region)
                logger.info("S3 client initialized successfully using AWS credential chain")
                
        except NoCredentialsError:
            logger.error("No AWS credentials found. Please configure credentials using one of the following methods:")
            logger.error("1. Set environment variables: AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY")
            logger.error("2. Configure AWS CLI: aws configure")
            logger.error("3. Use IAM roles (for EC2 instances)")
            logger.error("4. Use AWS credential files (~/.aws/credentials)")
            raise
        except Exception as e:
            logger.error(f"Failed to initialize S3 client: {e}")
            raise
    
    def list_buckets(self):
        """
        List all S3 buckets accessible with the current credentials.
        
        Returns:
            list: List of bucket names
        """
        try:
            response = self.s3_client.list_buckets()
            buckets = [bucket['Name'] for bucket in response['Buckets']]
            logger.info(f"Found {len(buckets)} buckets")
            return buckets
        except ClientError as e:
            logger.error(f"Error listing buckets: {e}")
            return []
    
    def list_objects(self, bucket_name, prefix=""):
        """
        List objects in a specific S3 bucket.
        
        Args:
            bucket_name (str): Name of the S3 bucket
            prefix (str): Prefix to filter objects (optional)
            
        Returns:
            list: List of object keys
        """
        try:
            response = self.s3_client.list_objects_v2(
                Bucket=bucket_name,
                Prefix=prefix
            )
            
            if 'Contents' in response:
                objects = [obj['Key'] for obj in response['Contents']]
                logger.info(f"Found {len(objects)} objects in bucket '{bucket_name}'")
                return objects
            else:
                logger.info(f"No objects found in bucket '{bucket_name}'")
                return []
                
        except ClientError as e:
            logger.error(f"Error listing objects in bucket '{bucket_name}': {e}")
            return []
    
    def download_file(self, bucket_name, object_key, local_file_path):
        """
        Download a file from S3 to local filesystem.
        
        Args:
            bucket_name (str): Name of the S3 bucket
            object_key (str): S3 object key (file path in bucket)
            local_file_path (str): Local path where file will be saved
            
        Returns:
            bool: True if download successful, False otherwise
        """
        try:
            self.s3_client.download_file(bucket_name, object_key, local_file_path)
            logger.info(f"Successfully downloaded '{object_key}' to '{local_file_path}'")
            return True
        except ClientError as e:
            logger.error(f"Error downloading file '{object_key}': {e}")
            return False
    
    def get_object_content(self, bucket_name, object_key):
        """
        Get the content of an S3 object as string.
        
        Args:
            bucket_name (str): Name of the S3 bucket
            object_key (str): S3 object key
            
        Returns:
            str: Content of the object, or None if error
        """
        try:
            response = self.s3_client.get_object(Bucket=bucket_name, Key=object_key)
            content = response['Body'].read().decode('utf-8')
            logger.info(f"Successfully retrieved content from '{object_key}'")
            return content
        except ClientError as e:
            logger.error(f"Error getting object content '{object_key}': {e}")
            return None
        except UnicodeDecodeError as e:
            logger.error(f"Error decoding content from '{object_key}': {e}")
            return None
    
    def get_object_metadata(self, bucket_name, object_key):
        """
        Get metadata for an S3 object.
        
        Args:
            bucket_name (str): Name of the S3 bucket
            object_key (str): S3 object key
            
        Returns:
            dict: Object metadata, or None if error
        """
        try:
            response = self.s3_client.head_object(Bucket=bucket_name, Key=object_key)
            metadata = {
                'ContentLength': response.get('ContentLength'),
                'ContentType': response.get('ContentType'),
                'LastModified': response.get('LastModified'),
                'ETag': response.get('ETag'),
                'Metadata': response.get('Metadata', {})
            }
            logger.info(f"Successfully retrieved metadata for '{object_key}'")
            return metadata
        except ClientError as e:
            logger.error(f"Error getting object metadata '{object_key}': {e}")
            return None


def main():
    """
    Main function to demonstrate S3 fetching functionality.
    """
    # Initialize S3 fetcher with hardcoded credentials
    s3_fetcher = S3Fetcher()
    
    # Example usage
    print("=== S3 Fetch Demo ===")
    
    # List all buckets
    print("\n1. Listing all buckets:")
    buckets = s3_fetcher.list_buckets()
    for bucket in buckets:
        print(f"  - {bucket}")
    
    if buckets:
        # Use the first bucket for demonstration
        demo_bucket = buckets[0]
        print(f"\n2. Listing objects in bucket '{demo_bucket}':")
        
        # List objects in the first bucket
        objects = s3_fetcher.list_objects(demo_bucket)
        for obj in objects[:10]:  # Show first 10 objects
            print(f"  - {obj}")
        
        if objects:
            # Get metadata for the first object
            demo_object = objects[0]
            print(f"\n3. Getting metadata for '{demo_object}':")
            metadata = s3_fetcher.get_object_metadata(demo_bucket, demo_object)
            if metadata:
                for key, value in metadata.items():
                    print(f"  {key}: {value}")
            
            # Try to get content (for text files)
            print(f"\n4. Getting content preview for '{demo_object}':")
            content = s3_fetcher.get_object_content(demo_bucket, demo_object)
            if content:
                # Show first 200 characters
                preview = content[:200] + "..." if len(content) > 200 else content
                print(f"  Content preview: {preview}")
            
            # Download file example
            local_path = f"downloaded_{demo_object.replace('/', '_')}"
            print(f"\n5. Downloading '{demo_object}' to '{local_path}':")
            success = s3_fetcher.download_file(demo_bucket, demo_object, local_path)
            if success:
                print(f"  File downloaded successfully to {local_path}")
            else:
                print("  Download failed")
    else:
        print("No buckets found or access denied")


if __name__ == "__main__":
    main()