import boto3
import time
import logging
from botocore.exceptions import ClientError, NoCredentialsError, TokenRetrievalError
from django.conf import settings

logger = logging.getLogger('powerpulse.aws')

_session_cache = {'session': None, 'timestamp': 0, 'ttl': 600}  
_refresh_attempts = {'count': 0, 'last_attempt': 0}
_MAX_REFRESH_ATTEMPTS = 3
_REFRESH_COOLDOWN = 300  


def get_boto3_session():
   

    current_time = time.time()

    if _session_cache['session'] and (current_time - _session_cache['timestamp']) < _session_cache['ttl']:
        try:
            sts = _session_cache['session'].client('sts')
            sts.get_caller_identity() 
            return _session_cache['session']
        except (ClientError, NoCredentialsError, TokenRetrievalError) as e:
            if isinstance(e, ClientError):
                code = e.response['Error']['Code']
                if code in ['ExpiredToken', 'UnrecognizedClientException']:
                    logger.warning("Cached AWS session expired or invalid. Refreshing credentials...")
                else:
                    logger.warning(f"Cached session invalid ({code}). Refreshing...")
            else:
                logger.warning(f"Cached session invalid: {e}. Refreshing...")

            _clear_session_cache()

    if (_refresh_attempts['count'] >= _MAX_REFRESH_ATTEMPTS and
            (current_time - _refresh_attempts['last_attempt']) < _REFRESH_COOLDOWN):
        logger.error(f"Too many refresh attempts. Waiting {_REFRESH_COOLDOWN}s cooldown.")
        raise Exception("AWS credential refresh rate limited")

    return _create_new_session()


def _create_new_session():
    current_time = time.time()
    _refresh_attempts['count'] += 1
    _refresh_attempts['last_attempt'] = current_time

    try:
        session = boto3.Session(region_name=settings.AWS_REGION)
        sts = session.client('sts')
        identity = sts.get_caller_identity()

        creds = session.get_credentials()
        exp_time = getattr(creds, 'expiry_time', None)
        logger.info(f"AWS credentials detected - User: {identity.get('UserId', 'N/A')}, "
                    f"ARN: {identity.get('Arn', 'N/A')[:70]}...")
        if exp_time:
            logger.info(f"Credential Expiration: {exp_time}")

        _session_cache['session'] = session
        _session_cache['timestamp'] = current_time
        _refresh_attempts['count'] = 0
        _refresh_attempts['last_attempt'] = 0

        return session

    except (ClientError, NoCredentialsError, TokenRetrievalError) as e:
        logger.warning(f"Default AWS credentials failed: {e}")

        if hasattr(settings, 'AWS_ACCESS_KEY_ID') and settings.AWS_ACCESS_KEY_ID:
            try:
                logger.info("Attempting explicit credentials from Django settings...")
                session = boto3.Session(
                    aws_access_key_id=settings.AWS_ACCESS_KEY_ID,
                    aws_secret_access_key=settings.AWS_SECRET_ACCESS_KEY,
                    aws_session_token=getattr(settings, 'AWS_SESSION_TOKEN', None),
                    region_name=settings.AWS_REGION
                )
                sts = session.client('sts')
                identity = sts.get_caller_identity()
                logger.info(f"Explicit AWS credentials valid. ARN: {identity.get('Arn', 'N/A')}")

                _session_cache['session'] = session
                _session_cache['timestamp'] = current_time
                _refresh_attempts['count'] = 0
                _refresh_attempts['last_attempt'] = 0
                return session

            except Exception as settings_error:
                logger.error(f"Explicit credentials failed: {settings_error}")

        logger.error(f"All AWS credential sources failed. Original error: {e}")
        raise Exception(f"No valid AWS credentials found: {e}")


def _clear_session_cache():
    global _session_cache
    _session_cache['session'] = None
    _session_cache['timestamp'] = 0


def _get_aws_resource(service_name, **kwargs):
    max_retries = 2
    for attempt in range(max_retries):
        try:
            session = get_boto3_session()
            if service_name == 'dynamodb':
                return session.resource('dynamodb', **kwargs)
            return session.client(service_name, **kwargs)

        except (ClientError, NoCredentialsError, TokenRetrievalError) as e:
            code = getattr(e, 'response', {}).get('Error', {}).get('Code', 'Unknown')
            logger.warning(f"AWS {service_name} access failed (attempt {attempt + 1}/{max_retries}): {e} [{code}]")
            if attempt == max_retries - 1:
                logger.error(f"Failed to get {service_name} after {max_retries} attempts.")
                raise
            _clear_session_cache()
            time.sleep(1)


def get_dynamodb_resource():
    return _get_aws_resource('dynamodb')


def get_dynamodb_client():
    return _get_aws_resource('dynamodb')


def get_s3_client():
    return _get_aws_resource('s3')


def get_sns_client():
    return _get_aws_resource('sns')


def get_sqs_client():
    return _get_aws_resource('sqs')


def get_cognito_client():
    return _get_aws_resource('cognito-idp')


def get_lambda_client():
    return _get_aws_resource('lambda')


def refresh_credentials():
    logger.info("Manually refreshing AWS credentials...")
    _clear_session_cache()
    _refresh_attempts['count'] = 0
    _refresh_attempts['last_attempt'] = 0
    return get_boto3_session()


def is_aws_available():
    try:
        session = get_boto3_session()
        sts = session.client('sts')
        sts.get_caller_identity()
        return True
    except Exception as e:
        logger.error(f"AWS not available: {e}")
        return False


def get_aws_identity():
    try:
        session = get_boto3_session()
        sts = session.client('sts')
        return sts.get_caller_identity()
    except Exception as e:
        logger.error(f"Could not get AWS identity: {e}")
        return None
        