package main

import (
	"log"
	"net"
	"time"
)

const (
	silenceTimeout = 5 * time.Second
	reconnectDelay = 2 * time.Second
	readBufSize    = 65536
)

// UDPClient manages the UDP connection to the MARTe2 streamer.
type UDPClient struct {
	serverAddr string
	localPort  int
	hub        *Hub

	stopCh chan struct{}
}

// NewUDPClient creates a UDPClient. Call Run() in a goroutine.
func NewUDPClient(serverAddr string, localPort int, hub *Hub) *UDPClient {
	return &UDPClient{
		serverAddr: serverAddr,
		localPort:  localPort,
		hub:        hub,
		stopCh:     make(chan struct{}),
	}
}

// Stop asks the client to disconnect and exit.
func (u *UDPClient) Stop() {
	close(u.stopCh)
}

// Run is the main loop; it reconnects automatically if the server goes silent.
func (u *UDPClient) Run() {
	for {
		select {
		case <-u.stopCh:
			return
		default:
		}

		log.Printf("udp: connecting to %s (local port %d)", u.serverAddr, u.localPort)
		err := u.runSession()
		if err != nil {
			log.Printf("udp: session ended: %v", err)
		}

		// Check stop before sleeping.
		select {
		case <-u.stopCh:
			return
		case <-time.After(reconnectDelay):
		}
	}
}

// runSession opens a UDP socket, sends CONNECT, reads data until the server
// goes silent or an error occurs.
func (u *UDPClient) runSession() error {
	localAddr := &net.UDPAddr{Port: u.localPort}
	conn, err := net.ListenUDP("udp4", localAddr)
	if err != nil {
		return err
	}
	defer conn.Close()

	serverAddr, err := net.ResolveUDPAddr("udp4", u.serverAddr)
	if err != nil {
		return err
	}

	// Send CONNECT.
	if _, err := conn.WriteToUDP(BuildConnectPacket(), serverAddr); err != nil {
		return err
	}
	log.Printf("udp: sent CONNECT")

	reassembler := NewReassembler(2 * time.Second)
	buf := make([]byte, readBufSize)
	var currentSigs []SignalInfo

	for {
		// Silence timeout.
		conn.SetReadDeadline(time.Now().Add(silenceTimeout))

		n, _, err := conn.ReadFromUDP(buf)
		arrivalTime := time.Now()
		if err != nil {
			// Send DISCONNECT best-effort before returning.
			conn.WriteToUDP(BuildDisconnectPacket(), serverAddr)
			return err
		}

		if n < HeaderSize {
			log.Printf("udp: short datagram (%d bytes), skipping", n)
			continue
		}

		hdr, err := ParseHeader(buf[:n])
		if err != nil {
			log.Printf("udp: parse header: %v", err)
			continue
		}

		payload := make([]byte, n-HeaderSize)
		copy(payload, buf[HeaderSize:n])

		complete, ok := reassembler.AddFragment(hdr, payload)
		if !ok {
			continue
		}

		switch hdr.Type {
		case PktConfig:
			sigs, err := ParseConfig(complete)
			if err != nil {
				log.Printf("udp: parse config: %v", err)
				continue
			}
			currentSigs = sigs
			log.Printf("udp: received CONFIG (%d signals)", len(sigs))
			u.hub.UpdateConfig(sigs)

		case PktData:
			if len(currentSigs) == 0 {
				// We haven't received a CONFIG yet – ignore.
				continue
			}
			sample, err := ParseData(complete, currentSigs, arrivalTime)
			if err != nil {
				log.Printf("udp: parse data: %v", err)
				continue
			}
			u.hub.PushData(sample)

		case PktACK:
			log.Printf("udp: received ACK (counter=%d)", hdr.Counter)

		case PktDisconnect:
			log.Printf("udp: server sent DISCONNECT")
			return nil

		default:
			log.Printf("udp: unknown packet type %d", hdr.Type)
		}

		// Check stop request.
		select {
		case <-u.stopCh:
			conn.WriteToUDP(BuildDisconnectPacket(), serverAddr)
			return nil
		default:
		}
	}
}