import logging
import boto3
from datetime import datetime
import os
import threading
from flask import g, has_request_context


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)

        # put_log_events_batch already handles thread safety
        try:
            self.put_log_events_batch([{
                'timestamp': timestamp,
                'message': log_message
            }])
        except Exception as e:
            print(f"Error sending log to CloudWatch: {e}")

    def put_log_events_batch(self, log_events):
        """
        Sends a batch of log events to CloudWatch.
        This handles the sequence token updates and batch size limits.
        """
        if not log_events:
            return

        with self._lock:
            try:
                # Basic batching if too many events (CloudWatch limit is 10,000)
                # For simplicity, we just send them as provided, assuming they are within limits
                kwargs = {
                    'logGroupName': self.log_group_name,
                    'logStreamName': self.log_stream_name,
                    'logEvents': log_events
                }
                if self.sequence_token:
                    kwargs['sequenceToken'] = self.sequence_token

                response = self.cloudwatch_logs.put_log_events(**kwargs)
                self.sequence_token = response['nextSequenceToken']
            except self.cloudwatch_logs.exceptions.InvalidSequenceTokenException as e:
                # If sequence token is invalid, refresh it and retry once
                print(f"Invalid sequence token, refreshing... {e}")
                if self._log_stream_exists():
                    if self.sequence_token:
                        kwargs['sequenceToken'] = self.sequence_token
                    else:
                        kwargs.pop('sequenceToken', None)
                    response = self.cloudwatch_logs.put_log_events(**kwargs)
                    self.sequence_token = response['nextSequenceToken']
            except Exception as e:
                print(f"Error sending log batch to CloudWatch: {e}")


def get_logger():
    logger = logging.getLogger('json-bre-engine')
    if not logger.hasHandlers():
        logger.setLevel(logging.INFO)
        log_group_name = '/aws/json-bre-engine'
        log_stream_name = 'json-bre-engine'
        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):
    if os.environ.get("LOGGING") == 'cloudwatch':
        timestamp = int(datetime.now().timestamp() * 1000)
        log_entry = {
            'timestamp': timestamp,
            'message': message
        }

        # Check if we are in a Flask request context to collate logs
        if has_request_context():
            if not hasattr(g, 'cloudwatch_logs'):
                g.cloudwatch_logs = []
            g.cloudwatch_logs.append(log_entry)
        else:
            # Fallback for logs outside of request context
            logger = get_logger()
            logger.info(message)
    else:
        print(message)


def flush_cloudwatch_logs():
    """
    Pushes collated logs from g.cloudwatch_logs to CloudWatch in a single batch.
    To be called at the end of request processing.
    """
    if os.environ.get("LOGGING") == 'cloudwatch' and has_request_context():
        logs = getattr(g, 'cloudwatch_logs', [])
        if logs:
            logger = get_logger()
            # Find the CloudWatchLogHandler
            for handler in logger.handlers:
                if isinstance(handler, CloudWatchLogHandler):
                    formatted_messages = []
                    last_timestamp = logs[0]['timestamp']
                    for log_entry in logs:
                        now = datetime.fromtimestamp(log_entry['timestamp'] / 1000.0)
                        time_str = now.strftime('%Y-%m-%d %H:%M:%S,%f')[:-3]
                        formatted_message = f"{time_str} - INFO - {log_entry['message']}"
                        formatted_messages.append(formatted_message)
                        last_timestamp = max(last_timestamp, log_entry['timestamp'])
                    
                    # Join all messages with newlines and send as a single record
                    combined_message = "\n".join(formatted_messages)
                    
                    # Ensure it's within CloudWatch limits (256KB)
                    if len(combined_message.encode('utf-8')) > 250000:
                        combined_message = combined_message[:200000] + "\n... [TRUNCATED]"
                    
                    handler.put_log_events_batch([{
                        'timestamp': last_timestamp,
                        'message': combined_message
                    }])
                    break
            # Clear logs after flushing
            g.cloudwatch_logs = []