"""
Modbus RTU client for chocolate tempering machine hardware communication.
Provides async interface to industrial hardware via RS-485/serial.
"""

import asyncio
import logging
from typing import Dict, List, Optional, Union, Any
from datetime import datetime, timedelta
from dataclasses import dataclass
from enum import Enum

from pymodbus.client import AsyncModbusSerialClient
from pymodbus.exceptions import ModbusException, ConnectionException
from pymodbus.pdu import ExceptionResponse

from ...shared.config import settings


logger = logging.getLogger(__name__)


class ModbusFunction(Enum):
    """Supported Modbus function codes."""
    READ_COILS = 0x01
    READ_DISCRETE_INPUTS = 0x02
    READ_HOLDING_REGISTERS = 0x03
    READ_INPUT_REGISTERS = 0x04
    WRITE_SINGLE_COIL = 0x05
    WRITE_SINGLE_REGISTER = 0x06
    WRITE_MULTIPLE_COILS = 0x0F
    WRITE_MULTIPLE_REGISTERS = 0x10


@dataclass
class ModbusRequest:
    """Modbus operation request."""
    function: ModbusFunction
    address: int
    count: Optional[int] = None
    value: Optional[Union[int, List[int], bool, List[bool]]] = None
    slave_id: int = 1
    retry_count: int = 0
    max_retries: int = 3
    timeout: float = 2.0
    priority: int = 1  # Lower number = higher priority


@dataclass
class ModbusResponse:
    """Modbus operation response."""
    success: bool
    value: Optional[Union[int, List[int], bool, List[bool]]] = None
    error: Optional[str] = None
    timestamp: Optional[datetime] = None
    duration_ms: Optional[float] = None


class ModbusClientError(Exception):
    """Custom exception for Modbus client errors."""
    pass


