package main

import (
	// context: Используется для управления жизненным циклом приложения.
	// Позволяет мягко (Graceful Shutdown) передать сигнал отмены всем воркерам.
	"context"

	// encoding/binary: Позволяет читать и писать числа в сетевом байтовом порядке
	// (BigEndian / Network Byte Order). Незаменим для парсинга и сборки STUN-заголовков.
	"encoding/binary"

	// errors: Предоставляет базовые инструменты для создания и сравнения ошибок.
	// Используется для создания pre-allocated (заранее созданных) констант ошибок.
	"errors"

	// flag: Реализует парсинг флагов командной строки (аргументов запуска).
	// Нужен для ручной передачи внешнего IP через параметр, например, `--public-ip=91.210.106.64`.
	"flag"

	// log/slog: Современный пакет структурированного логирования.
	// Позволяет выводить системные сообщения в формате "ключ=значение" с высокой скоростью.
	"log/slog"

	// net: Базовый сетевой пакет Go. Используется для открытия сетевого сокета
	// через net.ListenUDP и работы с традиционными UDP-соединениями.
	"net"

	// net/netip: Сверхбыстрый и современный пакет для работы с IP-адресами (с Go 1.18).
	// В отличие от старого net.IP, структуры netip.Addr распределяются в стеке (Zero-Alloc).
	"net/netip"

	// os: Позволяет взаимодействовать с операционной системой хоста.
	// Используется для чтения переменных окружения (os.Getenv) и выходов в случае критических сбоев.
	"os"

	// os/signal: Перехватывает системные сигналы от ОС (например, нажатие Ctrl+C).
	// Нужен для того, чтобы сервер не умирал мгновенно, а корректно закрывал сокеты при остановке.
	"os/signal"

	// runtime: Позволяет управлять планировщиком Go и исследовать среду выполнения.
	// Вызов runtime.NumCPU() определяет количество доступных ядер для пула воркеров.
	"runtime"

	// sync: Предоставляет примитивы синхронизации горутин.
	// Структура sync.WaitGroup гарантирует, что сервер не закроется, пока все воркеры не завершат работу.
	"sync"

	// syscall: Низкоуровневый интерфейс к системным вызовам операционной системы.
	// Нужен для точного перехвата сигнала завершения процесса (syscall.SIGTERM) от Docker.
	"syscall"
)

// Константы протокола STUN (RFC 5389 / RFC 8489)
const (
	Port               = 3478       // Стандартный UDP-порт для STUN-сервисов
	BufferSize         = 1500       // Размер буфера под MTU (максимальный пакет без фрагментации IP)
	STUNHeaderSize     = 20         // Фиксированный размер заголовка STUN (20 байт)
	MagicCookie        = 0x2112A442 // Магическое число STUN для идентификации протокола
	AttrResponseOrigin = 0x802B     // Код атрибута, указывающий исходный IP/порт сервера
)

// Преаллоцированные ошибки для защиты от лишних выделений памяти в рантайме
var (
	errPacketTooShort = errors.New("packet too short")
	errInvalidMethod  = errors.New("invalid message type or method")
	errInvalidCookie  = errors.New("invalid magic cookie")
)

