// Copyright 2021 Martin Dosch.
// Use of this source code is governed by the BSD-2-clause
// license that can be found in the LICENSE file.

package xmppsrv

import (
	"errors"
	"math/rand"
	"net"
	"sort"
	"time"
)

// The SRV type is basically the same as net.SRV (see
// https://pkg.go.dev/net#SRV but it adds the string Type
// which contains either xmpp-client, xmpps-client,
// xmpp-server or xmpps-server.
// This is especially useful for the function LookupXmppClient
// which returns a mix of xmpp-client and xmpps-client
// records.
type SRV struct {
	Type     string
	Target   string
	Port     uint16
	Priority uint16
	Weight   uint16
}
type byPriority []SRV

func (o byPriority) Len() int      { return len(o) }
func (o byPriority) Swap(i, j int) { o[i], o[j] = o[j], o[i] }
func (o byPriority) Less(i, j int) bool {
	return o[i].Priority < o[j].Priority
}

func getSRV(server string, srvType string) ([]SRV, error) {
	var records []SRV
	var err error

	// Look up SRV records.
	_, addr, err := net.LookupSRV(srvType, "tcp", server)
	if err != nil {
		return records, err
	}
	if len(addr) > 0 {
		for _, adr := range addr {
			records = append(records, SRV{srvType,
				adr.Target, adr.Port, adr.Priority,
				adr.Weight})
		}
	}
	return records, err
}

// LookupXmppServer returns the xmpp-server SRV records.
func LookupXmppServer(server string) ([]SRV, error) {
	// Look up xmpp-server SRV
	records, err := getSRV(server, "xmpp-server")
	return records, err
}

// LookupXmppsServer returns the xmpps-server SRV records.
func LookupXmppsServer(server string) ([]SRV, error) {
	// Look up xmpps-server SRV
	records, err := getSRV(server, "xmpps-server")
	return records, err
}

// LookupXmppClient returns the xmpp-server SRV records.
func LookupXmppClient(server string) ([]SRV, error) {
	// Look up xmpp-client SRV
	records, err := getSRV(server, "xmpp-client")
	return records, err
}

// LookupXmppsClient returns the xmpp-server SRV records.
func LookupXmppsClient(server string) ([]SRV, error) {
	// Look up xmpps-client SRV
	records, err := getSRV(server, "xmpps-client")
	return records, err
}

// LookupClient returns the xmpp-client and xmpps-client SRV records sorted by
// priority and weight.
func LookupClient(server string) ([]SRV, error) {
	var err error
	rand.Seed(time.Now().UnixNano())
	// Look up xmpp-client SRV
	records, _ := getSRV(server, "xmpp-client")
	// Look up xmpps-client SRV records.
	records2, _ := getSRV(server, "xmpps-client")
	switch rand.Intn(2) {
	case 0:
		records = append(records, records2...)
	case 1:
		records = append(records2, records...)
	default:
		records = append(records, records2...)
	}
	switch len(records) {
	case 0:
		return records, errors.New("No client records found.")
	case 1:
		return records, err
	default:
		// Sort xmpp- and xmpps-client SRV records according to the priority
		// and weight.
		sort.Sort(byPriority(records))
	}
	return records, err
}

// LookupServer returns the xmpp-server and xmpps-server SRV records sorted by
// priority and weight.
func LookupServer(server string) ([]SRV, error) {
	var err error
	rand.Seed(time.Now().UnixNano())
	// Look up xmpp-client SRV
	records, _ := getSRV(server, "xmpp-server")
	// Look up xmpps-client SRV records.
	records2, _ := getSRV(server, "xmpps-server")
	switch rand.Intn(2) {
	case 0:
		records = append(records, records2...)
	case 1:
		records = append(records2, records...)
	default:
		records = append(records, records2...)
	}
	switch len(records) {
	case 0:
		return records, errors.New("No server records found.")
	case 1:
		return records, err
	default:
		// Sort xmpp- and xmpps-server SRV records according to the priority
		// and weight.
		sort.Sort(byPriority(records))
	}
	return records, err
}
