# External imports
from pathlib import Path
import os
import boto3
import tempfile
from botocore.exceptions import NoCredentialsError, ClientError
import logging
import tempfile
from dotenv import load_dotenv
import os
from fastapi import UploadFile
import time

# Load the environment variables
load_dotenv()


class S3FileManager:
    """
    A class to interact with AWS S3.

    Attributes:
    -----------
    aws_access_key_id: str
        The AWS access key ID.
    aws_secret_access_key: str
        The AWS secret access key.
    bucket_name: str
        The name of the bucket.
    s3_client: S3 client
        The S3 client.


    Methods:
    --------
    upload_file(file_path, key)
        Upload a file to S3.
    upload_temp_file(file, key)
        Upload a temporary file to S3.
    list_files(key)
        List all files in the S3 bucket with the given key.
    download_file(key, download_path)
        Download a file from S3.
    delete_file(key)
        Delete a file from S3.
    upload_file_from_bytes(data, key)
        Upload a file to S3 from bytes.
    download_file_to_bytes(key)
        Download a file from S3 to bytes.
    get_object(key)
        Get an object from S3.

    """

    def __init__(self):
        """
        Constructor for the S3FileManager class.
        """

        # Initialize AWS credentials and S3 client
        self.aws_access_key_id = os.environ.get("AWS_ACCESS_KEY")
        self.aws_secret_access_key = os.environ.get("AWS_SECRET_KEY")
        self.bucket_name = os.environ.get("AWS_BUCKET_NAME")
        self.s3_client = boto3.client(
            's3',
            aws_access_key_id=self.aws_access_key_id,
            aws_secret_access_key=self.aws_secret_access_key
        )

    def upload_file_obj(self, file_obj, key):
        """
        Upload a file object to S3

        Args:
        file_obj: file object - file object to be uploaded
        key: str - key to be used in the S3 bucket

        Returns:
        bool: True if the file was uploaded successfully, False otherwise
        """
        try:
            self.s3_client.upload_fileobj(file_obj, self.bucket_name, key)
            self.make_object_public(key)
            return True
        except FileNotFoundError:
            logging.error("The file was not found")
            return False
        except NoCredentialsError:
            logging.error("Credentials not available")
            return False
        except ClientError as e:
            logging.error(e)
            return False
        
    async def upload_file_from_frontend(self, file: UploadFile, key):
        try:
            file_content = await file.read()
            self.s3_client.put_object(Bucket=self.bucket_name, Key=key, Body=file_content)
            self.make_object_public(key)
            return True
        except FileNotFoundError:
            logging.error("The file was not found")
            return False
        except NoCredentialsError:
            logging.error("Credentials not available")
            return False
        except ClientError as e:
            logging.error(e)
            return False
        

    async def upload_file(self, file_path, key, content_type=None):
        """
        Upload a file to S3

        Args:
        file_path: str - path to the file to be uploaded
        key: str - key to be used in the S3 bucket

        Returns:
        bool: True if the file was uploaded successfully, False otherwise
        """
        try:
            if content_type:
                self.s3_client.upload_file(file_path, self.bucket_name, key, ExtraArgs={'ContentType': content_type})
            else:
                self.s3_client.upload_file(file_path, self.bucket_name, key)
            self.make_object_public(key)
            return True
        except FileNotFoundError:
            logging.error("The file was not found")
            return False
        except NoCredentialsError:
            logging.error("Credentials not available")
            return False
        except ClientError as e:
            logging.error(e)
            return False

    def upload_temp_file(self, file, key):
        """
        Upload a temporary file to S3

        Args:
        file: bytes - file data to be uploaded
        key: str - key to be used in the S3 bucket

        Returns:
        bool: True if the file was uploaded successfully, False otherwise
        """
        try:
            with tempfile.NamedTemporaryFile(delete=False) as temp_file:
                temp_file.write(file)
                temp_file.close()
                self.upload_file(temp_file.name, key, content_type='application/octet-stream')
                os.unlink(temp_file.name)
            return True
        except NoCredentialsError:
            logging.error("Credentials not available")
            return False
        except ClientError as e:
            logging.error(e)
            return False
        
    def make_object_public(self, key):
        try:
            self.s3_client.put_object_acl(
                ACL='public-read', Bucket=self.bucket_name, Key=key)
            logging.info(
                f"Object '{key}' made public in S3 bucket '{self.bucket_name}'")
            return True
        except NoCredentialsError:
            logging.error("AWS credentials not available or incorrect.")
            return False
        except Exception as e:
            logging.error(f"An error occurred: make_object_public: {e}")
            return False
  

    def list_files(self, key):
        """
        List all files in the S3 bucket with the given key

        Args:
        key: str - key of the files in the S3 bucket

        Returns:
        list: List of files in the S3 bucket with the given key
        """
        try:
            response = self.s3_client.list_objects_v2(
                Bucket=self.bucket_name, Prefix=key)
            return response.get("Contents")
        except NoCredentialsError:
            logging.error("Credentials not available")
            return False
        except ClientError as e:
            logging.error(e)
            return False

    def download_file(self, key, download_path):
        """
        Download a file from S3

        Args:
        key: str - key of the file in the S3 bucket
        download_path: str - path to download the file

        Returns:
        bool: True if the file was downloaded successfully, False otherwise
        """
        try:
            with open(download_path, 'wb') as f:
                self.s3_client.download_fileobj(self.bucket_name, key, f)
            return True
        except NoCredentialsError:
            logging.error("Credentials not available")
            return False
        except ClientError as e:
            logging.error(e)
            return False

    def delete_file(self, key):
        """
        Delete a file from S3

        Args:
        key: str - key of the file in the S3 bucket

        Returns:
        bool: True if the file was deleted successfully, False otherwise
        """
        try:
            self.s3_client.delete_object(Bucket=self.bucket_name, Key=key)
            return True
        except NoCredentialsError:
            logging.error("Credentials not available")
            return False
        except ClientError as e:
            logging.error(e)
            return False

    def upload_file_from_bytes(self, data, key):
        """
        Upload a file to S3 from bytes

        Args:
        data: bytes - data to be uploaded
        key: str - key to be used in the S3 bucket

        Returns:
        bool: True if the file was uploaded successfully, False otherwise
        """
        try:
            with tempfile.NamedTemporaryFile(delete=False) as temp_file:
                temp_file.write(data)
                temp_file.close()
                self.upload_file(temp_file.name, key)
                os.unlink(temp_file.name)
            return True
        except NoCredentialsError:
            logging.error("Credentials not available")
            return False
        except ClientError as e:
            logging.error(e)
            return False

    def download_file_to_bytes(self, key):
        """
        Download a file from S3 to bytes

        Args:
        key: str - key of the file in the S3 bucket

        Returns:
        bytes: data of the file
        """
        try:
            with tempfile.NamedTemporaryFile(delete=False) as temp_file:
                temp_file.close()
                self.download_file(key, temp_file.name)
                with open(temp_file.name, 'rb') as f:
                    data = f.read()
                os.unlink(temp_file.name)
            return data
        except NoCredentialsError:
            logging.error("Credentials not available")
            return False
        except ClientError as e:
            logging.error(e)
            return False

    def get_object(self, key):
        """
        Get an object from S3

        Args:
        key: str - key of the object in the S3 bucket

        Returns:
        bytes: data of the object
        """
        try:
            response = self.s3_client.get_object(
                Bucket=self.bucket_name, Key=key)
            return response
        except NoCredentialsError:
            logging.error("Credentials not available")
            return None
        except ClientError as e:
            logging.error(e)
            return None

    def upload_directory(self, directory_path, key):
        """
        Upload a directory to S3

        Args:
        directory_path: str - path to the directory to be uploaded
        key: str - key to be used in the S3 bucket

        Returns:
        bool: True if the directory was uploaded successfully, False otherwise
        """
        try:
            for root, dirs, files in os.walk(directory_path):
                for file in files:
                    file_path = os.path.join(root, file)
                    s3_key = key + file_path[len(directory_path):]
                    self.upload_file(file_path, s3_key)
            return True
        except NoCredentialsError:
            logging.error("Credentials not available")
            return False
        except ClientError as e:
            logging.error(e)
            return False
        

    def download_directory(self, key, download_path):
        """
        Download a directory from S3

        Args:
        key: str - key of the directory in the S3 bucket
        download_path: str - path to download the directory

        Returns:
        bool: True if the directory was downloaded successfully, False otherwise
        """
        try:
            for obj in self.list_files(key):
                file_key = obj['Key']
                file_path = os.path.join(download_path, file_key[len(key):])
                self.download_file(file_key, file_path)
            return True
        except NoCredentialsError:
            logging.error("Credentials not available")
            return False
        except ClientError as e:
            logging.error(e)
            return False

    async def save_mp3_and_upload(self, audio_data, key):
        """
        Save the MP3 data to a file and upload it to S3.
        
        Args:
        audio_data: bytes - The audio data (MP3).
        key: str - The key under which to store the file in S3.
        
        Returns:
        bool: True if the upload succeeded, False otherwise.
        """
        try:
            # Define the folder where you want to save the file
            folder_path = "mp3_files"
            Path(folder_path).mkdir(parents=True, exist_ok=True)  # Create the folder if it doesn't exist
            
            # Create a unique filename using the current timestamp
            file_name = f"podcast_audio_{int(time.time())}.mp3"
            file_path = os.path.join(folder_path, file_name)
            
            # Write the audio data to an MP3 file
            with open(file_path, 'wb') as f:
                f.write(audio_data)

            # Upload the file asynchronously
            success = await self.async_upload_file(file_path, key)
            
            if success:
                logging.info(f"Successfully uploaded file to S3 with key: {key}")
                return True
            else:
                logging.error(f"Failed to upload file to S3 with key: {key}")
                return False

        except Exception as e:
            logging.error(f"Error saving and uploading MP3: {e}")
            return False  # Return False if there's an exception

        finally:
            # Clean up: Delete the file after uploading
            if 'file_path' in locals() and os.path.exists(file_path):
                os.remove(file_path)


    def async_upload_file(self, file_path, key):
        """
        Upload a file to S3.

        Args:
        file_path (str): Path to the file on your local machine.
        key (str): Key (path) in the S3 bucket where the file will be stored.

        Returns:
        bool: True if the file was uploaded successfully, False otherwise.
        """
        try:
            # Upload the file to S3
            response = self.s3_client.upload_file(file_path, self.bucket_name, key)
            logging.info(f"Successfully uploaded file to S3 with key: {key}")
            return True
        except Exception as e:
            logging.error(f"Error uploading file to S3: {e}")
            return False

        
    def delete_folder(self, folder_path):
        """
        Delete a folder and all its contents from S3.
        
        Args:
        folder_path (str): Path to the folder in the S3 bucket.
        
        Returns:
        bool: True if the folder was deleted successfully, False otherwise.
        """
        try:
            # List all objects in the folder
            objects = self.s3_client.list_objects_v2(Bucket=self.bucket_name, Prefix=folder_path)
            
            # Delete each object in the folder
            for obj in objects.get('Contents', []):
                self.s3_client.delete_object(Bucket=self.bucket_name, Key=obj['Key'])
            
            logging.info(f"Successfully deleted folder from S3 with path: {folder_path}")
            return True
        except NoCredentialsError:
            logging.error("Credentials not available")
            return False
        except ClientError as e:
            logging.error(e)
            return False
        
    def ensure_folder_exists(self, folder_key):
        """
        Ensure that a folder exists in S3 (S3 doesn't have actual folders, but this creates a placeholder)
        
        Args:
            folder_key (str): The folder key to create (should end with a '/')
        
        Returns:
            bool: True if successful
        """
        try:
            # Ensure the folder key ends with a slash
            if not folder_key.endswith('/'):
                folder_key += '/'
                
            # Create an empty object with the folder key to represent the folder
            self.s3_client.put_object(
                Bucket=self.bucket_name,
                Key=folder_key,
                Body=''
            )
            logging.info(f"Ensured folder exists: {folder_key}")
            return True
        except Exception as e:
            logging.error(f"Error ensuring folder exists: {e}")
            return False

    def copy_file(self, source_key, destination_key):
        """
        Copy a file in S3, ensuring the destination directory structure exists

        Args:
        source_key: str - key of the source file in the S3 bucket
        destination_key: str - key of the destination file in the S3 bucket

        Returns:
        bool: True if the file was copied successfully, False otherwise
        """
        try:
            # Ensure the destination directory exists by creating a folder marker if needed
            destination_folder = '/'.join(destination_key.split('/')[:-1])
            if destination_folder:
                self.ensure_folder_exists(destination_folder)
                
            # Copy the object
            self.s3_client.copy_object(
                Bucket=self.bucket_name, 
                CopySource=f"{self.bucket_name}/{source_key}", 
                Key=destination_key
            )
            
            # Make the new object public
            self.make_object_public(destination_key)
            
            logging.info(f"Successfully copied file: {source_key} → {destination_key}")
            return True
        except NoCredentialsError:
            logging.error("Credentials not available")
            return False
        except ClientError as e:
            logging.error(f"Client error copying file {source_key} to {destination_key}: {e}")
            return False
        except Exception as e:
            logging.error(f"Unexpected error copying file {source_key} to {destination_key}: {e}")
            return False

    def copy_folder_structure(self, source_prefix, destination_prefix):
        """
        Create a matching folder structure in the destination based on the source structure
        
        Args:
            source_prefix (str): The source folder prefix
            destination_prefix (str): The destination folder prefix
        
        Returns:
            bool: True if successful
        """
        try:
            # List all "folders" in the source
            response = self.s3_client.list_objects_v2(
                Bucket=self.bucket_name,
                Prefix=source_prefix,
                Delimiter='/'
            )
            
            # Create the main destination folder
            self.ensure_folder_exists(destination_prefix)
            
            # Create each subfolder in the destination
            if 'CommonPrefixes' in response:
                for prefix in response['CommonPrefixes']:
                    source_subfolder = prefix['Prefix']
                    # Replace only the first occurrence of source_prefix with destination_prefix
                    destination_subfolder = source_subfolder.replace(source_prefix, destination_prefix, 1)
                    self.ensure_folder_exists(destination_subfolder)
                    
                    # Recursively copy subfolders
                    self.copy_folder_structure(source_subfolder, destination_subfolder)
                    
            return True
        except Exception as e:
            logging.error(f"Error copying folder structure: {e}")
            return False

    def copy_folder(self, source_prefix, destination_prefix):
        """
        Copy all files from one folder to another within the same S3 bucket.
        Also creates the destination folder structure if it doesn't exist.

        Args:
            source_prefix (str): The S3 folder path to copy from (e.g., "qu-lab-design/<old_id>").
            destination_prefix (str): The S3 folder path to copy to (e.g., "qu-lab-design/<new_id>").

        Returns:
            bool: True if all files were copied successfully, False otherwise.
        """
        try:
            logging.info(f"[S3 Cloning] Copying folder: {source_prefix} → {destination_prefix}")
            
            # First ensure destination folder structure exists
            self.copy_folder_structure(source_prefix, destination_prefix)
            
            # Get list of all objects in the source folder
            objects = self.s3_client.list_objects_v2(Bucket=self.bucket_name, Prefix=source_prefix)

            if 'Contents' not in objects:
                logging.warning(f"[S3 Cloning] No files found in source folder: {source_prefix}")
                return True  # Nothing to copy is still a success

            # Copy each object to the destination
            for obj in objects['Contents']:
                source_key = obj['Key']
                # Derive the destination key by replacing the source prefix with the destination prefix
                destination_key = source_key.replace(source_prefix, destination_prefix, 1)
                success = self.copy_file(source_key, destination_key)

                if success:
                    logging.info(f"[S3 Cloning] Copied: {source_key} → {destination_key}")
                else:
                    logging.error(f"[S3 Cloning] Failed to copy: {source_key}")
                    return False

            return True

        except Exception as e:
            logging.error(f"[S3 Cloning] Exception while copying folder: {e}")
            return False
