#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
TFTP сервер для передачи прошивки с компьютера на коммутатор.

- Создает сервер по заданному IP и порту.
- Расшаривает папку Firmware.
- Показывает текущее состояние сервера и статус передачи файла:
  - кому (IP устройства),
  - сколько осталось байт,
  - сколько передано байт,
  - время передачи.
"""

import os
import socket
import struct
import threading
import time

class TFTPServer:
    def __init__(self, share_folder):
        """
        Инициализация TFTP сервера.

        :param share_folder: Путь к папке, содержащей файлы (например, папка 'Firmware')
        """
        self.share_folder = share_folder
        self.log_callback = None
        self.running = False
        self.server_socket = None
        self.lock = threading.Lock()
        self.transfer_sockets = set()  # Множество для хранения всех активных сокетов передачи
        # Словарь активных передач для мониторинга их статуса.
        # Ключ – адрес клиента, значение – словарь с информацией о передаче.
        self.active_transfers = {}

    def set_log_callback(self, callback):
        """
        Установка функции обратного вызова для логирования сообщений.

        :param callback: Функция, принимающая строку сообщения.
        """
        self.log_callback = callback

    def log(self, message):
        """
        Функция логирования: вызывает callback (если задан) или выводит сообщение в консоль.
        
        :param message: Строка с сообщением для логирования.
        """
        if self.log_callback:
            self.log_callback(message)
        else:
            print(message)

    def start_server(self, ip, port):
        """
        Запуск TFTP сервера на указанном IP и порту.

        :param ip: IP-адрес для привязки сервера.
        :param port: Порт для TFTP сервера.
        """
        if self.running:
            self.log("[WARN] Сервер уже запущен")
            return

        self.running = True
        try:
            self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            self.server_socket.bind((ip, port))
            self.log(f"[INFO] TFTP сервер запущен на {ip}:{port}")
            
            while self.running:
                try:
                    self.server_socket.settimeout(1.0)
                    data, client_addr = self.server_socket.recvfrom(2048)
                    if data and self.running:
                        threading.Thread(target=self.handle_request, args=(data, client_addr), daemon=True).start()
                except socket.timeout:
                    continue
                except socket.error as e:
                    if self.running:  # Логируем ошибку только если сервер еще запущен
                        self.log(f"[ERROR] Ошибка получения данных: {str(e)}")
                    break
                except Exception as e:
                    if self.running:  # Логируем ошибку только если сервер еще запущен
                        self.log(f"[ERROR] Ошибка получения данных: {str(e)}")
                    break
        except Exception as e:
            self.log(f"[ERROR] Ошибка запуска сервера: {str(e)}")
        finally:
            self.running = False
            if self.server_socket:
                try:
                    self.server_socket.close()
                except:
                    pass
                self.server_socket = None

    def stop_server(self):
        """
        Остановка TFTP сервера.
        """
        if not self.running:
            return

        self.log("[INFO] Остановка TFTP сервера...")
        self.running = False

        try:
            # Закрываем основной сокет сервера первым
            if self.server_socket:
                try:
                    # Создаем временный сокет и отправляем пакет самому себе,
                    # чтобы разблокировать recvfrom
                    temp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
                    try:
                        server_address = self.server_socket.getsockname()
                        temp_socket.sendto(b'', server_address)
                    except:
                        pass
                    finally:
                        try:
                            temp_socket.close()
                        except:
                            pass
                    
                    try:
                        self.server_socket.close()
                    except Exception as e:
                        self.log(f"[WARN] Ошибка при закрытии основного сокета: {str(e)}")
                finally:
                    self.server_socket = None

            # Закрываем все активные сокеты передачи
            with self.lock:
                active_sockets = list(self.transfer_sockets)
                self.transfer_sockets.clear()
                active_transfers = dict(self.active_transfers)
                self.active_transfers.clear()

            # Закрываем сокеты передачи после очистки множества
            for sock in active_sockets:
                try:
                    if sock:
                        sock.close()
                except Exception as e:
                    self.log(f"[WARN] Ошибка при закрытии сокета передачи: {str(e)}")

            # Отправляем сообщения об остановке для активных передач
            for client_addr, transfer_info in active_transfers.items():
                try:
                    self.send_error(client_addr, 0, "Сервер остановлен")
                except:
                    pass

        except Exception as e:
            self.log(f"[ERROR] Ошибка при остановке сервера: {str(e)}")
        finally:
            self.running = False  # Гарантируем, что флаг running будет False
            self.log("[INFO] TFTP сервер остановлен")

    def handle_request(self, data, client_addr):
        """
        Обработка входящего запроса от клиента.

        :param data: Полученные данные (UDP-пакет).
        :param client_addr: Адрес клиента, отправившего пакет.
        """
        if len(data) < 2:
            self.log(f"[WARN] Получен некорректный пакет от {client_addr}")
            return
        opcode = struct.unpack("!H", data[:2])[0]
        if opcode == 1:  # RRQ (Read Request) – запрос на чтение файла
            self.handle_rrq(data, client_addr)
        else:
            self.log(f"[WARN] Неподдерживаемый запрос (опкод {opcode}) от {client_addr}")

    def handle_rrq(self, data, client_addr):
        """
        Обработка запроса на чтение файла (RRQ).

        :param data: Данные запроса.
        :param client_addr: Адрес клиента.
        """
        try:
            # RRQ формата: 2 байта опкода, затем строка имени файла, за которой следует 0,
            # затем строка режима (например, "octet"), и завершается 0.
            parts = data[2:].split(b'\0')
            if len(parts) < 2:
                self.log(f"[WARN] Некорректный RRQ пакет от {client_addr}")
                return
            filename = parts[0].decode('utf-8')
            mode = parts[1].decode('utf-8').lower()
            self.log(f"[INFO] Получен RRQ от {client_addr}: файл '{filename}', режим '{mode}'")
            if mode != "octet":
                self.send_error(client_addr, 0, "Поддерживается только octet режим")
                return
            file_path = os.path.join(self.share_folder, filename)
            if not os.path.isfile(file_path):
                self.send_error(client_addr, 1, "Файл не найден")
                return
            # Запускаем передачу файла в новом потоке.
            threading.Thread(target=self.send_file, args=(file_path, client_addr), daemon=True).start()
        except Exception as e:
            self.log(f"[ERROR] Ошибка обработки RRQ: {str(e)}")

    def send_error(self, client_addr, error_code, error_message):
        """
        Отправка сообщения об ошибке клиенту.

        :param client_addr: Адрес клиента.
        :param error_code: Код ошибки.
        :param error_message: Текст ошибки.
        """
        # Формируем TFTP пакет ошибки: 2 байта опкода (5), 2 байта кода ошибки, сообщение об ошибке и завершающий 0.
        packet = struct.pack("!HH", 5, error_code) + error_message.encode('utf-8') + b'\0'
        self.server_socket.sendto(packet, client_addr)
        self.log(f"[INFO] Отправлено сообщение об ошибке '{error_message}' клиенту {client_addr}")

    def send_file(self, file_path, client_addr):
        """
        Передача файла клиенту по протоколу TFTP.
        """
        BLOCK_SIZE = 512
        transfer_socket = None
        try:
            if not os.path.exists(file_path):
                self.log(f"[ERROR] Файл '{file_path}' не существует")
                self.send_error(client_addr, 1, "Файл не найден")
                return

            filesize = os.path.getsize(file_path)
            if filesize == 0:
                self.log(f"[ERROR] Файл '{file_path}' пуст")
                self.send_error(client_addr, 0, "Файл пуст")
                return

            start_time = time.time()
            file_basename = os.path.basename(file_path)
            
            # Регистрируем активную передачу
            with self.lock:
                self.active_transfers[client_addr] = {
                    'filename': file_basename,
                    'filesize': filesize,
                    'bytes_sent': 0,
                    'start_time': start_time
                }
            
            self.log(f"[INFO] Начало передачи файла '{file_basename}' клиенту {client_addr}. Размер файла: {filesize} байт.")
            
            # Создаем отдельный сокет для передачи файла
            transfer_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            transfer_socket.settimeout(5.0)  # Таймаут ожидания ACK
            
            # Добавляем сокет в множество активных сокетов
            with self.lock:
                self.transfer_sockets.add(transfer_socket)
            
            block_num = 1
            bytes_sent = 0
            last_progress_time = time.time()

            try:
                with open(file_path, "rb") as f:
                    while self.running:  # Проверяем флаг running
                        try:
                            data_block = f.read(BLOCK_SIZE)
                            if not data_block:  # Достигнут конец файла
                                break

                            # Проверяем флаг running перед отправкой блока
                            if not self.running:
                                raise Exception("Передача прервана: сервер остановлен")

                            # Формируем TFTP пакет данных
                            packet = struct.pack("!HH", 3, block_num) + data_block
                            attempts = 0
                            ack_received = False

                            # Попытка отправки текущего блока (до 3 повторных попыток)
                            while attempts < 3 and not ack_received and self.running:
                                if transfer_socket is None:
                                    raise Exception("Сокет передачи закрыт")
                                    
                                try:
                                    transfer_socket.sendto(packet, client_addr)
                                    
                                    # Логируем прогресс каждую секунду
                                    current_time = time.time()
                                    if current_time - last_progress_time >= 1.0:
                                        elapsed_time = current_time - start_time
                                        remaining = filesize - bytes_sent
                                        self.log(f"[STATUS] Клиент: {client_addr} | Файл: {file_basename} | "
                                               f"Отправлено: {bytes_sent}/{filesize} байт | "
                                               f"Осталось: {remaining} байт | "
                                               f"Время: {elapsed_time:.2f} сек.")
                                        last_progress_time = current_time

                                    # Ожидаем подтверждение
                                    ack_data, addr = transfer_socket.recvfrom(4)
                                    if addr == client_addr:
                                        ack_opcode, ack_block = struct.unpack("!HH", ack_data)
                                        if ack_opcode == 4 and ack_block == block_num:
                                            ack_received = True
                                            bytes_sent += len(data_block)
                                            with self.lock:
                                                if client_addr in self.active_transfers:
                                                    self.active_transfers[client_addr]['bytes_sent'] = bytes_sent
                                        else:
                                            self.log(f"[WARN] Неверный ACK от {client_addr}. "
                                                   f"Ожидался блок {block_num}, получен {ack_block}.")
                                except socket.timeout:
                                    attempts += 1
                                    self.log(f"[WARN] Таймаут ожидания ACK для блока {block_num} "
                                           f"от {client_addr}. Попытка {attempts+1}.")
                                except socket.error as e:
                                    if not self.running:
                                        raise Exception("Передача прервана: сервер остановлен")
                                    self.log(f"[ERROR] Ошибка сокета при отправке блока {block_num}: {str(e)}")
                                    attempts += 1
                                except Exception as e:
                                    if not self.running:
                                        raise Exception("Передача прервана: сервер остановлен")
                                    self.log(f"[ERROR] Ошибка при отправке блока {block_num}: {str(e)}")
                                    attempts += 1

                            if not ack_received:
                                raise Exception(f"Не удалось получить подтверждение для блока {block_num}")

                            block_num = (block_num + 1) % 65536

                        except Exception as e:
                            if not self.running:
                                raise Exception("Передача прервана: сервер остановлен")
                            raise

                if bytes_sent == filesize:
                    elapsed_time = time.time() - start_time
                    self.log(f"[INFO] Передача файла '{file_basename}' клиенту {client_addr} "
                            f"завершена за {elapsed_time:.2f} сек. Всего отправлено {bytes_sent} байт.")
            except Exception as e:
                if not self.running:
                    self.log(f"[INFO] Передача файла '{file_basename}' клиенту {client_addr} прервана: сервер остановлен")
                raise

        except Exception as e:
            if not self.running:
                return  # Не логируем повторно о прерывании передачи
            self.log(f"[ERROR] Ошибка при передаче файла '{os.path.basename(file_path)}' "
                    f"клиенту {client_addr}: {str(e)}")
            try:
                self.send_error(client_addr, 0, str(e))
            except:
                pass
        finally:
            # Закрываем сокет передачи
            if transfer_socket:
                try:
                    with self.lock:
                        self.transfer_sockets.discard(transfer_socket)
                    transfer_socket.close()
                    transfer_socket = None
                except:
                    pass
            
            # Удаляем информацию о передаче
            with self.lock:
                if client_addr in self.active_transfers:
                    del self.active_transfers[client_addr] 