import json
import sys
import time
import logging
from logging import ERROR, INFO, DEBUG, WARN
from positron_job_runner.runner_env import runner_env
from positron_job_runner.cloudwatch_client import create_client
from positron_job_runner.logging_config import get_logger

console_logger = get_logger()

class CloudLogger:
  """A class to log messages to stdout, stderr, and CloudWatch."""
  stdout = []
  stderr = []
  # cloudwatch_client = None

  def __init__(self):
    # Connect to cloudwatch if enabled
    if runner_env.AWS_JOB_LOG_GROUP_NAME:
      self.__cloudwatch_client = create_client()
    else:
      console_logger.info("CloudWatch logging is disabled.")

  def error(self, message: str):
    self.log(message, ERROR)

  def warn(self, message: str):
    self.log(message, WARN)

  def info(self, message: str):
    self.log(message, INFO)
  
  def debug(self, message: str):
    self.log(message, DEBUG)
  
  def exception(self, *args, **kwargs):
    console_logger.exception(*args, **kwargs)
    self.log(str(args[0]), ERROR)

  def log(self, message: str, level = INFO):
    if isinstance(message, bytes):
        message = message.decode('utf-8')
    
    # Store in mem for later retrieval
    if level == ERROR:
        self.stderr.append(message)
    else:
        self.stdout.append(message)

    # Log to console at the provided level.
    console_logger.log(level = level, msg = message)

    # Log to CloudWatch if enabled.
    if runner_env.AWS_JOB_LOG_GROUP_NAME:
      self._log_to_cloudwatch(message, level)

    # This ensures logs in the right order etc, but may slow down the process
    sys.stdout.flush()
    sys.stderr.flush()

  def end_of_logs_signal(self):
    if runner_env.AWS_JOB_LOG_GROUP_NAME:
      self._log_to_cloudwatch("POSITRON_SIGNAL_EOS")

  def _log_to_cloudwatch(self, message: str, level = INFO):
    timestamp = int(round(time.time() * 1000))
    message = {'timestamp': timestamp, 'log_level': logging.getLevelName(level), 'message': message, 'app_name': 'remote_machine'}
    log_events = [{
        'timestamp': timestamp,
        'message': json.dumps(message),
    }]
    log_stream_name = f'positron-job-{runner_env.JOB_ID}'
    log_group_name = runner_env.AWS_JOB_LOG_GROUP_NAME

    try:
        self.__cloudwatch_client.put_log_events(
            logGroupName=log_group_name,
            logStreamName=log_stream_name,
            logEvents=log_events
        )
    except Exception as e:
        print(f'Failed to submit logs: {e}')
  

logger = CloudLogger()