package main

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

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

// UDPClient manages the UDP connection to one MARTe2 streamer source.
type UDPClient struct {
	serverAddr string
	sourceID   string
	hub        *Hub

	stopCh chan struct{}
}

// NewUDPClient creates a UDPClient bound to a specific source ID. Call Run() in a goroutine.
func NewUDPClient(serverAddr, sourceID string, hub *Hub) *UDPClient {
	return &UDPClient{
		serverAddr: serverAddr,
		sourceID:   sourceID,
		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:
		}

		u.hub.SetSourceState(u.sourceID, "connecting")
		log.Printf("[%s] udp: connecting to %s", u.sourceID, u.serverAddr)
		err := u.runSession()
		if err != nil {
			log.Printf("[%s] udp: session ended: %v", u.sourceID, err)
		}
		u.hub.SetSourceState(u.sourceID, "disconnected")

		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 {
	// Port 0 lets the OS pick a free local port.
	conn, err := net.ListenUDP("udp4", &net.UDPAddr{})
	if err != nil {
		return err
	}
	defer conn.Close()

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

	if _, err := conn.WriteToUDP(BuildConnectPacket(), serverAddr); err != nil {
		return err
	}
	log.Printf("[%s] udp: sent CONNECT", u.sourceID)

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

	for {
		conn.SetReadDeadline(time.Now().Add(silenceTimeout))

		n, _, err := conn.ReadFromUDP(buf)
		arrivalTime := time.Now()
		if err != nil {
			conn.WriteToUDP(BuildDisconnectPacket(), serverAddr)
			return err
		}

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

		hdr, err := ParseHeader(buf[:n])
		if err != nil {
			log.Printf("[%s] udp: parse header: %v", u.sourceID, 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("[%s] udp: parse config: %v", u.sourceID, err)
				continue
			}
			currentSigs = sigs
			log.Printf("[%s] udp: received CONFIG (%d signals)", u.sourceID, len(sigs))
			u.hub.SetSourceState(u.sourceID, "connected")
			u.hub.UpdateConfigForSource(u.sourceID, sigs)

		case PktData:
			if len(currentSigs) == 0 {
				continue
			}
			sample, err := ParseData(complete, currentSigs, arrivalTime)
			if err != nil {
				log.Printf("[%s] udp: parse data: %v", u.sourceID, err)
				continue
			}
			u.hub.PushDataForSource(u.sourceID, sample)

		case PktACK:
			log.Printf("[%s] udp: received ACK (counter=%d)", u.sourceID, hdr.Counter)

		case PktDisconnect:
			log.Printf("[%s] udp: server sent DISCONNECT", u.sourceID)
			return nil

		default:
			log.Printf("[%s] udp: unknown packet type %d", u.sourceID, hdr.Type)
		}

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