class AsyncModbusClient:
    """
    Asynchronous Modbus RTU client with connection management and error recovery.
    """
    
    def __init__(self):
        self.client: Optional[AsyncModbusSerialClient] = None
        self.is_connected = False
        self.connection_lock = asyncio.Lock()
        self.request_queue: asyncio.Queue = asyncio.Queue()
        self.response_cache: Dict[str, ModbusResponse] = {}
        self.last_successful_read = datetime.now()
        self.consecutive_errors = 0
        self.max_consecutive_errors = 10
        
        # Configuration
        self.port = settings.serial.port
        self.baudrate = settings.serial.baudrate
        self.timeout = settings.serial.timeout
        self.slave_address = settings.modbus.slave_address
        self.max_read_registers = settings.modbus.max_read_registers
        self.max_write_registers = settings.modbus.max_write_registers
        
        # Statistics
        self.stats = {
            "total_requests": 0,
            "successful_requests": 0,
            "failed_requests": 0,
            "connection_errors": 0,
            "timeout_errors": 0,
            "last_error": None,
            "uptime_start": datetime.now()
        }
    
    async def connect(self) -> bool:
        """Establish Modbus RTU connection."""
        async with self.connection_lock:
            try:
                if self.client:
                    await self.disconnect()
                
                logger.info(f"Connecting to Modbus RTU on {self.port} at {self.baudrate} baud")
                
                self.client = AsyncModbusSerialClient(
                    port=self.port,
                    baudrate=self.baudrate,
                    bytesize=settings.serial.bytesize,
                    parity=settings.serial.parity,
                    stopbits=settings.serial.stopbits,
                    timeout=self.timeout,
                )
                
                connected = await self.client.connect()
                
                if connected:
                    self.is_connected = True
                    self.consecutive_errors = 0
                    logger.info("Modbus RTU connection established successfully")
                    return True
                else:
                    self.is_connected = False
                    self.stats["connection_errors"] += 1
                    logger.error("Failed to establish Modbus RTU connection")
                    return False
                    
            except Exception as e:
                self.is_connected = False
                self.stats["connection_errors"] += 1
                logger.error(f"Modbus connection error: {e}")
                return False
    
    async def disconnect(self) -> None:
        """Close Modbus connection."""
        async with self.connection_lock:
            if self.client:
                try:
                    self.client.close()
                    logger.info("Modbus RTU connection closed")
                except Exception as e:
                    logger.error(f"Error closing Modbus connection: {e}")
                finally:
                    self.client = None
                    self.is_connected = False
    
    async def ensure_connected(self) -> bool:
        """Ensure Modbus connection is active, reconnect if necessary."""
        if not self.is_connected or not self.client:
            return await self.connect()
        
        # Test connection with a simple read
        try:
            # Try to read a single register to test connection
            result = await self.client.read_holding_registers(0, 1, slave=self.slave_address)
            if isinstance(result, ExceptionResponse):
                # Connection might be broken, try to reconnect
                logger.warning("Connection test failed, attempting reconnection")
                return await self.connect()
            return True
        except Exception as e:
            logger.warning(f"Connection test error: {e}, attempting reconnection")
            return await self.connect()
    
    async def read_coils(self, address: int, count: int, slave_id: Optional[int] = None) -> ModbusResponse:
        """Read coil states (digital outputs)."""
        return await self._execute_request(ModbusRequest(
            function=ModbusFunction.READ_COILS,
            address=address,
            count=count,
            slave_id=slave_id or self.slave_address
        ))
    
    async def read_discrete_inputs(self, address: int, count: int, slave_id: Optional[int] = None) -> ModbusResponse:
        """Read discrete input states (digital inputs)."""
        return await self._execute_request(ModbusRequest(
            function=ModbusFunction.READ_DISCRETE_INPUTS,
            address=address,
            count=count,
            slave_id=slave_id or self.slave_address
        ))
    
    async def read_holding_registers(self, address: int, count: int, slave_id: Optional[int] = None) -> ModbusResponse:
        """Read holding register values (analog outputs/settings)."""
        return await self._execute_request(ModbusRequest(
            function=ModbusFunction.READ_HOLDING_REGISTERS,
            address=address,
            count=count,
            slave_id=slave_id or self.slave_address
        ))
    
    async def read_input_registers(self, address: int, count: int, slave_id: Optional[int] = None) -> ModbusResponse:
        """Read input register values (analog inputs/sensors)."""
        return await self._execute_request(ModbusRequest(
            function=ModbusFunction.READ_INPUT_REGISTERS,
            address=address,
            count=count,
            slave_id=slave_id or self.slave_address
        ))
    
    async def write_coil(self, address: int, value: bool, slave_id: Optional[int] = None) -> ModbusResponse:
        """Write single coil state (digital output)."""
        return await self._execute_request(ModbusRequest(
            function=ModbusFunction.WRITE_SINGLE_COIL,
            address=address,
            value=value,
            slave_id=slave_id or self.slave_address
        ))
    
    async def write_register(self, address: int, value: int, slave_id: Optional[int] = None) -> ModbusResponse:
        """Write single register value (analog output/setting)."""
        return await self._execute_request(ModbusRequest(
            function=ModbusFunction.WRITE_SINGLE_REGISTER,
            address=address,
            value=value,
            slave_id=slave_id or self.slave_address
        ))
    
    async def write_multiple_coils(self, address: int, values: List[bool], slave_id: Optional[int] = None) -> ModbusResponse:
        """Write multiple coil states (digital outputs)."""
        return await self._execute_request(ModbusRequest(
            function=ModbusFunction.WRITE_MULTIPLE_COILS,
            address=address,
            value=values,
            slave_id=slave_id or self.slave_address
        ))
    
    async def write_multiple_registers(self, address: int, values: List[int], slave_id: Optional[int] = None) -> ModbusResponse:
        """Write multiple register values (analog outputs/settings)."""
        return await self._execute_request(ModbusRequest(
            function=ModbusFunction.WRITE_MULTIPLE_REGISTERS,
            address=address,
            value=values,
            slave_id=slave_id or self.slave_address
        ))
    
    async def _execute_request(self, request: ModbusRequest) -> ModbusResponse:
        """Execute a Modbus request with error handling and retries."""
        start_time = datetime.now()
        
        for attempt in range(request.max_retries + 1):
            try:
                # Ensure connection is active
                if not await self.ensure_connected():
                    raise ModbusClientError("Failed to establish connection")
                
                # Execute the request based on function code
                result = await self._perform_operation(request)
                
                # Process successful response
                if not isinstance(result, ExceptionResponse):
                    self.consecutive_errors = 0
                    self.last_successful_read = datetime.now()
                    self.stats["successful_requests"] += 1
                    
                    response = ModbusResponse(
                        success=True,
                        value=self._extract_value(result, request.function),
                        timestamp=datetime.now(),
                        duration_ms=(datetime.now() - start_time).total_seconds() * 1000
                    )
                    
                    logger.debug(f"Modbus request successful: {request.function.name} addr={request.address}")
                    return response
                else:
                    raise ModbusException(f"Modbus exception response: {result}")
                    
            except Exception as e:
                self.consecutive_errors += 1
                self.stats["failed_requests"] += 1
                
                error_msg = f"Modbus request failed (attempt {attempt + 1}/{request.max_retries + 1}): {e}"
                logger.error(error_msg)
                
                # Check if we should trigger connection recovery
                if self.consecutive_errors >= self.max_consecutive_errors:
                    logger.critical("Too many consecutive errors, forcing reconnection")
                    await self.disconnect()
                
                # If this was the last attempt, return error response
                if attempt >= request.max_retries:
                    self.stats["last_error"] = str(e)
                    return ModbusResponse(
                        success=False,
                        error=error_msg,
                        timestamp=datetime.now(),
                        duration_ms=(datetime.now() - start_time).total_seconds() * 1000
                    )
                
                # Wait before retry with exponential backoff
                retry_delay = settings.serial.retry_delay * (2 ** attempt)
                await asyncio.sleep(retry_delay)
        
        # Should never reach here
        return ModbusResponse(success=False, error="Unexpected error in request execution")
    
    async def _perform_operation(self, request: ModbusRequest):
        """Perform the actual Modbus operation."""
        if not self.client:
            raise ModbusClientError("Modbus client not initialized")
        
        slave = request.slave_id
        
        if request.function == ModbusFunction.READ_COILS:
            return await self.client.read_coils(request.address, request.count, slave=slave)
            
        elif request.function == ModbusFunction.READ_DISCRETE_INPUTS:
            return await self.client.read_discrete_inputs(request.address, request.count, slave=slave)
            
        elif request.function == ModbusFunction.READ_HOLDING_REGISTERS:
            return await self.client.read_holding_registers(request.address, request.count, slave=slave)
            
        elif request.function == ModbusFunction.READ_INPUT_REGISTERS:
            return await self.client.read_input_registers(request.address, request.count, slave=slave)
            
        elif request.function == ModbusFunction.WRITE_SINGLE_COIL:
            return await self.client.write_coil(request.address, request.value, slave=slave)
            
        elif request.function == ModbusFunction.WRITE_SINGLE_REGISTER:
            return await self.client.write_register(request.address, request.value, slave=slave)
            
        elif request.function == ModbusFunction.WRITE_MULTIPLE_COILS:
            return await self.client.write_coils(request.address, request.value, slave=slave)
            
        elif request.function == ModbusFunction.WRITE_MULTIPLE_REGISTERS:
            return await self.client.write_registers(request.address, request.value, slave=slave)
            
        else:
            raise ModbusClientError(f"Unsupported Modbus function: {request.function}")
    
    def _extract_value(self, result, function: ModbusFunction) -> Union[int, List[int], bool, List[bool]]:
        """Extract value from Modbus response."""
        if function in [ModbusFunction.READ_COILS, ModbusFunction.READ_DISCRETE_INPUTS]:
            # Boolean values
            if hasattr(result, 'bits'):
                return result.bits
            return []
            
        elif function in [ModbusFunction.READ_HOLDING_REGISTERS, ModbusFunction.READ_INPUT_REGISTERS]:
            # Register values
            if hasattr(result, 'registers'):
                return result.registers
            return []
            
        elif function in [ModbusFunction.WRITE_SINGLE_COIL, ModbusFunction.WRITE_SINGLE_REGISTER,
                         ModbusFunction.WRITE_MULTIPLE_COILS, ModbusFunction.WRITE_MULTIPLE_REGISTERS]:
            # Write operations - return success status
            return True
            
        return None
    
    def get_statistics(self) -> Dict[str, Any]:
        """Get client operation statistics."""
        uptime = datetime.now() - self.stats["uptime_start"]
        
        return {
            **self.stats,
            "is_connected": self.is_connected,
            "consecutive_errors": self.consecutive_errors,
            "uptime_hours": uptime.total_seconds() / 3600,
            "success_rate": (
                self.stats["successful_requests"] / max(1, self.stats["total_requests"]) * 100
                if self.stats["total_requests"] > 0 else 0
            ),
            "last_successful_read": self.last_successful_read.isoformat(),
        }
    
    def reset_statistics(self) -> None:
        """Reset operation statistics."""
        self.stats = {
            "total_requests": 0,
            "successful_requests": 0,
            "failed_requests": 0,
            "connection_errors": 0,
            "timeout_errors": 0,
            "last_error": None,
            "uptime_start": datetime.now()
        }
    
    async def health_check(self) -> bool:
        """Perform a health check on the Modbus connection."""
        try:
            # Try to read a register to test communication
            response = await self.read_holding_registers(0, 1)
            return response.success
        except Exception as e:
            logger.error(f"Health check failed: {e}")
            return False


# Global Modbus client instance
modbus_client = AsyncModbusClient()