package checker

import (
	"encoding/xml"
	"strings"
	"testing"
)

func TestReadFeatures_WithStartTLSAndSCRAM(t *testing.T) {
	doc := `<stream:features xmlns:stream="http://etherx.jabber.org/streams">
  <starttls xmlns='urn:ietf:params:xml:ns:xmpp-tls'><required/></starttls>
  <mechanisms xmlns='urn:ietf:params:xml:ns:xmpp-sasl'>
    <mechanism>SCRAM-SHA-256</mechanism>
    <mechanism>SCRAM-SHA-1</mechanism>
    <mechanism>PLAIN</mechanism>
  </mechanisms>
</stream:features>`
	feats := decodeFeaturesForTest(t, doc)
	if feats.StartTLS == nil {
		t.Fatal("expected starttls present")
	}
	if feats.StartTLS.Required == nil {
		t.Fatal("expected starttls <required/>")
	}
	if feats.Mechanisms == nil || len(feats.Mechanisms.Mechanism) != 3 {
		t.Fatalf("expected 3 mechanisms, got %+v", feats.Mechanisms)
	}
}

func TestReadFeatures_NoSTARTTLS(t *testing.T) {
	doc := `<stream:features xmlns:stream="http://etherx.jabber.org/streams">
  <mechanisms xmlns='urn:ietf:params:xml:ns:xmpp-sasl'>
    <mechanism>PLAIN</mechanism>
  </mechanisms>
</stream:features>`
	feats := decodeFeaturesForTest(t, doc)
	if feats.StartTLS != nil {
		t.Fatal("expected no starttls")
	}
}

func TestReadFeatures_S2SDialback(t *testing.T) {
	doc := `<stream:features xmlns:stream="http://etherx.jabber.org/streams">
  <starttls xmlns='urn:ietf:params:xml:ns:xmpp-tls'><required/></starttls>
  <dialback xmlns='urn:xmpp:features:dialback'/>
</stream:features>`
	feats := decodeFeaturesForTest(t, doc)
	if feats.Dialback == nil {
		t.Fatal("expected dialback feature")
	}
}

func decodeFeaturesForTest(t *testing.T, doc string) *streamFeatures {
	t.Helper()
	dec := xml.NewDecoder(strings.NewReader(doc))
	feats, err := readFeatures(dec)
	if err != nil {
		t.Fatalf("readFeatures: %v", err)
	}
	return feats
}

func TestApplyFeatures_SASLExternal(t *testing.T) {
	ep := EndpointProbe{Mode: ModeServer}
	applyFeatures(&ep, &streamFeatures{
		Mechanisms: &mechanismsEl{Mechanism: []string{"EXTERNAL"}},
	})
	if !ep.SASLExternal {
		t.Fatal("expected SASLExternal to be set")
	}
}

func TestDeriveIssues_NoSRV(t *testing.T) {
	d := &XMPPData{
		Domain: "example.com",
		SRV:    SRVLookup{FallbackProbed: true},
	}
	is := deriveIssues(d, true, true)
	if !containsCode(is, CodeNoSRV) {
		t.Fatalf("expected %s in %+v", CodeNoSRV, is)
	}
}

func TestDeriveIssues_LegacyJabber(t *testing.T) {
	d := &XMPPData{
		SRV: SRVLookup{
			Client: []SRVRecord{{Target: "xmpp.example.com", Port: 5222}},
			Jabber: []SRVRecord{{Target: "xmpp.example.com", Port: 5222}},
		},
	}
	is := deriveIssues(d, true, false)
	if !containsCode(is, CodeLegacyJabber) {
		t.Fatalf("expected %s in %+v", CodeLegacyJabber, is)
	}
}

func TestDeriveIssues_StartTLSMissing(t *testing.T) {
	d := &XMPPData{
		SRV: SRVLookup{Client: []SRVRecord{{Target: "x", Port: 5222}}},
		Endpoints: []EndpointProbe{{
			Mode: ModeClient, Address: "1.2.3.4:5222",
			TCPConnected: true, StreamOpened: true, STARTTLSOffered: false,
		}},
	}
	is := deriveIssues(d, true, false)
	if !containsCode(is, CodeStartTLSMissing) {
		t.Fatalf("expected %s in %+v", CodeStartTLSMissing, is)
	}
}

func TestDeriveIssues_PlainOnlyAndNoSCRAM(t *testing.T) {
	d := &XMPPData{
		SRV: SRVLookup{Client: []SRVRecord{{Target: "x", Port: 5222}}},
		Endpoints: []EndpointProbe{{
			Mode: ModeClient, Address: "1.2.3.4:5222",
			TCPConnected: true, StreamOpened: true,
			STARTTLSOffered: true, STARTTLSRequired: true, STARTTLSUpgraded: true,
			FeaturesRead: true, SASLMechanisms: []string{"PLAIN"},
		}},
	}
	is := deriveIssues(d, true, false)
	if !containsCode(is, CodeSASLPlainOnly) {
		t.Fatalf("expected %s in %+v", CodeSASLPlainOnly, is)
	}
	if !containsCode(is, CodeSASLNoSCRAM) {
		t.Fatalf("expected %s in %+v", CodeSASLNoSCRAM, is)
	}
}

