########################################################################################################################
# IMPORTS

import io
import logging
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional

import boto3
from botocore.signers import CloudFrontSigner
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import padding

########################################################################################################################
# CLASSES

logger = logging.getLogger(__name__)


class AWSInterface:
    _COMMON_CONFIG_PROFILE_NAME = "common-config"

    def __init__(self, config) -> None:
        self.profiles: List[Dict[str, Any]] = []
        self.config = config
        self.cloudfront_url = None
        self.key_pair_id = None
        self.private_key_path = None
        self._cached_signer: Optional[CloudFrontSigner] = None

        for section in getattr(self.config, "sections", lambda: [])():
            if section.startswith("aws:"):
                profile_name = section.split(":", 1)[1]

                if profile_name == self._COMMON_CONFIG_PROFILE_NAME:
                    self.cloudfront_url = self.config[section].get("cloudfront_url")
                    self.key_pair_id = self.config[section].get("key_pair_id")

                    pk_path = self.config[section].get("private_key_path")
                    self.private_key_path = Path(pk_path).expanduser() if pk_path else None
                    continue

                bucket_value = self.config[section].get("buckets", "")
                buckets = [b.strip() for b in bucket_value.split(",") if b.strip()]
                session = boto3.Session(profile_name=profile_name)

                self.profiles.append(
                    {
                        "profile": profile_name,
                        "buckets": buckets,
                        "session": session,
                    }
                )

        if not self.profiles:
            logger.warning("No AWS profiles found in config file")

        self.current_profile: Optional[Dict[str, Any]] = self.profiles[0] if self.profiles else None
        self._update_resources()
        self._initialize_cloudfront()

    def _initialize_cloudfront(self) -> None:
        # Initialize CloudFront Signer once to avoid repeated disk I/O
        if self.cloudfront_url and self.key_pair_id and self.private_key_path:
            try:
                with open(self.private_key_path, "rb") as key_file:
                    private_key = serialization.load_pem_private_key(
                        key_file.read(),
                        password=None,
                    )

                # CloudFront requires SHA-1 for RSA signatures in the standard signer context
                def rsa_signer(message):
                    return private_key.sign(message, padding.PKCS1v15(), hashes.SHA1())  # noqa: S303

                self._cached_signer = CloudFrontSigner(self.key_pair_id, rsa_signer)
            except FileNotFoundError:
                logger.error(f"CloudFront private key file not found at: {self.private_key_path}")
            except (ValueError, TypeError) as e:  # Add other specific crypto errors if needed
                logger.error(
                    f"Failed to load CloudFront private key. Check if the file is a valid PEM-encoded private key. Error: {e}"
                )
            except Exception as e:
                logger.error(f"An unexpected error occurred while initializing CloudFront signer: {e}")

    def _update_resources(self) -> None:
        """Refresh S3 resources for the current profile and set default bucket (first in list)"""
        if self.current_profile:
            self.s3 = self.current_profile["session"].resource("s3")
            self.s3_client = self.s3.meta.client
            buckets = self.current_profile.get("buckets", [])
            self.bucket = buckets[0] if buckets else None
        else:
            self.s3 = None
            self.s3_client = None
            self.bucket = None

    def switch_profile(self, profile_name: str) -> None:
        for profile in self.profiles:
            if profile["profile"] == profile_name:
                self.current_profile = profile
                self._update_resources()
                return
        logger.warning(f"Profile {profile_name} not found")

    def switch_bucket(self, bucket: str) -> None:
        if not self.current_profile:
            logger.warning("No current AWS profile to switch bucket on")
            return

        buckets = self.current_profile.get("buckets") or []
        if bucket not in buckets:
            logger.warning(f"Bucket {bucket} not found in profile {self.current_profile.get('profile')}")
            return

        self.bucket = bucket

    def switch_bucket_for_profile(self, profile_name: str, bucket: str) -> None:
        """
        Select a profile and then switch its active bucket.
        """
        for profile in self.profiles:
            if profile["profile"] == profile_name:
                self.current_profile = profile
                self._update_resources()  # sets default bucket & s3 clients
                self.switch_bucket(bucket)  # only sets self.bucket if valid
                return
        logger.warning(f"Profile {profile_name} not found")

    def get_bucket_url(self) -> Optional[str]:
        """Return active bucket URL."""
        if not self.bucket:
            logger.warning("No active bucket selected")
            return None
        region = self.s3_client.meta.region_name
        return f"https://{self.bucket}.s3.{region}.amazonaws.com"

    def get_cloudfront_url(self) -> Optional[str]:
        """Return cloudfront url if found in config under the aws:common-config profile"""
        return self.cloudfront_url

    def get_file(self, s3_path: str):
        if not self.bucket:
            logger.warning("No active bucket selected")
            return None
        try:
            return self.s3.Object(self.bucket, s3_path).get()
        except self.s3_client.exceptions.NoSuchKey:
            logger.info(f"{s3_path} does not exist")
            return None

    def file_exists(self, s3_path: str) -> bool:
        if not self.bucket:
            logger.warning("No active bucket selected")
            return False
        try:
            self.s3_client.head_object(Bucket=self.bucket, Key=s3_path)
            return True
        except self.s3_client.exceptions.NoSuchKey:
            return False
        except Exception as e:
            logger.error(f"Error checking existence of {s3_path}: {e}")
            raise

    def read_file_as_bytes(self, s3_path: str) -> Optional[io.BytesIO]:
        obj = self.get_file(s3_path)
        if not obj:
            return None
        return io.BytesIO(obj["Body"].read())

    def upload_file(self, local_path: str, s3_path: str, **kwargs) -> None:
        if not self.bucket:
            logger.warning("No active bucket selected")
            return
        self.s3.Bucket(self.bucket).upload_file(local_path, s3_path, **kwargs)

    def sign_cloudfront_url(
        self,
        s3_path: str,
        expiration_days: int = 90,
    ) -> Optional[str]:
        """
        Generate a CloudFront signed URL using the pre-cached RSA signer.
        """
        if not self._cached_signer:
            logger.warning("CloudFront signer not initialized. Check your aws:common-config.")
            return None

        try:
            resource_url = f"{self.cloudfront_url.rstrip('/')}/{s3_path.lstrip('/')}"

            # Botocore expects a datetime object for the expiration
            expire_date = datetime.now(timezone.utc) + timedelta(days=expiration_days)

            # generate_presigned_url handles Policy creation, Signature, and Base64 encoding
            return self._cached_signer.generate_presigned_url(resource_url, date_less_than=expire_date)

        except Exception as e:
            logger.error(f"Error signing CloudFront URL: {e}")
            return None
