package checker

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"net"
	"slices"
	"strconv"
	"strings"
	"sync"
	"time"

	ldapv3 "github.com/go-ldap/ldap/v3"
	"github.com/miekg/dns"

	sdk "git.happydns.org/checker-sdk-go/checker"
)

// tlsProbeConfig builds a permissive *tls.Config for probing: hostname
// verification is skipped because cert validation is the TLS checker's
// job. We only care that a TLS session can be established at all.
func tlsProbeConfig(serverName string) *tls.Config {
	return &tls.Config{
		ServerName:         serverName,
		InsecureSkipVerify: true, //nolint:gosec -- cert validation is the TLS checker's job
		MinVersion:         tls.VersionTLS10,
	}
}

// Collect runs the full LDAP probe for a domain.
func (p *ldapProvider) Collect(ctx context.Context, opts sdk.CheckerOptions) (any, error) {
	domain, _ := sdk.GetOption[string](opts, "domain")
	domain = strings.TrimSuffix(domain, ".")
	if domain == "" {
		return nil, fmt.Errorf("domain is required")
	}

	bindDN, _ := sdk.GetOption[string](opts, "bind_dn")
	bindPassword, _ := sdk.GetOption[string](opts, "bind_password")
	baseDN, _ := sdk.GetOption[string](opts, "base_dn")

	timeoutSecs := sdk.GetFloatOption(opts, "timeout", 10)
	if timeoutSecs < 1 {
		timeoutSecs = 10
	}
	perEndpoint := time.Duration(timeoutSecs * float64(time.Second))

	data := &LDAPData{
		Domain: domain,
		BaseDN: baseDN,
		RunAt:  time.Now().UTC().Format(time.RFC3339),
		SRV:    SRVLookup{Errors: map[string]string{}},
	}

	resolver := net.DefaultResolver

	var srvWG sync.WaitGroup
	var srvErrMu sync.Mutex
	srvWG.Add(2)
	lookup := func(prefix string, dst *[]SRVRecord) {
		defer srvWG.Done()
		records, err := lookupSRV(ctx, resolver, prefix, domain)
		if err != nil {
			srvErrMu.Lock()
			data.SRV.Errors[prefix] = err.Error()
			srvErrMu.Unlock()
			return
		}
		*dst = records
	}
	go lookup("_ldap._tcp.", &data.SRV.LDAP)
	go lookup("_ldaps._tcp.", &data.SRV.LDAPS)
	srvWG.Wait()

	totalSRV := len(data.SRV.LDAP) + len(data.SRV.LDAPS)
	if totalSRV == 0 {
		data.SRV.FallbackProbed = true
		data.SRV.LDAP = []SRVRecord{{Target: domain, Port: 389}}
		data.SRV.LDAPS = []SRVRecord{{Target: domain, Port: 636}}
	}

	resolveAll(ctx, resolver, data.SRV.LDAP, data.SRV.LDAPS)

	data.BindTested = bindDN != "" && bindPassword != ""

	var plainEndpoints, ldapsEndpoints []EndpointProbe
	var probeWG sync.WaitGroup
	probeWG.Add(2)
	go func() {
		defer probeWG.Done()
		plainEndpoints = probeSet(ctx, domain, ModePlain, "_ldap._tcp", data.SRV.LDAP, perEndpoint, bindDN, bindPassword, baseDN)
	}()
	go func() {
		defer probeWG.Done()
		ldapsEndpoints = probeSet(ctx, domain, ModeLDAPS, "_ldaps._tcp", data.SRV.LDAPS, perEndpoint, bindDN, bindPassword, baseDN)
	}()
	probeWG.Wait()
	data.Endpoints = append(plainEndpoints, ldapsEndpoints...)

	return data, nil
}

func probeSet(ctx context.Context, domain string, mode LDAPMode, prefix string, records []SRVRecord, timeout time.Duration, bindDN, bindPassword, baseDN string) []EndpointProbe {
	type task struct {
		rec  SRVRecord
		addr *probeAddr // nil means the SRV target has no A/AAAA records
	}
	var tasks []task
	for _, rec := range records {
		addrs := addressesForProbe(rec)
		if len(addrs) == 0 {
			tasks = append(tasks, task{rec: rec})
			continue
		}
		for _, a := range addrs {
			tasks = append(tasks, task{rec: rec, addr: &a})
		}
	}

	results := make([]EndpointProbe, len(tasks))
	var wg sync.WaitGroup
	for i, t := range tasks {
		wg.Add(1)
		go func(i int, t task) {
			defer wg.Done()
			if t.addr == nil {
				results[i] = EndpointProbe{
					Mode:      mode,
					SRVPrefix: prefix,
					Target:    t.rec.Target,
					Port:      t.rec.Port,
					Error:     "no A/AAAA records for target",
				}
				return
			}
			results[i] = probeEndpoint(ctx, domain, mode, prefix, t.rec, t.addr.ip, t.addr.isV6, timeout, bindDN, bindPassword, baseDN)
		}(i, t)
	}
	wg.Wait()
	return results
}