func main() {
	// Инициализируем быстрый структурированный логгер (slog) в текстовом формате
	logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelInfo}))
	slog.SetDefault(logger)

	// Регистрируем флаг командной строки для ручного ввода публичного IP
	flagStr := flag.String("public-ip", "", "Внешний публичный IPv4-адрес сервера (для логов)")
	flag.Parse()

	// Логика приоритета: если флаг пустой, ищем значение в переменных окружения (Docker/.env)
	publicIPStr := *flagStr
	if publicIPStr == "" {
		publicIPStr = os.Getenv("PUBLIC_IP")
	}

	// Парсим и валидируем IP-адрес
	var publicIP netip.Addr
	if publicIPStr != "" {
		var err error
		publicIP, err = netip.ParseAddr(publicIPStr)
		// Сервер строго IPv4, поэтому проверяем Is4()
		if err != nil || !publicIP.Is4() {
			slog.Error("Указан некорректный IPv4 public-ip", "val", publicIPStr)
			os.Exit(1)
		}
	} else {
		// Дефолтный фолбек на localhost, если адрес вообще нигде не указан
		publicIP = netip.AddrFrom4([4]byte{127, 0, 0, 1})
	}

	// Определяем количество логических ядер процессора для запуска оптимального пула воркеров
	numWorkers := runtime.NumCPU()

	// Настраиваем контекст, который слушает сигналы завершения работы ОС (Ctrl+C / SIGTERM от Docker)
	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer cancel()

	// Открываем единый UDP-сокет на всех интерфейсах (0.0.0.0) и порту 3478
	addr := &net.UDPAddr{IP: net.IPv4zero, Port: Port}
	conn, err := net.ListenUDP("udp4", addr)
	if err != nil {
		slog.Error("Не удалось запустить UDP сервер", "error", err)
		os.Exit(1)
	}
	defer conn.Close()

	// Настройка буферов сокета на уровне ядра ОС (4 Мегабайта).
	// Защищает от потери пакетов, пока Go-горутины переключают контекст.
	_ = conn.SetReadBuffer(4194304)
	_ = conn.SetWriteBuffer(4194304)

	// Запускаем пул параллельных воркеров по числу ядер CPU
	var wg sync.WaitGroup
	for i := 0; i < numWorkers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			startWorker(ctx, conn, publicIP)
		}()
	}

	slog.Info("STUN-сервер запущен (True Zero-Alloc, IPv4)", "port", Port, "workers", numWorkers)

	// Блокируем главный поток, ожидая системного сигнала на остановку
	<-ctx.Done()
	slog.Info("Завершение работы сервера...")

	// Принудительно закрываем сокет. Это мгновенно разблокирует методы Read во всех воркерах.
	conn.Close()
	wg.Wait() // Ждем, пока все воркеры корректно завершат текущую итерацию цикла
	slog.Info("Сервер успешно остановлен")
}

func startWorker(ctx context.Context, conn *net.UDPConn, publicIP netip.Addr) {
	// КЛЮЧЕВАЯ ОПТИМИЗАЦИЯ: Буфер выделяется один раз ПЕРЕД циклом для каждого воркера.
	// Переменная не "убегает" в кучу (heap alloc = 0) и переиспользуется для каждого пакета.
	buf := make([]byte, BufferSize)

	for {
		// Неблокирующая быстрая проверка — не приказал ли контекст завершить работу
		select {
		case <-ctx.Done():
			return
		default:
		}

		// Читаем пакет из сетевого сокета прямо в преаллоцированный буфер.
		// addrPort возвращает netip.AddrPort (аллокация структуры в стеке, zero-alloc).
		n, addrPort, err := conn.ReadFromUDPAddrPort(buf)
		if err != nil {
			if errors.Is(err, net.ErrClosed) {
				return // Нормальный выход, если сокет был закрыт главным потоком
			}
			slog.Debug("Ошибка чтения из UDP сокета", "error", err)
			continue
		}

		// Быстрая фильтрация невалидных пакетов или IPv6-запросов (так как сервер под IPv4)
		if !addrPort.IsValid() || !addrPort.Addr().Is4() {
			continue
		}

		// Передаем пакет на разбор и формирование ответа
		_ = handleSTUN(conn, addrPort, buf, n, publicIP)
	}
}

