package checker

import (
	"context"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"net"
	"strconv"
	"strings"
	"sync"
	"time"

	asn1 "github.com/jcmturner/gofork/encoding/asn1"
	"github.com/jcmturner/gokrb5/v8/config"
	"github.com/jcmturner/gokrb5/v8/iana/errorcode"
	"github.com/jcmturner/gokrb5/v8/iana/etypeID"
	"github.com/jcmturner/gokrb5/v8/iana/nametype"
	"github.com/jcmturner/gokrb5/v8/iana/patype"
	"github.com/jcmturner/gokrb5/v8/messages"
	"github.com/jcmturner/gokrb5/v8/types"

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

// srvPrefixes lists every SRV label we probe, in order. The first two are
// considered mandatory (no realm is usable without one of them).
var srvPrefixes = []struct {
	prefix string
	role   string
	proto  string
}{
	{"_kerberos._tcp.", "kdc", "tcp"},
	{"_kerberos._udp.", "kdc", "udp"},
	{"_kerberos-master._tcp.", "master", "tcp"},
	{"_kerberos-adm._tcp.", "kadmin", "tcp"},
	{"_kpasswd._tcp.", "kpasswd", "tcp"},
	{"_kpasswd._udp.", "kpasswd", "udp"},
}

// weakEnctypes lists the IDs considered weak per RFC 8429 + common guidance.
var weakEnctypes = map[int32]bool{
	etypeID.DES_CBC_CRC:   true,
	etypeID.DES_CBC_MD4:   true,
	etypeID.DES_CBC_MD5:   true,
	etypeID.DES_CBC_RAW:   true,
	etypeID.DES3_CBC_MD5:  true,
	etypeID.DES3_CBC_RAW:  true,
	etypeID.DES3_CBC_SHA1: true,
	etypeID.DES_HMAC_SHA1: true,
	etypeID.RC4_HMAC:      true,
	etypeID.RC4_HMAC_EXP:  true,
}

// preferredEnctypes is the enctype preference list used in anonymous probes
// and authenticated requests. RC4 is included last so the KDC's ETYPE-INFO2
// reveals whether it still supports it.
var preferredEnctypes = []int32{
	etypeID.AES256_CTS_HMAC_SHA1_96,
	etypeID.AES128_CTS_HMAC_SHA1_96,
	etypeID.AES256_CTS_HMAC_SHA384_192,
	etypeID.AES128_CTS_HMAC_SHA256_128,
	etypeID.RC4_HMAC,
}

// etypeName returns a human-friendly name for an enctype ID, falling back
// to its numeric value when unknown.
func etypeName(id int32) string {
	for name, nid := range etypeID.ETypesByName {
		if nid == id {
			// Prefer canonical "aes..." / "rc4-hmac" shape
			if !strings.Contains(name, "-CmsOID") && !strings.HasSuffix(name, "-EnvOID") {
				return name
			}
		}
	}
	return fmt.Sprintf("etype-%d", id)
}

func (p *kerberosProvider) Collect(ctx context.Context, opts sdk.CheckerOptions) (any, error) {
	realmRaw, _ := opts["realm"].(string)
	realmRaw = strings.TrimSuffix(realmRaw, ".")
	if realmRaw == "" {
		return nil, fmt.Errorf("realm is required")
	}
	realm := strings.ToUpper(realmRaw)
	domain := strings.ToLower(realmRaw)

	timeoutSec := optFloat(opts, "timeout", 5)
	if timeoutSec <= 0 {
		timeoutSec = 5
	}
	timeout := time.Duration(timeoutSec * float64(time.Second))

	now := time.Now().UTC()
	data := &KerberosData{
		Realm:       realm,
		CollectedAt: now,
		LocalTime:   now,
		Resolution:  map[string]HostResolution{},
	}

	// 1. SRV discovery ---------------------------------------------------------
	resolver := &net.Resolver{}
	data.SRV = make([]SRVBucket, len(srvPrefixes))
	var wg sync.WaitGroup
	for i, sp := range srvPrefixes {
		wg.Go(func() {
			bucket := SRVBucket{Prefix: sp.prefix, LookupName: sp.prefix + domain + "."}
			lctx, cancel := context.WithTimeout(ctx, timeout)
			_, addrs, err := resolver.LookupSRV(lctx, "", "", sp.prefix+domain)
			cancel()
			if err != nil {
				var dnsErr *net.DNSError
				if errors.As(err, &dnsErr) && dnsErr.IsNotFound {
					bucket.NXDomain = true
				} else {
					bucket.Error = err.Error()
				}
			}
			for _, srv := range addrs {
				target := strings.TrimSuffix(srv.Target, ".")
				bucket.Records = append(bucket.Records, SRVRecord{
					Target:   target,
					Port:     srv.Port,
					Priority: srv.Priority,
					Weight:   srv.Weight,
				})
			}
			data.SRV[i] = bucket
		})
	}
	wg.Wait()

	// 2. Forward resolution ----------------------------------------------------
	var targets []string
	seen := map[string]bool{}
	for _, b := range data.SRV {
		for _, r := range b.Records {
			if !seen[r.Target] {
				seen[r.Target] = true
				targets = append(targets, r.Target)
			}
		}
	}
	resolutions := make([]HostResolution, len(targets))
	for i, target := range targets {
		wg.Go(func() {
			hr := HostResolution{Target: target}
			lctx, cancel := context.WithTimeout(ctx, timeout)
			ips, err := resolver.LookupIPAddr(lctx, target)
			cancel()
			if err != nil {
				hr.Error = err.Error()
			}
			for _, ip := range ips {
				if ip.IP.To4() != nil {
					hr.IPv4 = append(hr.IPv4, ip.IP.String())
				} else {
					hr.IPv6 = append(hr.IPv6, ip.IP.String())
				}
			}
			resolutions[i] = hr
		})
	}
	wg.Wait()
	for _, hr := range resolutions {
		data.Resolution[hr.Target] = hr
	}

	// 3. L4 probes -------------------------------------------------------------
	kdcHosts := map[string]uint16{} // target -> tcp port, used to pick the AS-REQ destination
	for _, b := range data.SRV {
		role := roleForPrefix(b.Prefix)
		if strings.HasSuffix(b.Prefix, "._tcp.") {
			for _, r := range b.Records {
				probe := dialTCP(ctx, r.Target, r.Port, role, timeout)
				data.Probes = append(data.Probes, probe)
				if role == "kdc" && probe.OK {
					if _, ok := kdcHosts[r.Target]; !ok {
						kdcHosts[r.Target] = r.Port
					}
				}
			}
		}
	}

	// 4. AS-REQ probe (tries each reachable KDC TCP first, then UDP fallbacks).
	asProbe := ASProbeResult{Attempted: true}
	asReq, err := buildProbeASReq(realm)
	if err != nil {
		asProbe.Error = fmt.Sprintf("build AS-REQ: %v", err)
	} else {
		req, _ := asReq.Marshal()

		// Target list: reachable TCP KDCs, then UDP KDCs if nothing else.
		var tcpEps, udpEps []endpoint
		for _, b := range data.SRV {
			if !strings.HasPrefix(b.Prefix, "_kerberos.") {
				continue
			}
			proto := "tcp"
			if strings.HasSuffix(b.Prefix, "._udp.") {
				proto = "udp"
			}
			for _, r := range b.Records {
				e := endpoint{r.Target, r.Port, proto}
				if proto == "tcp" {
					tcpEps = append(tcpEps, e)
				} else {
					udpEps = append(udpEps, e)
				}
			}
		}
		eps := append(tcpEps[:len(tcpEps):len(tcpEps)], udpEps...)

		for _, e := range eps {
			start := time.Now()
			var reply []byte
			var perr error
			if e.proto == "tcp" {
				reply, perr = sendASReqTCP(ctx, e.target, e.port, req, timeout)
			} else {
				reply, perr = sendASReqUDP(ctx, e.target, e.port, req, timeout)
				// Track UDP reachability via this attempt.
				probe := KDCProbe{
					Target: e.target, Port: e.port, Proto: "udp", Role: "kdc",
					RTT: time.Since(start),
				}
				if perr == nil {
					probe.OK = true
					probe.KrbSeen = true
				} else {
					probe.Error = perr.Error()
				}
				data.Probes = append(data.Probes, probe)
			}
			if perr != nil || len(reply) == 0 {
				continue
			}
			asProbe.Target = e.target
			asProbe.Proto = e.proto
			parseASResponse(reply, &asProbe)
			break
		}
		if asProbe.Target == "" && asProbe.Error == "" {
			asProbe.Error = "no KDC answered our AS-REQ probe"
		}

		if !asProbe.ServerTime.IsZero() {
			asProbe.ClockSkew = time.Since(asProbe.ServerTime)
		}
	}
	data.AS = asProbe

	// 5. Roll up enctypes ------------------------------------------------------
	for _, e := range asProbe.Enctypes {
		if e.Weak {
			data.WeakEnctypes = append(data.WeakEnctypes, e)
		}
	}
	data.Enctypes = asProbe.Enctypes

	// 6. Optional authenticated round-trip ------------------------------------
	principal, _ := opts["principal"].(string)
	password, _ := opts["password"].(string)
	if principal != "" && password != "" {
		data.Auth = runAuthProbe(ctx, realm, principal, password,
			stringOpt(opts, "targetService"), kdcHosts, timeout)
	}

	return data, nil
}

func roleForPrefix(prefix string) string {
	switch {
	case strings.HasPrefix(prefix, "_kerberos-master."):
		return "master"
	case strings.HasPrefix(prefix, "_kerberos-adm."):
		return "kadmin"
	case strings.HasPrefix(prefix, "_kpasswd."):
		return "kpasswd"
	default:
		return "kdc"
	}
}

type endpoint struct {
	target string
	port   uint16
	proto  string
}

func dialTCP(ctx context.Context, target string, port uint16, role string, timeout time.Duration) KDCProbe {
	probe := KDCProbe{Target: target, Port: port, Proto: "tcp", Role: role}
	addr := net.JoinHostPort(target, strconv.Itoa(int(port)))
	start := time.Now()
	d := net.Dialer{Timeout: timeout}
	conn, err := d.DialContext(ctx, "tcp", addr)
	probe.RTT = time.Since(start)
	if err != nil {
		probe.Error = err.Error()
		return probe
	}
	_ = conn.Close()
	probe.OK = true
	return probe
}

// sendASReqTCP frames the AS-REQ with its 4-byte length prefix (RFC 4120 §7.2.2)
// and reads the response the same way.
func sendASReqTCP(ctx context.Context, target string, port uint16, body []byte, timeout time.Duration) ([]byte, error) {
	addr := net.JoinHostPort(target, strconv.Itoa(int(port)))
	d := net.Dialer{Timeout: timeout}
	conn, err := d.DialContext(ctx, "tcp", addr)
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	_ = conn.SetDeadline(time.Now().Add(timeout))

	hdr := make([]byte, 4)
	binary.BigEndian.PutUint32(hdr, uint32(len(body)))
	if _, err := conn.Write(append(hdr, body...)); err != nil {
		return nil, err
	}

	lenBuf := make([]byte, 4)
	if _, err := io.ReadFull(conn, lenBuf); err != nil {
		return nil, err
	}
	n := binary.BigEndian.Uint32(lenBuf)
	if n == 0 || n > 1<<20 {
		return nil, fmt.Errorf("suspicious reply length %d", n)
	}
	buf := make([]byte, n)
	if _, err := io.ReadFull(conn, buf); err != nil {
		return nil, err
	}
	return buf, nil
}

func sendASReqUDP(ctx context.Context, target string, port uint16, body []byte, timeout time.Duration) ([]byte, error) {
	addr := net.JoinHostPort(target, strconv.Itoa(int(port)))
	d := net.Dialer{Timeout: timeout}
	conn, err := d.DialContext(ctx, "udp", addr)
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	_ = conn.SetDeadline(time.Now().Add(timeout))
	if _, err := conn.Write(body); err != nil {
		return nil, err
	}
	buf := make([]byte, 4096)
	n, err := conn.Read(buf)
	if err != nil {
		return nil, err
	}
	return buf[:n], nil
}

// buildProbeASReq builds an AS-REQ for krbtgt/REALM@REALM with a throwaway
// cname. We don't want to depend on reading a local krb5.conf, so we craft
// the request fields by hand from a default config.
func buildProbeASReq(realm string) (messages.ASReq, error) {
	cfg := config.New()
	cfg.LibDefaults.DefaultRealm = realm
	cfg.LibDefaults.NoAddresses = true
	cfg.LibDefaults.TicketLifetime = 10 * time.Minute
	cfg.LibDefaults.DefaultTktEnctypeIDs = preferredEnctypes
	cname := types.NewPrincipalName(nametype.KRB_NT_PRINCIPAL, "probe-happydomain")
	return messages.NewASReqForTGT(realm, cfg, cname)
}

// parseASResponse inspects the raw KDC reply and fills the ASProbeResult.
// Expected replies: KRB-ERROR (PREAUTH_REQUIRED / C_PRINCIPAL_UNKNOWN) or,
// less commonly, an AS-REP (principal exists and doesn't require preauth —
// AS-REP roasting territory).
func parseASResponse(raw []byte, out *ASProbeResult) {
	// Try KRB-ERROR first.
	var krbErr messages.KRBError
	if err := krbErr.Unmarshal(raw); err == nil {
		out.ErrorCode = krbErr.ErrorCode
		out.ErrorName = errorcodeName(krbErr.ErrorCode)
		out.ServerRealm = krbErr.Realm
		if !krbErr.STime.IsZero() {
			out.ServerTime = krbErr.STime.UTC()
		}
		out.PreauthReq = krbErr.ErrorCode == errorcode.KDC_ERR_PREAUTH_REQUIRED
		out.Raw = fmt.Sprintf("KRB-ERROR %s", out.ErrorName)

		if len(krbErr.EData) > 0 {
			enctypes, pkinit := extractEData(krbErr.EData)
			out.Enctypes = enctypes
			out.PKINITOffered = pkinit
		}
		return
	}

	// Try AS-REP. If this succeeds, preauth wasn't required — surface it.
	var asRep messages.ASRep
	if err := asRep.Unmarshal(raw); err == nil {
		out.PrincipalFound = true
		out.ServerRealm = asRep.CRealm
		out.Raw = "AS-REP received without preauth (AS-REP roasting exposure)"
		return
	}

	out.Error = "unable to parse KDC reply (" + hex.EncodeToString(first(raw, 16)) + ")"
}

func first(b []byte, n int) []byte {
	if len(b) < n {
		return b
	}
	return b[:n]
}

// extractEData parses the METHOD-DATA / PADataSequence that KDCs attach to
// PREAUTH_REQUIRED errors and returns the advertised enctypes + pkinit hint.
func extractEData(edata []byte) ([]EnctypeEntry, bool) {
	var pas types.PADataSequence
	if _, err := asn1.Unmarshal(edata, &pas); err != nil {
		return nil, false
	}
	var out []EnctypeEntry
	pkinit := false
	for _, pa := range pas {
		switch pa.PADataType {
		case patype.PA_ETYPE_INFO2:
			var info types.ETypeInfo2
			if _, err := asn1.Unmarshal(pa.PADataValue, &info); err == nil {
				for _, e := range info {
					out = append(out, EnctypeEntry{
						ID:     e.EType,
						Name:   etypeName(e.EType),
						Weak:   weakEnctypes[e.EType],
						Salt:   e.Salt,
						Source: "etype-info2",
					})
				}
			}
		case patype.PA_ETYPE_INFO:
			var info types.ETypeInfo
			if _, err := asn1.Unmarshal(pa.PADataValue, &info); err == nil {
				for _, e := range info {
					if hasEnctype(out, e.EType) {
						continue
					}
					out = append(out, EnctypeEntry{
						ID:     e.EType,
						Name:   etypeName(e.EType),
						Weak:   weakEnctypes[e.EType],
						Salt:   string(e.Salt),
						Source: "etype-info",
					})
				}
			}
		case patype.PA_PK_AS_REQ, patype.PA_PK_AS_REQ_OLD:
			pkinit = true
		}
	}
	return out, pkinit
}

func hasEnctype(list []EnctypeEntry, id int32) bool {
	for _, e := range list {
		if e.ID == id {
			return true
		}
	}
	return false
}

func errorcodeName(code int32) string {
	s := errorcode.Lookup(code)
	if _, after, ok := strings.Cut(s, ") "); ok {
		s = after
	}
	if before, _, ok := strings.Cut(s, " "); ok {
		s = before
	}
	return s
}

// ---- helpers ----------------------------------------------------------------

func optFloat(opts sdk.CheckerOptions, key string, def float64) float64 {
	v, ok := opts[key]
	if !ok {
		return def
	}
	switch x := v.(type) {
	case float64:
		return x
	case float32:
		return float64(x)
	case int:
		return float64(x)
	case int64:
		return float64(x)
	case string:
		if f, err := strconv.ParseFloat(x, 64); err == nil {
			return f
		}
	}
	return def
}

func stringOpt(opts sdk.CheckerOptions, key string) string {
	s, _ := opts[key].(string)
	return s
}