type probeAddr struct {
	ip   string
	isV6 bool
}

func addressesForProbe(rec SRVRecord) []probeAddr {
	var out []probeAddr
	for _, ip := range rec.IPv4 {
		out = append(out, probeAddr{ip: ip, isV6: false})
	}
	for _, ip := range rec.IPv6 {
		out = append(out, probeAddr{ip: ip, isV6: true})
	}
	return out
}

// probeEndpoint runs the full probe on a single (host, ip, port) tuple:
// TCP connect → optional StartTLS or direct-TLS handshake → RootDSE read →
// plaintext-bind posture check (only on LDAP:389 before TLS) → optional
// authenticated bind + base-DN read.
func probeEndpoint(ctx context.Context, domain string, mode LDAPMode, prefix string, rec SRVRecord, ip string, isV6 bool, timeout time.Duration, bindDN, bindPassword, baseDN string) EndpointProbe {
	start := time.Now()
	result := EndpointProbe{
		Mode:      mode,
		SRVPrefix: prefix,
		Target:    rec.Target,
		Port:      rec.Port,
		Address:   net.JoinHostPort(ip, strconv.Itoa(int(rec.Port))),
		IsIPv6:    isV6,
	}
	defer func() { result.ElapsedMS = time.Since(start).Milliseconds() }()

	dialCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	dialer := &net.Dialer{}
	rawConn, err := dialer.DialContext(dialCtx, "tcp", result.Address)
	if err != nil {
		result.Error = "tcp: " + err.Error()
		return result
	}
	result.TCPConnected = true
	defer rawConn.Close()
	// renewDeadline keeps the underlying TCP deadline rolling per major
	// step so a slow TLS handshake doesn't starve the RootDSE / bind /
	// base-read calls that follow.
	renewDeadline := func() { _ = rawConn.SetDeadline(time.Now().Add(timeout)) }
	renewDeadline()

	// For the plaintext-bind posture check on Mode=ldap, we first spin up
	// a separate short-lived connection before upgrading this one to TLS.
	// A single raw connection can't both "test cleartext bind refusal" and
	// "then StartTLS" cleanly -- once we've bound we'd skew RootDSE results.
	if mode == ModePlain {
		result.PlaintextBindTested, result.PlaintextBindAccepted = probePlaintextBindRefusal(result.Address, timeout)
	}

	// Establish the LDAP session we'll use for the rest of the probe.
	var conn *ldapv3.Conn
	if mode == ModeLDAPS {
		tlsConn := tls.Client(rawConn, tlsProbeConfig(domain))
		if err := tlsConn.Handshake(); err != nil {
			result.Error = "tls-handshake: " + err.Error()
			return result
		}
		result.TLSEstablished = true
		state := tlsConn.ConnectionState()
		result.TLSVersion = tls.VersionName(state.Version)
		result.TLSCipher = tls.CipherSuiteName(state.CipherSuite)
		renewDeadline()
		conn = ldapv3.NewConn(tlsConn, true)
	} else {
		conn = ldapv3.NewConn(rawConn, false)
	}
	conn.Start()
	defer conn.Close()
	conn.SetTimeout(timeout)

	// Try RootDSE over the native transport first -- works on LDAPS straight
	// away, and on LDAP it reveals the supported extensions including
	// StartTLS capability before we attempt the upgrade.
	renewDeadline()
	readRootDSE(conn, &result)

	if mode == ModePlain {
		// Detect StartTLS advertisement in supportedExtension. RFC 4511
		// says the RootDSE MAY advertise "1.3.6.1.4.1.1466.20037".
		offersStartTLS := stringListContainsFold(result.SupportedExtension, "1.3.6.1.4.1.1466.20037")
		if offersStartTLS {
			result.StartTLSOffered = true
			if err := conn.StartTLS(tlsProbeConfig(domain)); err != nil {
				result.Error = "starttls: " + err.Error()
			} else {
				result.StartTLSUpgraded = true
				result.TLSEstablished = true
				renewDeadline()
				// go-ldap doesn't expose the *tls.ConnectionState directly.
				// Fall back to inspecting the underlying conn via TLSConnectionState.
				if state, ok := conn.TLSConnectionState(); ok {
					result.TLSVersion = tls.VersionName(state.Version)
					result.TLSCipher = tls.CipherSuiteName(state.CipherSuite)
				}
				// Post-TLS SASL refresh: some servers only publish the
				// strong mechanisms once the channel is encrypted. We only
				// need the mechanisms list -- naming contexts, LDAP
				// version and vendor strings don't change.
				refreshSASLMechanisms(conn, &result)
			}
		} else if !result.RootDSERead {
			result.Error = "rootdse-unreadable: RootDSE could not be read"
		}
	}

	// Anonymous bind + search -- we try unconditionally so we can flag
	// exposure.
	renewDeadline()
	anonBindOK := conn.Bind("", "") == nil
	result.AnonymousBindAllowed = anonBindOK
	if anonBindOK && len(result.NamingContexts) > 0 {
		// baseObject search returns 0 or 1 entries -- we only want to
		// detect whether an anonymous query can peek at DIT contents.
		sr, err := conn.Search(ldapv3.NewSearchRequest(
			result.NamingContexts[0],
			ldapv3.ScopeBaseObject,
			ldapv3.NeverDerefAliases,
			1, int(timeout.Seconds()), false,
			"(objectClass=*)",
			[]string{"1.1"}, // request no attributes
			nil,
		))
		if err == nil && sr != nil && len(sr.Entries) > 0 {
			result.AnonymousSearchAllowed = true
		}
	}

	// Authenticated bind + base DN read (only when caller provided creds
	// AND we are on an encrypted channel -- never ship a password over
	// cleartext).
	if bindDN != "" && bindPassword != "" && result.TLSEstablished {
		result.BindAttempted = true
		renewDeadline()
		err := conn.Bind(bindDN, bindPassword)
		if err == nil {
			result.BindOK = true
			if baseDN != "" {
				result.BaseReadAttempted = true
				renewDeadline()
				sr, err := conn.Search(ldapv3.NewSearchRequest(
					baseDN,
					ldapv3.ScopeBaseObject,
					ldapv3.NeverDerefAliases,
					1, int(timeout.Seconds()), false,
					"(objectClass=*)",
					[]string{"1.1"},
					nil,
				))
				if err != nil {
					result.BaseReadError = err.Error()
				} else if sr != nil {
					result.BaseReadOK = true
					result.BaseReadEntries = len(sr.Entries)
				}
			}
		} else {
			result.BindError = err.Error()
		}
	}

	return result
}

