"""
Main OCPP UDP client server implementation
"""
import asyncio
import logging
import signal
import sys
from typing import Optional, Tuple
from .config import get_config, validate_required_config
from .utils import parse_ocpp_message
from shared.dispatcher import get_unified_dispatcher

# Get configuration and configure logging
config = get_config()
logger = logging.getLogger(__name__)


class OCPPProtocol(asyncio.DatagramProtocol):
    """UDP Protocol handler for OCPP messages"""
    
    def __init__(self, dispatcher):
        self.transport: Optional[asyncio.DatagramTransport] = None
        self.dispatcher = dispatcher
        
    def connection_made(self, transport: asyncio.DatagramTransport):
        """Called when the transport is ready"""
        self.transport = transport
        logger.info("UDP transport ready")
    
    def datagram_received(self, data: bytes, addr: Tuple[str, int]):
        """Called when a UDP packet is received"""
        try:
            # Log source IP and raw message
            raw_message = data.decode('utf-8')
            logger.info(f"Received from {addr[0]}:{addr[1]} - Raw: {raw_message}")
            
            # Create asyncio task to handle message processing
            asyncio.create_task(self._handle_message(raw_message, addr))
            
        except UnicodeDecodeError:
            logger.error(f"Failed to decode message from {addr[0]}:{addr[1]} - Invalid UTF-8")
        except Exception as e:
            logger.error(f"Error processing message from {addr[0]}:{addr[1]}: {e}")
    
    async def _handle_message(self, raw_message: str, addr: Tuple[str, int]):
        """Handle message processing asynchronously"""
        try:
            # Use dispatcher's handle_udp_packet method for comprehensive message handling
            response_bytes = await self.dispatcher.handle_udp_packet(raw_message.encode('utf-8'), addr)
            
            if response_bytes and self.transport:
                # Send response back to client
                self.transport.sendto(response_bytes, addr)
                response_str = response_bytes.decode('utf-8')
                logger.debug(f"Sent response to {addr[0]}:{addr[1]}: {response_str}")
                    
        except Exception as e:
            logger.error(f"Error handling message from {addr[0]}:{addr[1]}: {e}")
    
    def error_received(self, exc: Exception):
        """Called when a send or receive operation raises an OSError"""
        logger.error(f"UDP protocol error: {exc}")
    
    def connection_lost(self, exc: Optional[Exception]):
        """Called when the transport is closed"""
        if exc:
            logger.error(f"UDP connection lost with error: {exc}")
        else:
            logger.info("UDP connection closed gracefully")


class OCPPServer:
    """OCPP UDP Client Server"""
    
    def __init__(self, host: str = None, port: int = None):
        self.host = host or config.udp_host
        self.port = port or config.udp_port
        self.transport: Optional[asyncio.DatagramTransport] = None
        self.protocol: Optional[OCPPProtocol] = None
        self.dispatcher = get_unified_dispatcher()
        self._shutdown_event = asyncio.Event()
        
        # Setup signal handlers for graceful shutdown
        self._setup_signal_handlers()
        
    def _setup_signal_handlers(self):
        """Setup signal handlers for graceful shutdown"""
        if sys.platform != 'win32':
            # Unix/Linux signal handling
            def signal_handler(signum, frame):
                logger.info(f"Received signal {signum}, initiating graceful shutdown...")
                asyncio.create_task(self.shutdown())
            
            signal.signal(signal.SIGINT, signal_handler)
            signal.signal(signal.SIGTERM, signal_handler)
        
    async def start_server(self):
        """Start the UDP server"""
        try:
            logger.info(f"Starting UDP server on {self.host}:{self.port}")
            
            loop = asyncio.get_running_loop()
            
            # Create UDP endpoint with our protocol
            self.transport, self.protocol = await loop.create_datagram_endpoint(
                lambda: OCPPProtocol(self.dispatcher),
                local_addr=(self.host, self.port)
            )
            
            logger.info(f"UDP server running on {self.host}:{self.port}")
            
            # Wait for shutdown signal
            await self._shutdown_event.wait()
            
        except OSError as e:
            if e.errno == 48:  # Address already in use
                logger.error(f"Port {self.port} is already in use. Please choose a different port or stop the conflicting service.")
            else:
                logger.error(f"Failed to bind to {self.host}:{self.port}: {e}")
            raise
        except Exception as e:
            logger.error(f"Error starting UDP server: {e}")
            raise
        finally:
            await self.cleanup()
    
    async def shutdown(self):
        """Initiate graceful shutdown"""
        logger.info("Shutting down UDP server...")
        self._shutdown_event.set()
    
    async def cleanup(self):
        """Clean up resources"""
        if self.transport:
            logger.info("Closing UDP transport...")
            self.transport.close()
            
            # Wait a bit for the transport to close cleanly
            try:
                await asyncio.wait_for(asyncio.sleep(0.1), timeout=1.0)
            except asyncio.TimeoutError:
                pass
                
        logger.info("UDP server stopped")


async def main():
    """Main entry point"""
    try:
        # Validate configuration
        validate_required_config()
        
        # Print configuration at startup
        config.print_config()
        
        logger.info("Voltie OCPP UDP Client starting...")
        
        server = OCPPServer(config.udp_host, config.udp_port)
        await server.start_server()
        
    except KeyboardInterrupt:
        logger.info("Received Ctrl+C, shutting down...")
    except Exception as e:
        logger.error(f"Failed to start server: {e}")
        raise


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\nGraceful shutdown complete.")
        sys.exit(0)
    except Exception as e:
        logger.error(f"Server failed: {e}")
        sys.exit(1)