package acme

import (
	"crypto/sha256"
	"encoding/base64"
	"errors"
	"fmt"
	"log"
	"net"
	"strings"
	"time"

	"github.com/miekg/dns"
)

type preCheckDNSFunc func(domain, fqdn, value string) error

var preCheckDNS preCheckDNSFunc = checkDnsPropagation

var recursionMaxDepth = 10

// DNS01Record returns a DNS record which will fulfill the `dns-01` challenge
func DNS01Record(domain, keyAuth string) (fqdn string, value string, ttl int) {
	keyAuthShaBytes := sha256.Sum256([]byte(keyAuth))
	// base64URL encoding without padding
	keyAuthSha := base64.URLEncoding.EncodeToString(keyAuthShaBytes[:sha256.Size])
	value = strings.TrimRight(keyAuthSha, "=")
	ttl = 120
	fqdn = fmt.Sprintf("_acme-challenge.%s.", domain)
	return
}

// dnsChallenge implements the dns-01 challenge according to ACME 7.5
type dnsChallenge struct {
	jws      *jws
	validate validateFunc
	provider ChallengeProvider
}

func (s *dnsChallenge) Solve(chlng challenge, domain string) error {
	logf("[INFO][%s] acme: Trying to solve DNS-01", domain)

	if s.provider == nil {
		return errors.New("No DNS Provider configured")
	}

	// Generate the Key Authorization for the challenge
	keyAuth, err := getKeyAuthorization(chlng.Token, &s.jws.privKey.PublicKey)
	if err != nil {
		return err
	}

	err = s.provider.Present(domain, chlng.Token, keyAuth)
	if err != nil {
		return fmt.Errorf("Error presenting token %s", err)
	}
	defer func() {
		err := s.provider.CleanUp(domain, chlng.Token, keyAuth)
		if err != nil {
			log.Printf("Error cleaning up %s %v ", domain, err)
		}
	}()

	fqdn, value, _ := DNS01Record(domain, keyAuth)

	logf("[INFO][%s] Checking DNS record propagation...", domain)

	if err = preCheckDNS(domain, fqdn, value); err != nil {
		return err
	}

	return s.validate(s.jws, domain, chlng.URI, challenge{Resource: "challenge", Type: chlng.Type, Token: chlng.Token, KeyAuthorization: keyAuth})
}

// checkDnsPropagation checks if the expected TXT record has been propagated to
// all authoritative nameservers. If not it waits and retries for some time.
func checkDnsPropagation(domain, fqdn, value string) error {
	authoritativeNss, err := lookupNameservers(toFqdn(domain))
	if err != nil {
		return err
	}

	if err = waitFor(30, 2, func() (bool, error) {
		return checkAuthoritativeNss(fqdn, value, authoritativeNss)
	}); err != nil {
		return err
	}

	return nil
}

// checkAuthoritativeNss queries each of the given nameservers for the expected TXT record.
func checkAuthoritativeNss(fqdn, value string, nameservers []string) (bool, error) {
	for _, ns := range nameservers {
		r, err := dnsQuery(fqdn, dns.TypeTXT, ns, false)
		if err != nil {
			return false, err
		}

		if r.Rcode != dns.RcodeSuccess {
			return false, fmt.Errorf("%s returned RCode %s", ns, dns.RcodeToString[r.Rcode])
		}

		var found bool
		for _, rr := range r.Answer {
			if txt, ok := rr.(*dns.TXT); ok {
				if strings.Join(txt.Txt, "") == value {
					found = true
					break
				}
			}
		}

		if !found {
			return false, fmt.Errorf("%s did not return the expected TXT record", ns)
		}
	}

	return true, nil
}

// dnsQuery sends a DNS query to the given nameserver.
func dnsQuery(fqdn string, rtype uint16, nameserver string, recursive bool) (in *dns.Msg, err error) {
	m := new(dns.Msg)
	m.SetQuestion(fqdn, rtype)
	m.SetEdns0(4096, false)
	if !recursive {
		m.RecursionDesired = false
	}

	in, err = dns.Exchange(m, net.JoinHostPort(nameserver, "53"))
	if err == dns.ErrTruncated {
		tcp := &dns.Client{Net: "tcp"}
		in, _, err = tcp.Exchange(m, nameserver)
	}

	return
}

// lookupNameservers returns the authoritative nameservers for the given domain name.
func lookupNameservers(fqdn string) ([]string, error) {
	var err error
	var r *dns.Msg
	var authoritativeNss []string
	resolver := "google-public-dns-a.google.com"

	r, err = dnsQuery(fqdn, dns.TypeSOA, resolver, true)
	if err != nil {
		return nil, err
	}

	// If there is a SOA RR in the Answer section then fqdn is the root domain.
	for _, rr := range r.Answer {
		if soa, ok := rr.(*dns.SOA); ok {
			r, err = dnsQuery(soa.Hdr.Name, dns.TypeNS, resolver, true)
			if err != nil {
				return nil, err
			}

			for _, rr := range r.Answer {
				if ns, ok := rr.(*dns.NS); ok {
					authoritativeNss = append(authoritativeNss, strings.ToLower(ns.Ns))
				}
			}

			return authoritativeNss, nil
		}
	}

	// Strip of the left most label to get the parent domain.
	offset, _ := dns.NextLabel(fqdn, 0)
	next := fqdn[offset:]
	// Only the TLD label left. This should not happen if the domain DNS is healthy.
	if dns.CountLabel(next) < 2 {
		return nil, fmt.Errorf("Could not determine root domain")
	}
	
	return lookupNameservers(fqdn[offset:])
}

// toFqdn converts the name into a fqdn appending a trailing dot.
func toFqdn(name string) string {
	n := len(name)
	if n == 0 || name[n-1] == '.' {
		return name
	}
	return name + "."
}

// unFqdn converts the fqdn into a name removing the trailing dot.
func unFqdn(name string) string {
	n := len(name)
	if n != 0 && name[n-1] == '.' {
		return name[:n-1]
	}
	return name
}

// waitFor polls the given function 'f', once every 'interval' seconds, up to 'timeout' seconds.
func waitFor(timeout, interval int, f func() (bool, error)) error {
	var lastErr string
	timeup := time.After(time.Duration(timeout) * time.Second)
	for {
		select {
		case <-timeup:
			return fmt.Errorf("Time limit exceeded. Last error: %s", lastErr)
		default:
		}

		stop, err := f()
		if stop {
			return nil
		}
		if err != nil {
			lastErr = err.Error()
		}

		time.Sleep(time.Duration(interval) * time.Second)
	}
}