// probePlaintextBindRefusal opens a short-lived, fresh TCP connection and
// attempts a simple bind with a fixed, syntactically-safe DN over cleartext.
// We are not probing credentials -- we want to learn whether the server
// refuses authentication on an unprotected link (RFC 4513 §5.1.2 calls this
// "confidentialityRequired" / resultCode 13). Any response other than
// resultCode 13 means the server will accept cleartext bind attempts.
func probePlaintextBindRefusal(address string, timeout time.Duration) (tested, accepted bool) {
	dialer := &net.Dialer{Timeout: timeout}
	raw, err := dialer.Dial("tcp", address)
	if err != nil {
		return false, false
	}
	defer raw.Close()
	_ = raw.SetDeadline(time.Now().Add(timeout))

	conn := ldapv3.NewConn(raw, false)
	conn.Start()
	defer conn.Close()
	conn.SetTimeout(timeout)

	tested = true
	// Fixed probe DN -- caller-supplied domain is not interpolated to avoid
	// LDAP DN injection. The server is expected to reject this DN regardless
	// of value; we only care whether it returns confidentialityRequired (13)
	// or attempts the bind anyway.
	err = conn.Bind("cn=checker-probe", "x-not-a-real-password-x")
	if err == nil {
		return tested, true
	}
	// resultCode 13 (confidentialityRequired) is the only response that
	// means the server actively refused to authenticate over cleartext.
	// Anything else (49 invalidCredentials, 32 noSuchObject, …) means the
	// server was willing to attempt the bind, which is the insecure
	// posture we want to flag.
	var lerr *ldapv3.Error
	if errors.As(err, &lerr) && lerr.ResultCode == ldapv3.LDAPResultConfidentialityRequired {
		return tested, false
	}
	return tested, true
}

