from aiokafka import AIOKafkaConsumer
import json
import asyncio
from typing import Dict, Any
from ...core.config import settings
from ...application.use_cases.process_search_results import ProcessSearchResultsUseCase
from ...infrastructure.repositories.mongodb_search_repository import MongoDBSearchRepository
from ...core.database import get_database
import logging

logger = logging.getLogger(__name__)

class SearchResultsConsumer:
    def __init__(self):
        self.consumer = None
        self.is_running = False
        self.process_results_use_case = None
    
    async def start(self):
        """Start the Kafka consumer"""
        try:
            # Initialize use case
            database = await get_database()
            search_repository = MongoDBSearchRepository(database)
            self.process_results_use_case = ProcessSearchResultsUseCase(search_repository)
            
            # Create consumer
            self.consumer = AIOKafkaConsumer(
                'search-results',  # Topic for search results
                bootstrap_servers=settings.KAFKA_BOOTSTRAP_SERVERS,
                group_id=f"{settings.KAFKA_CONSUMER_GROUP}-results",
                value_deserializer=lambda m: json.loads(m.decode('utf-8'))
            )
            
            await self.consumer.start()
            self.is_running = True
            
            logger.info("Search results consumer started")
            
            # Start consuming messages
            await self._consume_messages()
            
        except Exception as e:
            logger.error(f"Error starting search results consumer: {e}")
            raise
    
    async def stop(self):
        """Stop the Kafka consumer"""
        self.is_running = False
        if self.consumer:
            await self.consumer.stop()
            logger.info("Search results consumer stopped")
    
    async def _consume_messages(self):
        """Consume messages from Kafka"""
        try:
            async for message in self.consumer:
                if not self.is_running:
                    break
                
                await self._process_message(message.value)
                
        except Exception as e:
            logger.error(f"Error consuming messages: {e}")
    
    async def _process_message(self, message: Dict[str, Any]):
        """Process a single message"""
        try:
            message_type = message.get("type")
            task_id = message.get("task_id")
            
            if message_type == "task_completed":
                results = message.get("results", {})
                await self.process_results_use_case.process_task_completion(task_id, results)
                
            elif message_type == "task_failed":
                error = message.get("error", "Unknown error")
                await self.process_results_use_case.process_task_failure(task_id, error)
            
            else:
                logger.warning(f"Unknown message type: {message_type}")
                
        except Exception as e:
            logger.error(f"Error processing message: {e}")

search_results_consumer = SearchResultsConsumer()