import redis
import json
import time
import uuid
import logging

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


class RWLQueueManager:
    """
    Manages a Redis-backed priority queue for RWL (Read Write Locate) jobs.
    """

    PRIORITIES = {
        'critical': 4,
        'high': 3,
        'normal': 2,
        'low': 1
    }

    STATUSES = {
        'queued': 'queued',
        'processing': 'processing',
        'completed': 'completed',
        'failed': 'failed',
        'retrying': 'retrying'
    }

    MAX_RETRIES = 3

    def __init__(self, redis_host='localhost', redis_port=6379, redis_db=0, queue_name='rwl_queue'):
        """
        Initializes the RWLQueueManager with Redis connection details.

        Args:
            redis_host (str): Redis host address.
            redis_port (int): Redis port number.
            redis_db (int): Redis database number.
            queue_name (str): Name of the Redis sorted set used as the queue.
        """
        self.redis_client = redis.Redis(host=redis_host, port=redis_port, db=redis_db)
        self.queue_name = queue_name
        self.logger = logging.getLogger(__name__)  # Use a class-specific logger
        self.logger.info(f"RWLQueueManager initialized with queue name: {self.queue_name}")


    def enqueue(self, prd_path, priority='normal'):
        """
        Enqueues a new RWL job with the specified priority.

        Args:
            prd_path (str): Path to the PRD file.
            priority (str): Priority of the job (critical, high, normal, low).  Defaults to 'normal'.

        Returns:
            str: The job ID.

        Raises:
            ValueError: If an invalid priority is provided.
        """
        if priority not in self.PRIORITIES:
            self.logger.error(f"Invalid priority: {priority}.  Must be one of {', '.join(self.PRIORITIES.keys())}")
            raise ValueError(f"Invalid priority: {priority}. Must be one of {', '.join(self.PRIORITIES.keys())}")

        job_id = str(uuid.uuid4())
        job_data = {
            'id': job_id,
            'prd_path': prd_path,
            'priority': priority,
            'status': self.STATUSES['queued'],
            'attempts': 0,
            'created_at': time.time()
        }

        score = self.PRIORITIES[priority]
        self.redis_client.zadd(self.queue_name, {json.dumps(job_data): score})
        self.logger.info(f"Enqueued job {job_id} with priority {priority} and PRD path {prd_path}")
        return job_id


    def dequeue(self):
        """
        Dequeues the highest priority RWL job from the queue.

        Returns:
            dict: The job data as a dictionary, or None if the queue is empty.
        """
        job_data = self.redis_client.zrange(self.queue_name, 0, 0, desc=True, withscores=False)

        if not job_data:
            return None

        job_data_str = job_data[0].decode('utf-8')
        job_data = json.loads(job_data_str)

        # Remove the job from the queue
        self.redis_client.zrem(self.queue_name, job_data_str)

        job_id = job_data['id']
        self.set_status(job_id, self.STATUSES['processing'])
        self.logger.info(f"Dequeued job {job_id} for processing")
        return job_data


    def set_status(self, job_id, status):
        """
        Updates the status of a job.

        Args:
            job_id (str): The ID of the job.
            status (str): The new status of the job (queued, processing, completed, failed, retrying).

        Raises:
            ValueError: If an invalid status is provided.
            KeyError: If the job is not found (status not updated).
        """
        if status not in self.STATUSES.values():
            self.logger.error(f"Invalid status: {status}. Must be one of {', '.join(self.STATUSES.values())}")
            raise ValueError(f"Invalid status: {status}. Must be one of {', '.join(self.STATUSES.values())}")

        job_data = self._get_job_data(job_id)
        if job_data is None:
             self.logger.error(f"Job not found: {job_id}. Cannot update status.")
             raise KeyError(f"Job not found: {job_id}")

        job_data['status'] = status

        #Re-enqueue the job with the updated status.  We need to remove, then add because sorted set.
        self._remove_job_by_id(job_id)
        score = self.PRIORITIES[job_data['priority']]
        self.redis_client.zadd(self.queue_name, {json.dumps(job_data): score})

        self.logger.info(f"Updated status of job {job_id} to {status}")

    def _get_job_data(self, job_id):
      """
      Internal method to retrieve job data by ID. This does not remove the job.

      Args:
          job_id (str): The ID of the job to retrieve.

      Returns:
          dict: The job data as a dictionary, or None if the job is not found.
      """
      for job_data_str in self.redis_client.zrange(self.queue_name, 0, -1):
        job_data = json.loads(job_data_str.decode('utf-8'))
        if job_data['id'] == job_id:
          return job_data
      return None

    def _remove_job_by_id(self, job_id):
      """
      Internal method to remove a job from the queue by its ID.

      Args:
          job_id (str): The ID of the job to remove.

      Returns:
          bool: True if the job was found and removed, False otherwise.
      """
      for job_data_str in self.redis_client.zrange(self.queue_name, 0, -1):
        job_data = json.loads(job_data_str.decode('utf-8'))
        if job_data['id'] == job_id:
          self.redis_client.zrem(self.queue_name, job_data_str)
          return True
      return False


    def job_completed(self, job_id):
        """
        Marks a job as completed.

        Args:
            job_id (str): The ID of the job.
        """
        try:
          self.set_status(job_id, self.STATUSES['completed'])
          self.logger.info(f"Job {job_id} completed successfully")
        except KeyError:
          self.logger.warning(f"Tried to complete job {job_id}, but it was not found.")

    def job_failed(self, job_id):
        """
        Handles job failure, retrying if possible.

        Args:
            job_id (str): The ID of the job.
        """
        try:
          job_data = self._get_job_data(job_id)

          if job_data is None:
              self.logger.warning(f"Tried to fail job {job_id}, but it was not found.")
              return

          job_data['attempts'] += 1
          if job_data['attempts'] <= self.MAX_RETRIES:
              self.set_status(job_id, self.STATUSES['retrying'])
              self.logger.info(f"Job {job_id} failed, retrying (attempt {job_data['attempts']}/{self.MAX_RETRIES})")
              # Re-enqueue the job with the same priority
              self._remove_job_by_id(job_id)  # Remove the old entry
              score = self.PRIORITIES[job_data['priority']]
              self.redis_client.zadd(self.queue_name, {json.dumps(job_data): score})
          else:
              self.set_status(job_id, self.STATUSES['failed'])
              self.logger.error(f"Job {job_id} failed after {self.MAX_RETRIES} attempts")

        except KeyError:
          self.logger.warning(f"Tried to fail job {job_id}, but it was not found.")

    def get_job_status(self, job_id):
        """
        Retrieves the status of a job.

        Args:
            job_id (str): The ID of the job.

        Returns:
            str: The status of the job, or None if the job is not found.
        """
        job_data = self._get_job_data(job_id)
        if job_data:
          return job_data['status']
        else:
          return None

    def get_queue_length(self):
        """
        Returns the number of jobs in the queue.
        """
        return self.redis_client.zcard(self.queue_name)

    def clear_queue(self):
        """
        Clears all jobs from the queue.  USE WITH CAUTION!
        """
        self.redis_client.delete(self.queue_name)
        self.logger.warning("Queue cleared! All jobs have been removed.")

    def list_jobs(self, start=0, end=-1):
      """
      Lists all jobs in the queue.

      Args:
          start (int): The starting index (inclusive) of the range to retrieve. Default is 0.
          end (int): The ending index (inclusive) of the range to retrieve. Default is -1 (all jobs).

      Returns:
          list: A list of job data dictionaries.
      """
      jobs = []
      for job_data_str in self.redis_client.zrange(self.queue_name, start, end):
          job_data = json.loads(job_data_str.decode('utf-8'))
          jobs.append(job_data)
      return jobs