// readRootDSE performs a single RootDSE lookup and fills the matching
// fields on ep. Failures are not fatal -- many hardened servers refuse
// anonymous RootDSE reads; we just note that we couldn't read it.
func readRootDSE(conn *ldapv3.Conn, ep *EndpointProbe) {
	sr, err := conn.Search(ldapv3.NewSearchRequest(
		"",
		ldapv3.ScopeBaseObject,
		ldapv3.NeverDerefAliases,
		1, 5, false,
		"(objectClass=*)",
		[]string{
			"supportedLDAPVersion",
			"supportedSASLMechanisms",
			"supportedControl",
			"supportedExtension",
			"namingContexts",
			"vendorName",
			"vendorVersion",
		},
		nil,
	))
	if err != nil || sr == nil || len(sr.Entries) == 0 {
		return
	}
	ep.RootDSERead = true
	e := sr.Entries[0]
	ep.SupportedLDAPVersion = unique(append(ep.SupportedLDAPVersion, e.GetAttributeValues("supportedLDAPVersion")...))
	ep.SupportedSASLMechanisms = unique(append(ep.SupportedSASLMechanisms, e.GetAttributeValues("supportedSASLMechanisms")...))
	ep.SupportedControl = unique(append(ep.SupportedControl, e.GetAttributeValues("supportedControl")...))
	ep.SupportedExtension = unique(append(ep.SupportedExtension, e.GetAttributeValues("supportedExtension")...))
	ep.NamingContexts = unique(append(ep.NamingContexts, e.GetAttributeValues("namingContexts")...))
	if v := e.GetAttributeValue("vendorName"); v != "" {
		ep.VendorName = v
	}
	if v := e.GetAttributeValue("vendorVersion"); v != "" {
		ep.VendorVersion = v
	}
}

// refreshSASLMechanisms re-queries the RootDSE after StartTLS to pick up any
// SASL mechanisms the server only advertises over an encrypted channel.
func refreshSASLMechanisms(conn *ldapv3.Conn, ep *EndpointProbe) {
	sr, err := conn.Search(ldapv3.NewSearchRequest(
		"",
		ldapv3.ScopeBaseObject,
		ldapv3.NeverDerefAliases,
		1, 5, false,
		"(objectClass=*)",
		[]string{"supportedSASLMechanisms"},
		nil,
	))
	if err != nil || sr == nil || len(sr.Entries) == 0 {
		return
	}
	ep.SupportedSASLMechanisms = unique(append(ep.SupportedSASLMechanisms, sr.Entries[0].GetAttributeValues("supportedSASLMechanisms")...))
}

func lookupSRV(ctx context.Context, r *net.Resolver, prefix, domain string) ([]SRVRecord, error) {
	name := prefix + dns.Fqdn(domain)
	_, records, err := r.LookupSRV(ctx, "", "", name)
	if err != nil {
		var dnsErr *net.DNSError
		if errors.As(err, &dnsErr) && dnsErr.IsNotFound {
			return nil, nil
		}
		return nil, err
	}
	// RFC 2782: single record "." with port 0 means "service explicitly not
	// available at this domain". Treat that as "no records" for probing.
	if len(records) == 1 && (records[0].Target == "." || records[0].Target == "") && records[0].Port == 0 {
		return nil, nil
	}
	out := make([]SRVRecord, 0, len(records))
	for _, r := range records {
		out = append(out, SRVRecord{
			Target:   strings.TrimSuffix(r.Target, "."),
			Port:     r.Port,
			Priority: r.Priority,
			Weight:   r.Weight,
		})
	}
	return out, nil
}

// resolveAll resolves A/AAAA for every record across all sets concurrently.
// Each goroutine writes only to its own record, so no lock is needed.
func resolveAll(ctx context.Context, r *net.Resolver, sets ...[]SRVRecord) {
	var wg sync.WaitGroup
	for _, records := range sets {
		for i := range records {
			wg.Add(1)
			go func(rec *SRVRecord) {
				defer wg.Done()
				ips, err := r.LookupIPAddr(ctx, rec.Target)
				if err != nil {
					return
				}
				for _, ip := range ips {
					if v4 := ip.IP.To4(); v4 != nil {
						rec.IPv4 = append(rec.IPv4, v4.String())
					} else {
						rec.IPv6 = append(rec.IPv6, ip.IP.String())
					}
				}
			}(&records[i])
		}
	}
	wg.Wait()
}

func stringListContainsFold(list []string, want string) bool {
	return slices.ContainsFunc(list, func(s string) bool { return strings.EqualFold(s, want) })
}

func unique(list []string) []string {
	if len(list) <= 1 {
		return list
	}
	seen := make(map[string]struct{}, len(list))
	out := make([]string, 0, len(list))
	for _, s := range list {
		if _, ok := seen[s]; ok {
			continue
		}
		seen[s] = struct{}{}
		out = append(out, s)
	}
	return out
}