import time

from pymodbus.client import ModbusTcpClient
from pymodbus.exceptions import ModbusException
from typing import Dict, Set, Tuple, Optional, Callable, Any
from pprint import pprint

from src.core.models.modbus_scan import ModbusScanResult
from src.utils.logger import get_logger



class ModbusScanner:
    def __init__(
            self,
            host: str,
            port: int = 502,
            coil_range: Tuple[int, int] = (0, 100),
            discrete_input_range: Tuple[int, int] = (0, 100),
            holding_register_range: Tuple[int, int] = (0, 100),
            input_register_range: Tuple[int, int] = (0, 100),
            scan_duration: int = 10,
            block_size: int = 50
    ):
        """
        Сканер Modbus устройств

        :param host: IP-адрес устройства
        :param port: Порт Modbus (по умолчанию 502)
        :param coil_range: Диапазон сканирования Coils (start, end)
        :param discrete_input_range: Диапазон сканирования Discrete Inputs (start, end)
        :param holding_register_range: Диапазон сканирования Holding Registers (start, end)
        :param input_register_range: Диапазон сканирования Input Registers (start, end)
        :param scan_duration: Длительность сканирования в секундах
        :param block_size: Размер блока для чтения за один запрос
        """
        self.host = host
        self.port = port
        self.coil_range = coil_range
        self.discrete_input_range = discrete_input_range
        self.holding_register_range = holding_register_range
        self.input_register_range = input_register_range
        self.scan_duration = scan_duration
        self.block_size = block_size
        self.client = None
        self.logger = get_logger("modbus_scanner")

    def connect(self) -> bool:
        """Подключение к устройству"""
        self.logger.debug(f"Connecting to {self.host}:{self.port}")
        try:
            self.client = ModbusTcpClient(
                host=self.host,
                port=self.port,
                timeout=2
            )
            connected = self.client.connect()
            if connected:
                self.logger.success(f"Connected to {self.host}:{self.port}")
            else:
                self.logger.error(f"Connection failed to {self.host}:{self.port}")
            return connected
        except Exception as e:
            self.logger.opt(exception=True).error(f"Connection error: {e}")
            return False

    def disconnect(self):
        """Отключение от устройства"""
        if self.client:
            self.client.close()
            self.logger.debug(f"Disconnected from {self.host}:{self.port}")

    def scan_registers(self) -> Optional[ModbusScanResult]:
        """Основной метод сканирования всех типов регистров"""
        self.logger.info(f"Starting Modbus scan on {self.host}:{self.port}")
        self.logger.debug(f"Scan parameters: "
                          f"coils={self.coil_range}, "
                          f"discrete_inputs={self.discrete_input_range}, "
                          f"holding_registers={self.holding_register_range}, "
                          f"input_registers={self.input_register_range}, "
                          f"duration={self.scan_duration}s, "
                          f"block_size={self.block_size}")

        if not self.connect():
            return None

        start_time = time.time()

        # Словари для хранения активных регистров (адрес -> последнее значение)
        active_coils: Dict[int, bool] = {}
        active_discrete_inputs: Dict[int, bool] = {}
        active_holding_registers: Dict[int, int] = {}
        active_input_registers: Dict[int, int] = {}

        # Множества для отслеживания активности (адрес -> было ли ненулевое значение)
        activity_tracker: Dict[str, Set[int]] = {
            'coil': set(),
            'discrete_input': set(),
            'holding_register': set(),
            'input_register': set()
        }

        try:
            scan_count = 0
            while time.time() - start_time < self.scan_duration:
                scan_count += 1
                self.logger.debug(f"Scan iteration #{scan_count}")

                # Сканируем Coils (0x)
                self._scan_block(
                    read_func=self.client.read_coils,
                    address_range=self.coil_range,
                    active_values=active_coils,
                    activity_tracker=activity_tracker['coil'],
                    reg_type="coil",
                    is_register=False
                )

                # Сканируем Discrete Inputs (1x)
                self._scan_block(
                    read_func=self.client.read_discrete_inputs,
                    address_range=self.discrete_input_range,
                    active_values=active_discrete_inputs,
                    activity_tracker=activity_tracker['discrete_input'],
                    reg_type="discrete_input",
                    is_register=False
                )

                # Сканируем Holding Registers (4x)
                self._scan_block(
                    read_func=self.client.read_holding_registers,
                    address_range=self.holding_register_range,
                    active_values=active_holding_registers,
                    activity_tracker=activity_tracker['holding_register'],
                    reg_type="holding_register",
                    is_register=True
                )

                # Сканируем Input Registers (3x) - Holding Registers (Read-Only)
                self._scan_block(
                    read_func=self.client.read_input_registers,
                    address_range=self.input_register_range,
                    active_values=active_input_registers,
                    activity_tracker=activity_tracker['input_register'],
                    reg_type="input_register",
                    is_register=True
                )

                time.sleep(0.5)

            # Формируем результаты
            result = ModbusScanResult(
                ip_address=f"{self.host}:{self.port}",
                active_coils=active_coils,
                active_discrete_inputs=active_discrete_inputs,
                active_holding_registers=active_holding_registers,
                active_input_registers=active_input_registers
            )

            self.logger.success(
                f"Scan completed in {scan_count} iterations. "
                f"Active coils: {len(active_coils)}, "
                f"Active discrete inputs: {len(active_discrete_inputs)}, "
                f"Active holding registers: {len(active_holding_registers)}, "
                f"Active input registers: {len(active_input_registers)}"
            )
            return result

        except ModbusException as e:
            self.logger.opt(exception=True).error(f"Modbus error: {e}")
            return None
        finally:
            self.disconnect()

    def _scan_block(
            self,
            read_func: Callable,
            address_range: Tuple[int, int],
            active_values: Dict[int, Any],
            activity_tracker: Set[int],
            reg_type: str,
            is_register: bool = False
    ):
        """
        Сканирование блока адресов

        :param read_func: Функция чтения регистров
        :param address_range: Диапазон адресов (start, end)
        :param active_values: Словарь для хранения активных регистров
        :param activity_tracker: Множество для отслеживания активности
        :param reg_type: Тип регистра (для логирования)
        :param is_register: Флаг для регистров (16-битные значения)
        """
        start_addr, end_addr = address_range
        current_addr = start_addr

        while current_addr <= end_addr:
            count_val = min(self.block_size, end_addr - current_addr + 1)

            # Проверяем активность подключения
            if not self.client.is_socket_open():
                self.logger.warning("Connection lost during scan, reconnecting...")
                if not self.connect():
                    self.logger.error("Reconnection failed, aborting scan")
                    break

            try:
                # Чтение блока регистров
                response = read_func(address=current_addr, count=count_val)

                if response.isError():
                    self.logger.warning(
                        f"Error reading {reg_type} at address {current_addr}"
                    )
                    current_addr += count_val
                    continue

                # Получаем значения
                if is_register:
                    values = response.registers
                else:
                    values = response.bits

                # Обрабатываем каждое значение в блоке
                for i, value in enumerate(values):
                    addr = current_addr + i

                    # Если адрес выходит за верхнюю границу, прерываем
                    if addr > end_addr:
                        break

                    # Для регистров преобразуем в знаковое целое
                    if is_register and value >= 0x8000:
                        value = value - 0x10000

                    # Определяем активность
                    is_active = False

                    # Регистр считается активным если:
                    # 1. Текущее значение не нулевое
                    # 2. Или ранее было ненулевое значение (даже если сейчас 0)
                    if value != 0:
                        is_active = True
                        # Запоминаем что этот адрес когда-либо был активен
                        activity_tracker.add(addr)
                    elif addr in activity_tracker:
                        # Если ранее было ненулевое значение
                        is_active = True

                    # Обновляем значение только для активных регистров
                    if is_active:
                        active_values[addr] = value
                    elif addr in active_values:
                        # Удаляем неактивные регистры
                        del active_values[addr]

            except ModbusException as e:
                self.logger.opt(exception=True).error(
                    f"Error reading {reg_type} at {current_addr}: {e}"
                )
                current_addr += count_val
                continue

            current_addr += count_val

if __name__ == "__main__":
    host = "192.168.1.55"
    logger = get_logger("modbus_scan")
    scanner = ModbusScanner(
        host=host,
        coil_range=(0, 1000),      # 0x - Coils
        discrete_input_range=(0, 1000), # 1x - Discrete Inputs
        holding_register_range=(0, 1000), # 4x - Holding Registers
        input_register_range=(0, 1000), # 3x - Input Registers
        scan_duration=15,
        block_size=125               # Максимальный размер блока для Modbus
    )
    result = scanner.scan_registers()
    if result:
        logger.success(
            f"✅ Modbus scan completed. "
            f"Coils: {len(result.active_coils)}, "
            f"Discrete Inputs: {len(result.active_discrete_inputs)}, "
            f"Holding Registers: {len(result.active_holding_registers)}, "
            f"Input Registers: {len(result.active_input_registers)}"
        )
    else:
        logger.error("❌ Modbus scan failed")

    pprint(result)