import json
import asyncio
from typing import Any, Dict, Optional
from jazzmine.logging.sinks.base_sink import BaseSink
from jazzmine.logging.exceptions import LoggerDependencyError


class KafkaLogger(BaseSink):
    """Apache Kafka sink with both sync and async support"""
    
    def initialize(self) -> None:
        self.bootstrap_servers = self.sink_config.get('bootstrap_servers', ['localhost:9092'])
        self.topic = self.sink_config.get('topic', 'logs')
        self.use_async = self.sink_config.get('async', True)
        
        if self.use_async:
            try:
                from aiokafka import AIOKafkaProducer
                self.producer_class = AIOKafkaProducer
            except ImportError:
                raise LoggerDependencyError(
                    "aiokafka is required for async Kafka sink. Install: pip install aiokafka"
                )
        else:
            try:
                from kafka import KafkaProducer
                self.producer_class = KafkaProducer
            except ImportError:
                raise LoggerDependencyError(
                    "kafka-python is required for Kafka sink. Install: pip install kafka-python"
                )
        
        self.producer: Optional[Any] = None
        self._initialized = True
        
        if not self.use_async:
            self.producer = self.producer_class(
                bootstrap_servers=self.bootstrap_servers,
                value_serializer=lambda v: json.dumps(v).encode('utf-8')
            )
    
    def shutdown(self) -> None:
        if self.producer:
            if self.use_async:
                asyncio.create_task(self.producer.stop())
            else:
                self.producer.close()
    
    def _format_message(self, level: str, msg: str, **kwargs) -> Dict[str, Any]:
        import time
        return {
            'timestamp': time.time(),
            'level': level,
            'logger': self.logger_name,
            'message': msg,
            **kwargs
        }
    
    def log(self, level: str, msg: str, **kwargs) -> None:
        if not self._initialized or self.use_async:
            return
        
        payload = self._format_message(level, msg, **kwargs)
        
        try:
            self.producer.send(self.topic, value=payload)
        except Exception:
            pass
    
    async def alog(self, level: str, msg: str, **kwargs) -> None:
        if not self._initialized:
            return
        
        if not self.producer:
            self.producer = self.producer_class(bootstrap_servers=self.bootstrap_servers)
            await self.producer.start()
        
        payload = self._format_message(level, msg, **kwargs)
        
        try:
            await self.producer.send_and_wait(
                self.topic,
                value=json.dumps(payload).encode('utf-8')
            )
        except Exception:
            pass
