package dns

import (
	"fmt"

	"github.com/miekg/dns"

	"code.crute.me/mcrute/go_ddns_manager/bind"
)

type DNSClient struct {
	Server string
}

type DNSTransaction struct {
	zone *bind.Zone
	key  *bind.Key
	msg  *dns.Msg
}

func (t *DNSTransaction) Upsert(rrs ...RR) *DNSTransaction {
	t.RemoveAll(rrs...).Insert(rrs...)
	return t
}

func (t *DNSTransaction) Insert(rrs ...RR) *DNSTransaction {
	t.msg.Insert(toRRSet(t.zone, rrs...))
	return t
}

func (t *DNSTransaction) Remove(rrs ...RR) *DNSTransaction {
	t.msg.Remove(toRRSet(t.zone, rrs...))
	return t
}

func (t *DNSTransaction) RemoveAll(rrs ...RR) *DNSTransaction {
	t.msg.RemoveRRset(toRRSet(t.zone, rrs...))
	return t
}

func (c *DNSClient) AXFR(zone *bind.Zone) (chan *dns.Envelope, error) {
	k := zone.Keys()[0]
	t := &dns.Transfer{TsigSecret: k.AsMap()} // Always uses tcp

	m := &dns.Msg{}
	m.SetAxfr(zone.Name)
	k.Sign(m)

	return t.In(m, c.Server)
}

func (c *DNSClient) ReadRemoteZone(zone *bind.Zone) ([]RR, error) {
	rrs := []RR{}
	seenSoa := false

	data, err := c.AXFR(zone)
	if err != nil {
		return nil, err
	}

	for rd := range data {
		for _, r := range rd.RR {
			switch dr := FromDNS(r).(type) {
			case *SOA:
				// Transfers have 2 SOA records, exclude the last one
				if !seenSoa {
					rrs = append(rrs, dr)
					seenSoa = true
				}
			case RR:
				rrs = append(rrs, dr)
			default:
				// This should only be possible if we somehow are
				// missing generated DNS data types.
				return nil, fmt.Errorf("Invalid return type")
			}
		}
	}

	return rrs, nil
}

func (c *DNSClient) StartUpdate(zone *bind.Zone) *DNSTransaction {
	m := &dns.Msg{}
	m.SetUpdate(zone.Name)

	return &DNSTransaction{
		zone: zone,
		key:  zone.Keys()[0],
		msg:  m,
	}
}

func (c *DNSClient) QueryRecursive(zone *bind.Zone, fqdn string, rtype uint16) *DNSTransaction {
	m := &dns.Msg{}
	m.RecursionDesired = true
	m.SetQuestion(fqdn, rtype)

	return &DNSTransaction{
		zone: zone,
		key:  zone.Keys()[0],
		msg:  m,
	}
}

func (c *DNSClient) SendUpdate(t *DNSTransaction) error {
	udp := &dns.Client{Net: "udp", TsigSecret: t.key.AsMap()}
	tcp := &dns.Client{Net: "tcp", TsigSecret: t.key.AsMap()}

	t.msg.SetEdns0(4096, false)
	t.key.Sign(t.msg)

	in, _, err := udp.Exchange(t.msg, c.Server)
	if in != nil && in.Truncated {
		// If the TCP request succeeds, the err will reset to nil
		in, _, err = tcp.Exchange(t.msg, c.Server)
	}

	if err != nil {
		return err
	}

	return nil
}

func (c *DNSClient) SendQuery(t *DNSTransaction) ([]dns.RR, error) {
	udp := &dns.Client{Net: "udp", TsigSecret: t.key.AsMap()}
	tcp := &dns.Client{Net: "tcp", TsigSecret: t.key.AsMap()}

	t.msg.SetEdns0(4096, false)
	t.key.Sign(t.msg)

	in, _, err := udp.Exchange(t.msg, c.Server)
	if in != nil && in.Truncated {
		// If the TCP request succeeds, the err will reset to nil
		in, _, err = tcp.Exchange(t.msg, c.Server)
	}

	if err != nil {
		return nil, err
	}

	return in.Answer, nil
}