Imports System.IO
Imports System.Net
Imports System.Net.NetworkInformation
Imports System.Net.Sockets
Imports System.Threading
Imports Newtonsoft.Json.Linq

Public Class UDPManager
    Inherits ConnectionManagerBase

    Private ReadOnly client As UdpClient
    Private listenTask As Task = Nothing
    Private ReadOnly cancelTokenSource As New CancellationTokenSource
    Private ReadOnly cancelToken As CancellationToken = cancelTokenSource.Token
    Public ReadOnly Property MaxBufferSize As Integer = 8192

    Public Sub New(port As Integer)
        MyBase.New(port)
        client = New UdpClient(port)
    End Sub

    Protected Overrides Sub Finalize()
        MyBase.Finalize()
        client.Client.Shutdown(SocketShutdown.Both)
        client.Close()
    End Sub

    Public Overrides Sub Start()
        If Not IsListening Then
            StartInternal()
        End If
    End Sub

    Private Sub StartInternal()
        IsListening = True
        listenTask = Task.Run(
            Sub()
                Try
                    RetriveAnyData(cancelToken)
                Catch ex As Exception
                    IsListening = False
                End Try
            End Sub)
    End Sub

    Public Overrides Sub [Stop]()
        If IsListening Then
            IsListening = False
            cancelTokenSource.Cancel()
            listenTask.Wait()
        End If
    End Sub

    Protected Overrides Function GetBufferSize() As Integer
        Return MaxBufferSize
    End Function

    Private Sub RetriveAnyData(ct As CancellationToken)
        Dim doExit = Sub() ct.ThrowIfCancellationRequested()
        Dim receiveTask As Task(Of UdpReceiveResult) = client.ReceiveAsync()

        'Wait for the data and cancel if requested
        receiveTask.Wait(ct)

        Dim buf As Byte() = receiveTask.Result.Buffer
        Dim ip As String = receiveTask.Result.RemoteEndPoint.Address.ToString

        doExit()

        ProcessRetrivedData(ip, buf)

        doExit()

        StartInternal()
    End Sub

    Protected Overrides Sub SendData(ep As IPEndPoint, buf As Byte())
        Dim udp As New UdpClient
        udp.Connect(ep)
        udp.Send(buf, buf.Length)
        udp.Client.Shutdown(SocketShutdown.Both)
        udp.Close()
    End Sub

End Class