"""Redis worker support for Superviseur bot."""

import json
import uuid
import asyncio
import logging
from typing import Dict, Optional, Any

logger = logging.getLogger('Superviseur')


class RedisWorker:
    """Manages Redis-based LLM job queue."""
    
    def __init__(self, bot):
        self.bot = bot
        self.redis_client = None
        self._pending_jobs: Dict[str, asyncio.Future] = {}
        self._listener_task: Optional[asyncio.Task] = None
        self._response_queue = 'llm_queue:responses'
        self._job_queue = 'llm_queue:jobs'
    
    async def initialize(self, redis_url: str) -> bool:
        """Initialize Redis connection."""
        if not self.bot.redis_available:
            return False
        
        try:
            self.redis_client = self.bot.aioredis.from_url(redis_url)
            self._listener_task = asyncio.create_task(self._listen())
            logger.info("Redis LLM worker enabled; listening for responses")
            return True
        except Exception:
            logger.exception("Failed to initialize redis client")
            return False
    
    async def close(self) -> None:
        """Close Redis connection."""
        if self._listener_task:
            try:
                self._listener_task.cancel()
            except Exception:
                pass
        
        if self.redis_client:
            try:
                await self.redis_client.close()
            except Exception:
                pass
    
    async def enqueue_job(self, payload: Dict[str, Any], timeout: int = 60) -> Optional[Dict[str, Any]]:
        """Enqueue an LLM job and wait for result."""
        if not self.redis_client:
            return None
        
        job_id = str(uuid.uuid4())
        future = self.bot.loop.create_future()
        self._pending_jobs[job_id] = future
        payload['job_id'] = job_id
        
        await self.redis_client.lpush(self._job_queue, json.dumps(payload))
        logger.debug(f"LLM job {job_id} enqueued.")
        
        try:
            return await asyncio.wait_for(future, timeout=timeout)
        finally:
            self._pending_jobs.pop(job_id, None)
    
    async def _listen(self) -> None:
        """Listen for Redis responses."""
        logger.info(f"Starting Redis response listener on {self._response_queue}")
        
        while True:
            try:
                result = await self.redis_client.brpop(self._response_queue, timeout=1)
                if result is None:
                    await asyncio.sleep(0.1)
                    continue
                
                data_bytes = result[1]
                response = json.loads(data_bytes.decode('utf-8'))
                job_id = response.get('job_id')
                
                if job_id in self._pending_jobs:
                    future = self._pending_jobs[job_id]
                    if not future.done():
                        future.set_result(response)
                        logger.debug(f"LLM job {job_id} completed via Redis.")
                else:
                    logger.warning(f"Response for unknown job {job_id}")
            
            except asyncio.CancelledError:
                logger.info("Redis listener cancelled")
                break
            except Exception as e:
                logger.exception(f"Redis listener error: {e}")
                await asyncio.sleep(5)