func handleSTUN(conn *net.UDPConn, remoteAddr netip.AddrPort, buf []byte, packetLen int, publicIP netip.Addr) error {
	// Базовая проверка: STUN-пакет не может быть меньше своего заголовка
	if packetLen < STUNHeaderSize {
		return errPacketTooShort
	}

	// Читаем первые 2 байта (Message Type). Сервер обрабатывает только Binding Request (0x0001)
	messageType := binary.BigEndian.Uint16(buf[0:2])
	if messageType != 0x0001 {
		return errInvalidMethod
	}

	// Проверяем корректность Magic Cookie (байты 4-7)
	magicCookie := binary.BigEndian.Uint32(buf[4:8])
	if magicCookie != MagicCookie {
		return errInvalidCookie
	}

	// Извлекаем уникальный Transaction ID (12 байт), чтобы клиент понял, на какой запрос пришел ответ
	var transactionID [12]byte
	copy(transactionID[:], buf[8:20])

	// Получаем IP-адрес клиента в виде массива из 4 байт (стековая структура)
	clientIP := remoteAddr.Addr().As4()

	// Фокус оптимизации: Расширяем видимость среза до полной емкости буфера (1500 байт),
	// чтобы писать ответ прямо поверх данных запроса без аллокации новой памяти.
	buf = buf[:cap(buf)]

	// Смещение (указатель записи). Заголовок пока пропускаем, начнем писать атрибуты с 20-го байта.
	pos := STUNHeaderSize

	// --------------------------------------------------------------------------
	// АТРИБУТ 1: XOR-MAPPED-ADDRESS (0x0020) — возвращает клиенту его внешний IP/Порт с XOR-шифрованием
	// --------------------------------------------------------------------------
	binary.BigEndian.PutUint16(buf[pos:pos+2], 0x0020) // Type
	binary.BigEndian.PutUint16(buf[pos+2:pos+4], 8)    // Length (8 байт на данные)
	buf[pos+4] = 0                                     // Зарезервировано
	buf[pos+5] = 0x01                                  // Protocol Family: IPv4

	// Шифруем порт: делаем XOR порта клиента с первыми 16 битами Magic Cookie
	xorPort := remoteAddr.Port() ^ uint16(MagicCookie>>16)
	binary.BigEndian.PutUint16(buf[pos+6:pos+8], xorPort)

	// Шифруем IP: делаем XOR IP-адреса клиента с полным 32-битным Magic Cookie
	xorIP := binary.BigEndian.Uint32(clientIP[:]) ^ MagicCookie
	binary.BigEndian.PutUint32(buf[pos+8:pos+12], xorIP)
	pos += 12 // Сдвигаем указатель на 12 байт вперед

	// --------------------------------------------------------------------------
	// АТРИБУТ 2: MAPPED-ADDRESS (0x0001) — старый формат без XOR (для совместимости со старым софтом)
	// --------------------------------------------------------------------------
	binary.BigEndian.PutUint16(buf[pos:pos+2], 0x0001)              // Type
	binary.BigEndian.PutUint16(buf[pos+2:pos+4], 8)                 // Length
	buf[pos+4] = 0                                                  // Зарезервировано
	buf[pos+5] = 0x01                                               // Protocol Family: IPv4
	binary.BigEndian.PutUint16(buf[pos+6:pos+8], remoteAddr.Port()) // Порт в чистом виде
	copy(buf[pos+8:pos+12], clientIP[:])                            // IP в чистом виде
	pos += 12

	// --------------------------------------------------------------------------
	// АТРИБУТ 3: RESPONSE-ORIGIN (0x802B) — сообщает клиенту, с какого IP сервера отправлен ответ
	// --------------------------------------------------------------------------
	binary.BigEndian.PutUint16(buf[pos:pos+2], AttrResponseOrigin) // Type
	binary.BigEndian.PutUint16(buf[pos+2:pos+4], 8)                // Length
	buf[pos+4] = 0                                                 // Зарезервировано
	buf[pos+5] = 0x01                                              // Protocol Family: IPv4
	binary.BigEndian.PutUint16(buf[pos+6:pos+8], uint16(Port))     // Порт нашего STUN сервера
	serverIPBytes := publicIP.As4()
	copy(buf[pos+8:pos+12], serverIPBytes[:]) // Внешний IP нашего сервера из переменной окружения
	pos += 12

	// 计算 Длина ответа
	totalLen := pos                      // Полный размер пакета (20 байт заголовок + 36 байт атрибуты = 56 байт)
	attrLen := totalLen - STUNHeaderSize // Чистый размер только блока атрибутов (36 байт)

	// --------------------------------------------------------------------------
	// ЗАПОЛНЕНИЕ СТАРТОВОГО ЗАГОЛОВКА STUN (Первые 20 байт пакета)
	// --------------------------------------------------------------------------
	binary.BigEndian.PutUint16(buf[0:2], 0x0101)          // Message Type: Binding Success Response
	binary.BigEndian.PutUint16(buf[2:4], uint16(attrLen)) // Записываем длину атрибутов в заголовок
	binary.BigEndian.PutUint32(buf[4:8], MagicCookie)     // Возвращаем Magic Cookie
	copy(buf[8:20], transactionID[:])                     // Возвращаем оригинальный Transaction ID клиента

	// Отправляем сформированный срез байт обратно клиенту по UDP
	if _, err := conn.WriteToUDPAddrPort(buf[:totalLen], remoteAddr); err != nil {
		if !errors.Is(err, net.ErrClosed) {
			slog.Debug("Ошибка отправки ответа", "remote", remoteAddr.String(), "error", err.Error())
		}
		return err
	}

	return nil // Пакет успешно обработан без единой аллокации памяти в куче
}