import logging
import boto3
from datetime import datetime
import os
import threading


class CloudWatchLogHandler(logging.Handler):
    _lock = threading.Lock()

    def __init__(self, log_group_name, log_stream_name):
        super().__init__()
        self.log_group_name = log_group_name
        self.log_stream_name = log_stream_name
        self.max_batch_size = 256000  # Max size for CloudWatch log event
        self.sequence_token = None
        self.cloudwatch_logs = boto3.client(
            'logs',
            aws_access_key_id=os.environ.get('CLOUDWATCH_ACCESS_KEY'),
            aws_secret_access_key=os.environ.get('CLOUDWATCH_SECRET_KEY'),
            region_name='ap-south-1'
        )
        self._initialize_log_stream()

    def _initialize_log_stream(self):
        # Ensure the log group exists
        try:
            self.cloudwatch_logs.create_log_group(logGroupName=self.log_group_name)
        except self.cloudwatch_logs.exceptions.ResourceAlreadyExistsException:
            pass

        # Ensure the log stream exists
        if not self._log_stream_exists():
            self.cloudwatch_logs.create_log_stream(
                logGroupName=self.log_group_name,
                logStreamName=self.log_stream_name
            )

    def _log_stream_exists(self):
        response = self.cloudwatch_logs.describe_log_streams(
            logGroupName=self.log_group_name,
            logStreamNamePrefix=self.log_stream_name
        )
        for stream in response.get('logStreams', []):
            if stream['logStreamName'] == self.log_stream_name:
                self.sequence_token = stream.get('uploadSequenceToken')
                return True
        return False

    def emit(self, record):
        log_message = self.format(record)
        timestamp = int(datetime.now().timestamp() * 1000)

        # Ensure thread safety when sending logs to CloudWatch
        with self._lock:
            try:
                kwargs = {
                    'logGroupName': self.log_group_name,
                    'logStreamName': self.log_stream_name,
                    'logEvents': [{
                        'timestamp': timestamp,
                        'message': log_message
                    }]
                }
                if self.sequence_token:
                    kwargs['sequenceToken'] = self.sequence_token

                response = self.cloudwatch_logs.put_log_events(**kwargs)
                self.sequence_token = response['nextSequenceToken']
            except Exception as e:
                print(f"Error sending log to CloudWatch: {e}")


def get_logger():
    try:
        from flask import has_request_context, g
        if has_request_context():
            log_stream_name = getattr(g, 'log_stream_name', 'mconnect')
        else:
            log_stream_name = 'mconnect'
    except Exception:
        log_stream_name = 'mconnect'

    logger = logging.getLogger(f'api_wrapper_{log_stream_name}')
    if not logger.hasHandlers():
        logger.setLevel(logging.INFO)
        log_group_name = '/aws/api_wrapper'
        cloudwatch_handler = CloudWatchLogHandler(
            log_group_name=log_group_name,
            log_stream_name=log_stream_name
        )
        formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
        cloudwatch_handler.setFormatter(formatter)
        logger.addHandler(cloudwatch_handler)
    return logger


def print_log(message):
    print(message)
    logger = get_logger()
    logger.info(message)