import json
import socket
import struct
import asyncio
from typing import Any, Dict
import aiohttp
import requests
from jazzmine.logging.sinks.base_sink import BaseSink


class GraylogLogger(BaseSink):
    """Graylog GELF (Graylog Extended Log Format) sink"""
    
    def initialize(self) -> None:
        self.host = self.sink_config.get('host', 'localhost')
        self.port = self.sink_config.get('port', 12201)
        self.protocol = self.sink_config.get('protocol', 'http')  # http, udp, tcp
        self.session = None
        
        if self.protocol == 'http':
            self.url = f"http://{self.host}:{self.port}/gelf"
        elif self.protocol in ('udp', 'tcp'):
            self.sock = socket.socket(
                socket.AF_INET,
                socket.SOCK_DGRAM if self.protocol == 'udp' else socket.SOCK_STREAM
            )
            if self.protocol == 'tcp':
                self.sock.connect((self.host, self.port))
        
        self._initialized = True
    
    def shutdown(self) -> None:
        if self.protocol in ('udp', 'tcp') and hasattr(self, 'sock'):
            self.sock.close()
        if self.session:
            asyncio.create_task(self.session.close())
    
    def _format_gelf(self, level: str, msg: str, **kwargs) -> Dict[str, Any]:
        import time
        
        level_map = {
            'debug': 7,
            'info': 6,
            'warning': 4,
            'error': 3,
            'critical': 2
        }
        
        gelf_msg = {
            'version': '1.1',
            'host': socket.gethostname(),
            'short_message': msg,
            'timestamp': time.time(),
            'level': level_map.get(level, 6),
            '_logger': self.logger_name,
        }
        
        for key, value in kwargs.items():
            gelf_msg[f'_{key}'] = value
        
        return gelf_msg
    
    def log(self, level: str, msg: str, **kwargs) -> None:
        if not self._initialized:
            return
        
        payload = self._format_gelf(level, msg, **kwargs)
        
        try:
            if self.protocol == 'http':
                requests.post(self.url, json=payload, timeout=5)
            elif self.protocol == 'udp':
                data = json.dumps(payload).encode('utf-8')
                self.sock.sendto(data, (self.host, self.port))
            elif self.protocol == 'tcp':
                data = json.dumps(payload).encode('utf-8')
                # GELF TCP requires null byte terminator
                self.sock.sendall(data + b'\0')
        except Exception:
            pass
    
    async def alog(self, level: str, msg: str, **kwargs) -> None:
        if not self._initialized:
            return
        
        payload = self._format_gelf(level, msg, **kwargs)
        
        if self.protocol == 'http':
            if not self.session:
                self.session = aiohttp.ClientSession()
            
            try:
                await self.session.post(self.url, json=payload, timeout=aiohttp.ClientTimeout(total=5))
            except Exception:
                pass
        else:
            # For UDP/TCP, use sync method
            self.log(level, msg, **kwargs)
