package checker

import (
	"context"
	"crypto/sha256"
	"crypto/sha512"
	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"fmt"
	"net"
	"strconv"
	"strings"
	"time"

	"git.happydns.org/checker-tls/contract"
)

// buildChain returns CertInfo for each cert presented by the server, in the
// order the server sent them (leaf first). SPKI is extracted from the parsed
// certificate's RawSubjectPublicKeyInfo so we hash exactly the DER bytes
// DANE selector 1 refers to (RFC 6698 §1.1.3).
func buildChain(certs []*x509.Certificate) []CertInfo {
	out := make([]CertInfo, len(certs))
	for i, c := range certs {
		certSum256 := sha256.Sum256(c.Raw)
		certSum512 := sha512.Sum512(c.Raw)
		spkiSum256 := sha256.Sum256(c.RawSubjectPublicKeyInfo)
		spkiSum512 := sha512.Sum512(c.RawSubjectPublicKeyInfo)
		out[i] = CertInfo{
			DERBase64:     base64.StdEncoding.EncodeToString(c.Raw),
			Subject:       c.Subject.String(),
			Issuer:        c.Issuer.String(),
			NotAfter:      c.NotAfter,
			CertSHA256:    hex.EncodeToString(certSum256[:]),
			CertSHA512:    hex.EncodeToString(certSum512[:]),
			SPKISHA256:    hex.EncodeToString(spkiSum256[:]),
			SPKISHA512:    hex.EncodeToString(spkiSum512[:]),
			SPKIDERBase64: base64.StdEncoding.EncodeToString(c.RawSubjectPublicKeyInfo),
		}
	}
	return out
}

// probeTypeString renders the TLSProbe.Type string from a TLSEndpoint.
// Observation consumers already parse this field in its "tls" /
// "starttls-<proto>" shape; the contract-level split of direct vs.
// STARTTLS is collapsed back here so the wire format of tls_probes
// stays unchanged.
func probeTypeString(ep contract.TLSEndpoint) string {
	if ep.STARTTLS == "" {
		return "tls"
	}
	return "starttls-" + ep.STARTTLS
}

// probe performs a TLS handshake (or STARTTLS upgrade + handshake) on the
// given endpoint and returns a populated TLSProbe. It never returns an error:
// transport/handshake failures are recorded on the probe as raw fields so
// rules can classify them.
//
// This function MUST NOT decide severity or pass/fail: it only gathers
// observation data. All judgement happens in CheckRules (see rules_*.go).
func probe(ctx context.Context, ep contract.TLSEndpoint, timeout time.Duration) TLSProbe {
	start := time.Now()
	host := strings.TrimSuffix(ep.Host, ".")
	addr := net.JoinHostPort(host, strconv.Itoa(int(ep.Port)))
	sni := ep.SNI
	if sni == "" {
		sni = host
	}

	p := TLSProbe{
		Host:            host,
		Port:            ep.Port,
		Endpoint:        addr,
		Type:            probeTypeString(ep),
		SNI:             sni,
		RequireSTARTTLS: ep.RequireSTARTTLS,
		STARTTLSDialect: ep.STARTTLS,
	}

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

	d := &net.Dialer{}
	conn, err := d.DialContext(dialCtx, "tcp", addr)
	if err != nil {
		p.TCPError = err.Error()
		p.Error = "dial: " + err.Error()
		p.ElapsedMS = time.Since(start).Milliseconds()
		return p
	}
	defer conn.Close()

	if deadline, ok := dialCtx.Deadline(); ok {
		_ = conn.SetDeadline(deadline)
	}

	tlsConn, err := handshake(conn, ep, sni)
	if err != nil {
		p.HandshakeError = err.Error()
		p.Error = err.Error()
		if ep.STARTTLS != "" && isStartTLSUnsupported(err) {
			p.STARTTLSNotOffered = true
		}
		if errors.Is(err, errUnsupportedStartTLSProto) {
			p.STARTTLSUnsupportedProto = true
		}
		p.ElapsedMS = time.Since(start).Milliseconds()
		return p
	}
	defer tlsConn.Close()

	p.TLSHandshakeOK = true
	state := tlsConn.ConnectionState()
	p.TLSVersionNum = state.Version
	p.TLSVersion = tls.VersionName(state.Version)
	p.CipherSuite = tls.CipherSuiteName(state.CipherSuite)
	p.CipherSuiteID = state.CipherSuite

	if len(state.PeerCertificates) == 0 {
		p.NoPeerCert = true
		p.ElapsedMS = time.Since(start).Milliseconds()
		return p
	}

	leaf := state.PeerCertificates[0]
	p.NotAfter = leaf.NotAfter
	p.Issuer = leaf.Issuer.CommonName
	p.IssuerDN = leaf.Issuer.String()
	if len(leaf.AuthorityKeyId) > 0 {
		p.IssuerAKI = strings.ToUpper(hex.EncodeToString(leaf.AuthorityKeyId))
	}
	p.Subject = leaf.Subject.CommonName
	p.DNSNames = leaf.DNSNames
	p.Chain = buildChain(state.PeerCertificates)

	hostnameMatch := leaf.VerifyHostname(sni) == nil
	p.HostnameMatch = &hostnameMatch

	// Chain verification against system roots, using intermediates presented
	// by the server. Running it separately from tls.Config verification
	// means we can record it as a raw observation rather than aborting the
	// handshake, rules classify it afterwards.
	intermediates := x509.NewCertPool()
	for _, c := range state.PeerCertificates[1:] {
		intermediates.AddCert(c)
	}
	now := time.Now()
	_, verifyErr := leaf.Verify(x509.VerifyOptions{
		DNSName:       sni,
		Intermediates: intermediates,
		CurrentTime:   now,
	})
	chainValid := verifyErr == nil
	p.ChainValid = &chainValid
	if verifyErr != nil {
		p.ChainVerifyErr = verifyErr.Error()
	}

	p.ElapsedMS = time.Since(start).Milliseconds()
	return p
}

// handshake performs STARTTLS upgrade (when ep.STARTTLS is non-empty) and
// then a TLS handshake. InsecureSkipVerify is true on purpose: we verify
// the chain separately in probe so an invalid chain becomes a raw
// observation rather than aborting the handshake.
func handshake(conn net.Conn, ep contract.TLSEndpoint, sni string) (*tls.Conn, error) {
	cfg := &tls.Config{
		ServerName:         sni,
		InsecureSkipVerify: true, // #nosec G402 -- intentional: chain verified separately in probe()
	}

	if ep.STARTTLS == "" {
		tlsConn := tls.Client(conn, cfg)
		if err := tlsConn.Handshake(); err != nil {
			return nil, fmt.Errorf("tls-handshake: %w", err)
		}
		return tlsConn, nil
	}

	up, ok := starttlsUpgraders[ep.STARTTLS]
	if !ok {
		return nil, fmt.Errorf("%w: %q", errUnsupportedStartTLSProto, ep.STARTTLS)
	}
	if err := up(conn, sni); err != nil {
		return nil, fmt.Errorf("starttls-%s: %w", ep.STARTTLS, err)
	}
	tlsConn := tls.Client(conn, cfg)
	if err := tlsConn.Handshake(); err != nil {
		return nil, fmt.Errorf("tls-handshake-after-starttls: %w", err)
	}
	return tlsConn, nil
}

var (
	errStartTLSNotOffered       = errors.New("starttls not advertised by server")
	errUnsupportedStartTLSProto = errors.New("unsupported starttls protocol")
)

func isStartTLSUnsupported(err error) bool {
	return errors.Is(err, errStartTLSNotOffered)
}