package override

import (
	"net"

	"github.com/miekg/dns"
)

type Rule struct {
	Search   net.IP
	Override net.IP
}

type ResponseWrapper struct {
	dns.ResponseWriter
	Rules []Rule
}

func (r *ResponseWrapper) WriteMsg(res *dns.Msg) error {
	for _, r := range r.Rules {
		for _, rr := range res.Answer {
			overideRR(r, rr)
		}

		for _, rr := range res.Extra {
			overideRR(r, rr)
		}
	}

	return r.ResponseWriter.WriteMsg(res)
}

func overideRR(r Rule, rr dns.RR) {
	if rr.Header().Rrtype == dns.TypeA {
		a := rr.(*dns.A)
		if a.A.Equal(r.Search) {
			a.A = r.Override
		}
	} else if rr.Header().Rrtype == dns.TypeAAAA {
		a := rr.(*dns.AAAA)
		if a.AAAA.Equal(r.Search) {
			a.AAAA = r.Override
		}
	}
}

func (r *ResponseWrapper) Write(buf []byte) (int, error) {
	log.Warning("ResponseWrapper called with Write: not ensuring overide")
	return r.ResponseWriter.Write(buf)
}