package xmppsig

import (
	"bytes"
	"context"
	"encoding/xml"
	"fmt"
	"strconv"
	"time"

	"github.com/google/uuid"
	"github.com/pion/webrtc/v4"
	"gosrc.io/xmpp"
	"gosrc.io/xmpp/stanza"
)

// XEP-0215 v2; older v1 is dead on modern prosody/ejabberd
const NSExtDisco = "urn:xmpp:extdisco:2"

// one resolved STUN/TURN/TURNS entry; for TURN the user/pass are
// short-lived HMAC creds the server issues per-query
type Service struct {
	Type      string // "stun" | "turn" | "turns" | "stuns"
	Host      string
	Port      int
	Transport string // "udp" | "tcp"
	Username  string
	Password  string
}

// XEP-0215 wire encoding. Port is decoded as string because some servers
// have emitted non-numeric values; we surface the parse error.
type servicesQuery struct {
	XMLName  xml.Name      `xml:"urn:xmpp:extdisco:2 services"`
	Services []serviceElem `xml:"service"`
}

type serviceElem struct {
	XMLName   xml.Name `xml:"service"`
	Type      string   `xml:"type,attr"`
	Host      string   `xml:"host,attr"`
	Port      string   `xml:"port,attr"`
	Transport string   `xml:"transport,attr"`
	Username  string   `xml:"username,attr,omitempty"`
	Password  string   `xml:"password,attr,omitempty"`
}

// gosrc.io/xmpp IQ payload interface
func (servicesQuery) Namespace() string         { return NSExtDisco }
func (servicesQuery) Name() string              { return "services" }
func (servicesQuery) GetSet() *stanza.ResultSet { return nil }

// XEP-0215 services request; toJID is the c2s/parent server domain.
// Empty slice is a legitimate result.
func QueryServices(sender xmpp.Sender, fromJID, toJID string, timeout time.Duration) ([]Service, error) {
	iq, err := stanza.NewIQ(stanza.Attrs{
		Type: stanza.IQTypeGet,
		From: fromJID,
		To:   toJID,
		Id:   "extdisco-" + uuid.New().String(),
	})
	if err != nil {
		return nil, fmt.Errorf("build IQ: %w", err)
	}
	iq.Payload = &servicesQuery{}

	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	ch, err := sender.SendIQ(ctx, iq)
	if err != nil {
		return nil, fmt.Errorf("send IQ: %w", err)
	}

	var resp stanza.IQ
	select {
	case <-ctx.Done():
		return nil, fmt.Errorf("extdisco timeout after %s", timeout)
	case resp = <-ch:
	}

	if resp.Type == stanza.IQTypeError {
		// dump the marshalled IQ so the <error> condition is visible
		raw, _ := xml.Marshal(resp)
		return nil, fmt.Errorf("extdisco error from %s: %s", toJID, string(raw))
	}

	// gosrc.io decodes unrecognised payloads opaquely - re-marshal and
	// parse <services> ourselves rather than fight the stanza registry
	raw, err := xml.Marshal(resp)
	if err != nil {
		return nil, fmt.Errorf("re-marshal response: %w", err)
	}
	return extractServices(raw)
}

// unit-testable without gosrc.io; walks raw IQ XML
func extractServices(raw []byte) ([]Service, error) {
	dec := xml.NewDecoder(bytes.NewReader(raw))
	for {
		tok, err := dec.Token()
		if err != nil {
			return nil, fmt.Errorf("no <services> element in IQ")
		}
		se, ok := tok.(xml.StartElement)
		if !ok {
			continue
		}
		if se.Name.Space != NSExtDisco || se.Name.Local != "services" {
			continue
		}
		var q servicesQuery
		if err := dec.DecodeElement(&q, &se); err != nil {
			return nil, fmt.Errorf("decode services: %w", err)
		}
		out := make([]Service, 0, len(q.Services))
		for _, s := range q.Services {
			port, err := strconv.Atoi(s.Port)
			if err != nil {
				return nil, fmt.Errorf("service %s: bad port %q: %w", s.Type, s.Port, err)
			}
			out = append(out, Service{
				Type:      s.Type,
				Host:      s.Host,
				Port:      port,
				Transport: s.Transport,
				Username:  s.Username,
				Password:  s.Password,
			})
		}
		return out, nil
	}
}

// XEP-0215 services -> pion webrtc.ICEServer list
func ToICEServers(services []Service) []webrtc.ICEServer {
	out := make([]webrtc.ICEServer, 0, len(services))
	for _, s := range services {
		switch s.Type {
		case "stun", "stuns":
			out = append(out, webrtc.ICEServer{
				URLs: []string{fmt.Sprintf("%s:%s:%d", s.Type, s.Host, s.Port)},
			})
		case "turn", "turns":
			// skip turn entries without creds - pion rejects the whole PC
			// (InvalidAccessError) on any one missing-creds entry. Common
			// cause: prosody's mod_external_services only auto-maps
			// algorithms["turn"], so a `turns` entry needs an explicit
			// algorithm = "turn" to get creds generated.
			if s.Username == "" || s.Password == "" {
				continue
			}
			url := fmt.Sprintf("%s:%s:%d", s.Type, s.Host, s.Port)
			if s.Transport != "" {
				url += "?transport=" + s.Transport
			}
			out = append(out, webrtc.ICEServer{
				URLs:           []string{url},
				Username:       s.Username,
				Credential:     s.Password,
				CredentialType: webrtc.ICECredentialTypePassword,
			})
		default:
			// skip unknown service types (e.g. ftp from older XEPs)
		}
	}
	return out
}