import threading
import time
from pymodbus.client import ModbusTcpClient
from pymodbus.exceptions import ModbusException
from typing import List, Dict, Any, Optional
from queue import Queue
from src.utils.logger import get_logger
from src.core.models.modbus_load_test import LoadTestConfig, WriteTask, ModbusDataType

logger = get_logger("modbus_load_tester")


class ModbusLoadTester:
    def __init__(self, config: LoadTestConfig):
        """
        Инициализация нагрузочного тестера

        :param config: Конфигурация теста
        """
        self.config = config
        self.task_queue = Queue()
        self.threads = []
        self.stop_event = threading.Event()
        self.stats = {
            'total_requests': 0,
            'successful_requests': 0,
            'failed_requests': 0,
            'thread_errors': 0,
            'start_time': 0.0,
            'end_time': 0.0
        }

        # Заполняем очередь задач
        for task in config.tasks:
            for _ in range(task.count if task.count > 0 else 1):
                self.task_queue.put(task)

    def run(self) -> Dict[str, Any]:
        """
        Запуск нагрузочного теста
        """
        logger.info(f"🚀 Starting Modbus load test on {self.config.host}:{self.config.port}")
        logger.info(f"Threads: {self.config.threads}, Duration: {self.config.duration} sec")

        self.stats['start_time'] = time.time()
        self.stop_event.clear()

        # Создаем и запускаем потоки
        for i in range(self.config.threads):
            thread = threading.Thread(
                target=self._worker_thread,
                name=f"ModbusWorker-{i + 1}",
                daemon=True
            )
            thread.start()
            self.threads.append(thread)

        # Запускаем таймер остановки
        if self.config.duration > 0:
            timer = threading.Timer(
                self.config.duration,
                self.stop
            )
            timer.start()

        # Ожидаем завершения всех потоков
        for thread in self.threads:
            thread.join()

        self.stats['end_time'] = time.time()
        self._print_summary()

        return self.stats

    def stop(self):
        """
        Остановка нагрузочного теста
        """
        self.stop_event.set()
        logger.info("🛑 Stopping load test")

    def _worker_thread(self):
        """
        Рабочий поток для выполнения задач
        """
        thread_name = threading.current_thread().name
        logger.debug(f"{thread_name} started")

        client = None

        try:
            while not self.stop_event.is_set():
                try:
                    # Получаем задачу из очереди
                    task = self.task_queue.get_nowait()
                except Exception:
                    # Если очередь пуста, завершаем поток
                    break

                try:
                    # Создаем подключение при необходимости
                    if client is None or not client.is_socket_open():
                        client = self._create_client()

                    # Выполняем запись
                    success = self._execute_write(client, task)
                    self.stats['total_requests'] += 1

                    if success:
                        self.stats['successful_requests'] += 1
                    else:
                        self.stats['failed_requests'] += 1

                    # Возвращаем задачу в очередь, если нужно повторять
                    if task.count == 0:
                        self.task_queue.put(task)

                    # Задержка перед следующей записью
                    time.sleep(task.interval)

                except Exception as e:
                    logger.error(f"{thread_name} error: {e}")
                    self.stats['thread_errors'] += 1
                    client = None  # Принудительное переподключение
                    time.sleep(1)

        except Exception as e:
            logger.error(f"Critical error in {thread_name}: {e}")
            self.stats['thread_errors'] += 1
        finally:
            if client:
                client.close()
            logger.debug(f"{thread_name} exited")

    def _create_client(self) -> ModbusTcpClient:
        """
        Создание и подключение клиента Modbus
        """
        client = ModbusTcpClient(
            host=self.config.host,
            port=self.config.port,
            timeout=self.config.connection_timeout
        )

        if not client.connect():
            raise ConnectionError(f"Failed to connect to {self.config.host}:{self.config.port}")

        logger.debug(f"Connected to {self.config.host}:{self.config.port}")
        return client

    def _execute_write(self, client: ModbusTcpClient, task: WriteTask) -> bool:
        """
        Выполнение операции записи
        """
        try:
            if task.data_type == 'coil':
                result = client.write_coil(
                    address=task.address,
                    value=bool(task.value)
                )
            elif task.data_type == 'holding_register':
                result = client.write_register(
                    address=task.address,
                    value=int(task.value)
                )
            else:
                logger.warning(f"Unsupported data type: {task.data_type}")
                return False

            if result.isError():
                logger.warning(
                    f"Write error: {task.data_type}@{task.address} = {task.value}"
                )
                return False

            logger.debug(
                f"Write success: {task.data_type}@{task.address} = {task.value}"
            )
            return True

        except ModbusException as e:
            logger.error(
                f"Modbus error writing {task.data_type}@{task.address}: {e}"
            )
            return False

    def _print_summary(self):
        """
        Вывод статистики теста
        """
        duration = self.stats['end_time'] - self.stats['start_time']
        requests_per_sec = self.stats['total_requests'] / duration if duration > 0 else 0

        logger.success("📊 Load test summary:")
        logger.success(f"  Total time: {duration:.2f} seconds")
        logger.success(f"  Total requests: {self.stats['total_requests']}")
        logger.success(f"  Successful requests: {self.stats['successful_requests']}")
        logger.success(f"  Failed requests: {self.stats['failed_requests']}")
        logger.success(f"  Thread errors: {self.stats['thread_errors']}")
        logger.success(f"  Requests per second: {requests_per_sec:.2f}")