import json
import base64
from typing import Optional, Dict

# Import the Redis client that all workers will share
from .redis_client import redis_client

class ResponseManager:
    """
    Manages response state in a central Redis store, keyed by student_id.
    This solution is safe for multiple workers.
    """

    def __init__(self):
        """Initializes by connecting to the shared Redis client."""
        if redis_client is None:
            raise ConnectionError("ResponseManager requires a valid Redis connection. Check your REDIS_HOST/PORT environment variables.")
        self.redis = redis_client
        self.ttl_seconds = 600  # Responses will expire after 10 minutes

    def _get_key(self, student_id: str) -> str:
        """Creates a consistent key for Redis to avoid conflicts."""
        return f"student_response:{student_id}"

    def store_response(self, student_id: str, text: str, audio_filepath: Optional[str] = None, audio_bytes: Optional[bytes] = None) -> None:
        """Stores a response for a specific student_id in Redis."""
        key = self._get_key(student_id)
        
        # Encode binary audio data into a string (Base64) to store it in JSON
        encoded_audio = base64.b64encode(audio_bytes).decode('utf-8') if audio_bytes else None
        
        payload = {
            "text": text,
            "audio_filepath": audio_filepath,
            "audio_bytes_b64": encoded_audio
        }
        
        # Convert the dictionary to a JSON string and store it in Redis with an expiration time
        self.redis.setex(key, self.ttl_seconds, json.dumps(payload))

    def get_response(self, student_id: str) -> Dict:
        """
        Atomically gets the response for a student and removes it from Redis
        to ensure it's claimed only once.
        """
        key = self._get_key(student_id)
        
        # Use a pipeline to get and delete the key in a single, atomic operation
        pipe = self.redis.pipeline()
        pipe.get(key)
        pipe.delete(key)
        results = pipe.execute()
        json_value = results[0]

        if not json_value:
            # If nothing was found, return the same empty structure as the old class
            return {"text": None, "audio_filename": None, "audio_bytes": None}

        # If data was found, decode it
        payload = json.loads(json_value)
        
        # Decode the Base64 string back into binary audio data
        if payload.get("audio_bytes_b64"):
            payload["audio_bytes"] = base64.b64decode(payload["audio_bytes_b64"])
        else:
            payload["audio_bytes"] = None
        
        # Remove the temporary key before returning
        del payload["audio_bytes_b64"]
        
        return payload

    def clear_response(self, student_id: str) -> None:
        """Clears a response for a specific student from Redis."""
        key = self._get_key(student_id)
        self.redis.delete(key)

    def is_response_fresh(self, student_id: str, max_age_seconds: int = 300) -> bool:
        """Checks if a response exists in Redis for the given student."""
        key = self._get_key(student_id)
        # redis.exists() is the direct equivalent of checking if the key is present
        return self.redis.exists(key) > 0