package main

import (
	"io"
	"log"
	"net"
	"strconv"
	"time"
)

const (
	silenceTimeout = 5 * time.Second
	reconnectDelay = 2 * time.Second
	readBufSize    = 65536
	udpRcvBufSize  = 8 * 1024 * 1024 // 8 MB OS receive buffer — absorbs bursts at high data rates
)

// UDPClient manages the connection to one MARTe2 streamer source.
// In unicast mode (multicastGroup == "") it uses a single UDP socket for control and data.
// In multicast mode it uses TCP for control (CONNECT/CONFIG/DISCONNECT) and
// joins a UDP multicast group to receive DATA packets.
type UDPClient struct {
	serverAddr     string
	sourceID       string
	hub            *Hub
	multicastGroup string // "" = unicast mode
	dataPort       int    // UDP data port in multicast mode; 0 = serverAddr port+1

	stopCh chan struct{}
}

// NewUDPClient creates a UDPClient bound to a specific source ID. Call Run() in a goroutine.
// multicastGroup: multicast IP (e.g. "239.0.0.1") or "" for unicast.
// dataPort: UDP data port; 0 = serverAddr port+1.
func NewUDPClient(serverAddr, sourceID string, hub *Hub, multicastGroup string, dataPort int) *UDPClient {
	return &UDPClient{
		serverAddr:     serverAddr,
		sourceID:       sourceID,
		hub:            hub,
		multicastGroup: multicastGroup,
		dataPort:       dataPort,
		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] connecting to %s", u.sourceID, u.serverAddr)

		var err error
		if u.multicastGroup != "" {
			err = u.runMulticastSession()
		} else {
			err = u.runSession()
		}
		if err != nil {
			log.Printf("[%s] 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()

	// Increase OS receive buffer to reduce kernel-level packet drops at high data rates.
	if err := conn.SetReadBuffer(udpRcvBufSize); err != nil {
		log.Printf("[%s] udp: SetReadBuffer: %v (proceeding with OS default)", u.sourceID, err)
	}

	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
	var currentPublishMode uint8

	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 hdr.Type == PktData {
			u.hub.RecordDataFragment(u.sourceID, hdr.Counter, n, arrivalTime.UnixNano(), ok)
		}
		if !ok {
			continue
		}

		switch hdr.Type {
		case PktConfig:
			sigs, pm, err := ParseConfig(complete)
			if err != nil {
				log.Printf("[%s] udp: parse config: %v", u.sourceID, err)
				continue
			}
			currentSigs = sigs
			currentPublishMode = pm
			log.Printf("[%s] udp: received CONFIG (%d signals, publishMode=%d)", u.sourceID, len(sigs), pm)
			u.hub.SetSourceState(u.sourceID, "connected")
			u.hub.UpdateConfigForSource(u.sourceID, sigs)

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

		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:
		}
	}
}

// runMulticastSession handles the multicast mode session:
//   - TCP connection to serverAddr for control (CONNECT → CONFIG, DISCONNECT)
//   - UDP multicast socket joined to multicastGroup:dataPort for DATA packets
func (u *UDPClient) runMulticastSession() error {
	// 1. Resolve TCP control address
	tcpAddr, err := net.ResolveTCPAddr("tcp4", u.serverAddr)
	if err != nil {
		return err
	}

	// 2. Dial TCP control connection
	tcpConn, err := net.DialTCP("tcp4", nil, tcpAddr)
	if err != nil {
		return err
	}
	defer tcpConn.Close()

	// 3. Send CONNECT via TCP
	if _, err := tcpConn.Write(BuildConnectPacket()); err != nil {
		return err
	}
	log.Printf("[%s] tcp: sent CONNECT to %s", u.sourceID, u.serverAddr)

	// 4. Read CONFIG header via TCP (io.ReadFull to handle partial reads)
	hdrBuf := make([]byte, HeaderSize)
	if _, err := io.ReadFull(tcpConn, hdrBuf); err != nil {
		return err
	}
	cfgHdr, err := ParseHeader(hdrBuf)
	if err != nil {
		return err
	}
	if cfgHdr.Type != PktConfig {
		return net.ErrClosed // unexpected packet type
	}
	cfgPayload := make([]byte, cfgHdr.PayloadBytes)
	if cfgHdr.PayloadBytes > 0 {
		if _, err := io.ReadFull(tcpConn, cfgPayload); err != nil {
			return err
		}
	}
	currentSigs, currentPublishMode, err := ParseConfig(cfgPayload)
	if err != nil {
		return err
	}
	log.Printf("[%s] tcp: received CONFIG (%d signals, publishMode=%d)", u.sourceID, len(currentSigs), currentPublishMode)
	u.hub.SetSourceState(u.sourceID, "connected")
	u.hub.UpdateConfigForSource(u.sourceID, currentSigs)

	// 5. Determine multicast data port (dataPort or serverAddr port+1)
	mcastPort := u.dataPort
	if mcastPort == 0 {
		mcastPort = tcpAddr.Port + 1
	}

	// 6. Join multicast group for DATA
	mcastIP := net.ParseIP(u.multicastGroup)
	if mcastIP == nil {
		return &net.AddrError{Err: "invalid multicast group IP", Addr: u.multicastGroup}
	}
	mcastAddr := &net.UDPAddr{IP: mcastIP, Port: mcastPort}
	mcastConn, err := net.ListenMulticastUDP("udp4", nil, mcastAddr)
	if err != nil {
		return err
	}
	defer mcastConn.Close()
	if err := mcastConn.SetReadBuffer(udpRcvBufSize); err != nil {
		log.Printf("[%s] multicast SetReadBuffer: %v", u.sourceID, err)
	}
	log.Printf("[%s] joined multicast %s:%s", u.sourceID, u.multicastGroup,
		strconv.Itoa(mcastPort))

	// 7. Background goroutine: watch TCP conn for DISCONNECT or closure
	tcpDone := make(chan error, 1)
	go func() {
		buf := make([]byte, HeaderSize+64)
		for {
			n, readErr := tcpConn.Read(buf)
			if readErr != nil {
				tcpDone <- readErr
				return
			}
			if n >= HeaderSize {
				hdr, parseErr := ParseHeader(buf[:n])
				if parseErr == nil && hdr.Type == PktDisconnect {
					tcpDone <- nil
					return
				}
			}
		}
	}()

	// 8. Main loop: receive DATA packets from multicast group
	reassembler := NewReassembler(2 * time.Second)
	buf := make([]byte, readBufSize)

	for {
		mcastConn.SetReadDeadline(time.Now().Add(silenceTimeout))
		n, _, readErr := mcastConn.ReadFromUDP(buf)
		arrivalTime := time.Now()
		if readErr != nil {
			select {
			case <-tcpDone:
				return nil
			default:
			}
			tcpConn.Write(BuildDisconnectPacket())
			return readErr
		}

		if n < HeaderSize {
			continue
		}
		hdr, parseErr := ParseHeader(buf[:n])
		if parseErr != nil {
			log.Printf("[%s] multicast: parse header: %v", u.sourceID, parseErr)
			continue
		}

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

		complete, ok := reassembler.AddFragment(hdr, payload)
		if hdr.Type == PktData {
			u.hub.RecordDataFragment(u.sourceID, hdr.Counter, n,
				arrivalTime.UnixNano(), ok)
		}
		if !ok {
			continue
		}

		if hdr.Type == PktData {
			if len(currentSigs) == 0 {
				continue
			}
			samples, parseErr := ParseData(complete, currentSigs, currentPublishMode, arrivalTime)
			if parseErr != nil {
				log.Printf("[%s] multicast: parse data: %v", u.sourceID, parseErr)
				continue
			}
			for _, s := range samples {
				u.hub.PushDataForSource(u.sourceID, s)
			}
		}

		select {
		case <-u.stopCh:
			tcpConn.Write(BuildDisconnectPacket())
			return nil
		case tcpErr := <-tcpDone:
			log.Printf("[%s] tcp control closed: %v", u.sourceID, tcpErr)
			return nil
		default:
		}
	}
}