func TestDeriveIssues_S2SNoAuth(t *testing.T) {
	d := &XMPPData{
		SRV: SRVLookup{Server: []SRVRecord{{Target: "x", Port: 5269}}},
		Endpoints: []EndpointProbe{{
			Mode: ModeServer, Address: "1.2.3.4:5269",
			TCPConnected: true, StreamOpened: true,
			STARTTLSOffered: true, STARTTLSRequired: true, STARTTLSUpgraded: true,
			FeaturesRead:    true,
			DialbackOffered: false, SASLExternal: false,
		}},
	}
	is := deriveIssues(d, false, true)
	if !containsCode(is, CodeS2SNoAuth) {
		t.Fatalf("expected %s in %+v", CodeS2SNoAuth, is)
	}
}

func TestDeriveIssues_HappyPath(t *testing.T) {
	d := &XMPPData{
		SRV: SRVLookup{
			Client:       []SRVRecord{{Target: "x", Port: 5222}},
			Server:       []SRVRecord{{Target: "x", Port: 5269}},
			ClientSecure: []SRVRecord{{Target: "x", Port: 5223}},
		},
		Endpoints: []EndpointProbe{
			{
				Mode: ModeClient, Address: "1.2.3.4:5222",
				TCPConnected: true, StreamOpened: true,
				STARTTLSOffered: true, STARTTLSRequired: true, STARTTLSUpgraded: true,
				FeaturesRead:   true,
				SASLMechanisms: []string{"SCRAM-SHA-256", "SCRAM-SHA-256-PLUS"},
			},
			{
				Mode: ModeServer, Address: "1.2.3.4:5269",
				TCPConnected: true, StreamOpened: true,
				STARTTLSOffered: true, STARTTLSRequired: true, STARTTLSUpgraded: true,
				FeaturesRead:    true,
				DialbackOffered: true,
			},
		},
	}
	computeCoverage(d)
	is := deriveIssues(d, true, true)
	for _, i := range is {
		if i.Severity == SeverityCrit {
			t.Fatalf("unexpected crit issue %s: %s", i.Code, i.Message)
		}
	}
}

func TestComputeCoverage_Mixed(t *testing.T) {
	d := &XMPPData{
		Endpoints: []EndpointProbe{
			{Mode: ModeClient, TCPConnected: true, IsIPv6: false, STARTTLSUpgraded: true, SASLMechanisms: []string{"PLAIN"}},
			{Mode: ModeClient, TCPConnected: true, IsIPv6: true},
		},
	}
	computeCoverage(d)
	if !d.Coverage.HasIPv4 {
		t.Fatal("expected IPv4 true")
	}
	if !d.Coverage.HasIPv6 {
		t.Fatal("expected IPv6 true")
	}
	if !d.Coverage.WorkingC2S {
		t.Fatal("expected WorkingC2S")
	}
}

func TestDiscoverEndpoints_AllSets(t *testing.T) {
	d := &XMPPData{
		Domain: "example.com",
		SRV: SRVLookup{
			Client:       []SRVRecord{{Target: "xmpp.example.com", Port: 5222}},
			Server:       []SRVRecord{{Target: "xmpp.example.com", Port: 5269}},
			ClientSecure: []SRVRecord{{Target: "xmpp.example.com", Port: 5223}},
			ServerSecure: []SRVRecord{{Target: "xmpp.example.com", Port: 5270}},
			Jabber:       []SRVRecord{{Target: "xmpp.example.com", Port: 5222}}, // legacy, not emitted
		},
		Endpoints: []EndpointProbe{
			{Mode: ModeClient, Target: "xmpp.example.com", Port: 5222, STARTTLSRequired: true},
			{Mode: ModeServer, Target: "xmpp.example.com", Port: 5269, STARTTLSRequired: false},
		},
	}
	p := &xmppProvider{}
	eps, err := p.DiscoverEndpoints(d)
	if err != nil {
		t.Fatalf("DiscoverEndpoints: %v", err)
	}
	if len(eps) != 4 {
		t.Fatalf("expected 4 endpoints (legacy jabber excluded), got %d: %+v", len(eps), eps)
	}
	by := map[string]int{}
	for i, e := range eps {
		by[e.Type+":"+e.Host+":"+itoa(int(e.Port))] = i
		if e.SNI != "example.com" {
			t.Errorf("endpoint %d: SNI=%q, want example.com", i, e.SNI)
		}
	}
	c2s := eps[by["starttls-xmpp-client:xmpp.example.com:5222"]]
	if c2s.Meta["starttls"] != "required" {
		t.Errorf("c2s starttls meta = %v, want required", c2s.Meta)
	}
	s2s := eps[by["starttls-xmpp-server:xmpp.example.com:5269"]]
	if s2s.Meta["starttls"] != "opportunistic" {
		t.Errorf("s2s starttls meta = %v, want opportunistic", s2s.Meta)
	}
	direct := eps[by["tls:xmpp.example.com:5223"]]
	if direct.Meta != nil {
		t.Errorf("direct-TLS endpoint should not carry starttls meta, got %v", direct.Meta)
	}
}

func TestDiscoverEndpoints_WrongType(t *testing.T) {
	p := &xmppProvider{}
	eps, err := p.DiscoverEndpoints("not an XMPPData")
	if err != nil {
		t.Fatalf("expected nil error for wrong type, got %v", err)
	}
	if eps != nil {
		t.Fatalf("expected nil endpoints for wrong type, got %v", eps)
	}
}

func itoa(i int) string {
	if i == 0 {
		return "0"
	}
	var buf [6]byte
	pos := len(buf)
	for i > 0 {
		pos--
		buf[pos] = byte('0' + i%10)
		i /= 10
	}
	return string(buf[pos:])
}

func containsCode(is []Issue, code string) bool {
	for _, i := range is {
		if i.Code == code {
			return true
		}
	}
	return false
}