package station

import (
	"encoding/binary"
	"errors"
	"io"
	"net"
	"time"
)

func (s *Station) finishV3UDP(conn net.Conn, ex *exchange) error {
	udpConn, peer, err := dialAudio(conn)
	if err != nil {
		s.dropAudio()
		return err
	}
	ch := make(chan []byte, 32)
	errc := make(chan error, 1)
	go readUDP(udpConn, peer, ch, errc, ex.done)
	s.mu.Lock()
	s.udp = udpConn
	s.udpPeer = peer
	s.udpWait = true
	s.txCounter = 1
	s.rxCounter = 0
	err = s.sendRegistration()
	s.mu.Unlock()
	if err != nil {
		s.dropAudio()
		udpConn.Close()
		return err
	}
	ex.udpCh = ch
	ex.udpErr = errc
	ex.udpReady = true
	ex.ph = phBeat
	return nil
}

func dialAudio(conn net.Conn) (*net.UDPConn, *net.UDPAddr, error) {
	remote, ok := conn.RemoteAddr().(*net.TCPAddr)
	if !ok {
		return nil, nil, errors.New("tcp peer has no address")
	}
	local := &net.UDPAddr{Port: 0}
	if tcpLocal, ok := conn.LocalAddr().(*net.TCPAddr); ok && tcpLocal.IP != nil {
		local.IP = tcpLocal.IP
	}
	udpConn, err := net.ListenUDP("udp", local)
	if err != nil {
		return nil, nil, err
	}
	return udpConn, &net.UDPAddr{IP: append(net.IP(nil), remote.IP...), Port: remote.Port}, nil
}

func (s *Station) dropAudio() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.enc != nil {
		s.enc.Close()
		s.enc = nil
	}
	if s.dec != nil {
		s.dec.Close()
		s.dec = nil
	}
}

func readUDP(conn *net.UDPConn, peer *net.UDPAddr, out chan<- []byte, errc chan<- error, done <-chan struct{}) {
	buf := make([]byte, 2048)
	for {
		n, addr, err := conn.ReadFromUDP(buf)
		if err != nil {
			select {
			case errc <- err:
			case <-done:
			}
			return
		}
		if addr == nil || peer == nil || !addr.IP.Equal(peer.IP) || addr.Port != peer.Port {
			continue
		}
		pkt := append([]byte(nil), buf[:n]...)
		select {
		case out <- pkt:
		case <-done:
			return
		}
	}
}

func (s *Station) onUDP(w io.Writer, pkt []byte, ex *exchange, udpRx *time.Time, now time.Time) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.secure {
		plain, ok := s.openSecure(pkt)
		if !ok {
			return nil
		}
		*udpRx = now.Add(60 * time.Second)
		if len(plain) < 2 {
			return nil
		}
		typ := binary.BigEndian.Uint16(plain[:2])
		body := plain[2:]
		if s.udpWait {
			if typ != udpHeartbeat {
				return nil
			}
			s.udpWait = false
			s.mu.Unlock()
			var tg [4]byte
			binary.BigEndian.PutUint32(tg[:], uint32(s.cfg.talkGroup))
			err := writeFrame(w, tcpSelectTG, tg[:])
			s.mu.Lock()
			if err != nil {
				return err
			}
			ex.lastTCP = s.clock.Now()
			s.state.Connection = "Connected"
			s.up = true
			s.outage = false
			s.noteTransmitLocked()
			s.writeLog(s.clock.Now(), "connected")
			s.publish()
			ex.ph = phRun
			return nil
		}
		s.onPlain(typ, body, now)
		return nil
	}
	if !s.acceptV2(pkt) {
		return nil
	}
	*udpRx = now.Add(60 * time.Second)
	typ := binary.BigEndian.Uint16(pkt[:2])
	s.onPlain(typ, pkt[6:], now)
	return nil
}

func (s *Station) openSecure(pkt []byte) ([]byte, bool) {
	if len(pkt) < 12 || s.udpKey == nil {
		return nil, false
	}
	counter := binary.BigEndian.Uint32(pkt[:4])
	if counter < s.rxCounter {
		return nil, false
	}
	tag := pkt[4:12]
	ct := pkt[12:]
	plain, err := openDatagram(s.udpKey, s.ivPrefix, 0, counter, pkt[:4], tag, ct)
	if err != nil {
		return nil, false
	}
	s.rxCounter = counter + 1
	return plain, true
}

func (s *Station) acceptV2(pkt []byte) bool {
	if len(pkt) < 6 {
		return false
	}
	if binary.BigEndian.Uint16(pkt[2:4]) != s.clientID {
		return false
	}
	seq := binary.BigEndian.Uint16(pkt[4:6])
	diff := seq - s.rxSeq
	if diff > 32767 {
		return false
	}
	s.rxSeq = seq + 1
	return true
}

func (s *Station) onPlain(typ uint16, body []byte, now time.Time) {
	switch typ {
	case udpAllFlushed:
		s.flushWaiting = false
	case udpAudio:
		if s.dec == nil || s.transmitting() {
			return
		}
		packet, _, err := parseBlob(body)
		if err != nil {
			return
		}
		pcm, err := s.dec.Decode(packet)
		if err != nil {
			return
		}
		s.play = append(s.play, pcm...)
		s.roomBusy = true
		s.audioHeard = true
		s.lastAudio = now
		s.publish()
	}
}

func (s *Station) sendRegistration() error {
	return s.sendSecure(udpHeartbeat, nil, 0, true)
}

func (s *Station) sendUDP(typ uint16, body []byte) error {
	if s.udp == nil {
		return errors.New("udp is down")
	}
	if s.secure {
		return s.sendSecure(typ, body, s.txCounter, false)
	}
	buf := make([]byte, 6+len(body))
	binary.BigEndian.PutUint16(buf[0:2], typ)
	binary.BigEndian.PutUint16(buf[2:4], s.clientID)
	binary.BigEndian.PutUint16(buf[4:6], s.seq)
	s.seq++
	copy(buf[6:], body)
	if _, err := s.udp.WriteToUDP(buf, s.udpPeer); err != nil {
		return err
	}
	s.lastUDPSend = s.clock.Now()
	return nil
}

func (s *Station) sendSecure(typ uint16, body []byte, counter uint32, reg bool) error {
	plain := appendU16(nil, typ)
	plain = append(plain, body...)
	var aad []byte
	id := s.clientID
	if reg {
		aad = registrationAAD(s.clientID)
		counter = 0
	} else {
		aad = counterAAD(counter)
	}
	wire, err := sealDatagram(s.udpKey, s.ivPrefix, id, counter, aad, plain)
	if err != nil {
		return err
	}
	if _, err := s.udp.WriteToUDP(wire, s.udpPeer); err != nil {
		return err
	}
	s.lastUDPSend = s.clock.Now()
	if !reg {
		s.txCounter = counter + 1
	}
	return nil
}