package dns

import (
	"fmt"
	"net"
	"strings"

	"github.com/Dreamacro/clash/common/cache"
	"github.com/Dreamacro/clash/component/fakeip"
	"github.com/Dreamacro/clash/log"

	D "github.com/miekg/dns"
)

type handler func(w D.ResponseWriter, r *D.Msg)

func withFakeIP(cache *cache.Cache, pool *fakeip.Pool) handler {
	return func(w D.ResponseWriter, r *D.Msg) {
		q := r.Question[0]

		cacheItem := cache.Get("fakeip:" + q.String())
		if cacheItem != nil {
			msg := cacheItem.(*D.Msg).Copy()
			setMsgTTL(msg, 1)
			msg.SetReply(r)
			w.WriteMsg(msg)
			return
		}

		rr := &D.A{}
		rr.Hdr = D.RR_Header{Name: q.Name, Rrtype: D.TypeA, Class: D.ClassINET, Ttl: dnsDefaultTTL}
		ip := pool.Get()
		rr.A = ip
		msg := r.Copy()
		msg.Answer = []D.RR{rr}
		putMsgToCache(cache, "fakeip:"+q.String(), msg)
		putMsgToCache(cache, ip.String(), msg)

		setMsgTTL(msg, 1)
		return
	}
}

func withResolver(resolver *Resolver) handler {
	return func(w D.ResponseWriter, r *D.Msg) {
		msg, err := resolver.Exchange(r)

		if err != nil {
			q := r.Question[0]
			qString := fmt.Sprintf("%s %s %s", q.Name, D.Class(q.Qclass).String(), D.Type(q.Qtype).String())
			log.Debugln("[DNS Server] Exchange %s failed: %v", qString, err)
			D.HandleFailed(w, r)
			return
		}
		msg.SetReply(r)
		w.WriteMsg(msg)
		return
	}
}

func withHost(resolver *Resolver, next handler) handler {
	hosts := resolver.hosts
	if hosts == nil {
		panic("dns/withHost: hosts should not be nil")
	}

	return func(w D.ResponseWriter, r *D.Msg) {
		q := r.Question[0]
		if q.Qtype != D.TypeA && q.Qtype != D.TypeAAAA {
			next(w, r)
			return
		}

		domain := strings.TrimRight(q.Name, ".")
		host := hosts.Search(domain)
		if host == nil {
			next(w, r)
			return
		}

		ip := host.Data.(net.IP)
		if q.Qtype == D.TypeAAAA && ip.To16() == nil {
			next(w, r)
			return
		} else if q.Qtype == D.TypeA && ip.To4() == nil {
			next(w, r)
			return
		}

		var rr D.RR
		if q.Qtype == D.TypeAAAA {
			record := &D.AAAA{}
			record.Hdr = D.RR_Header{Name: q.Name, Rrtype: D.TypeAAAA, Class: D.ClassINET, Ttl: dnsDefaultTTL}
			record.AAAA = ip
			rr = record
		} else {
			record := &D.A{}
			record.Hdr = D.RR_Header{Name: q.Name, Rrtype: D.TypeA, Class: D.ClassINET, Ttl: dnsDefaultTTL}
			record.A = ip
			rr = record
		}

		msg := r.Copy()
		msg.Answer = []D.RR{rr}
		msg.SetReply(r)
		w.WriteMsg(msg)
		return
	}
}

func newHandler(resolver *Resolver) handler {
	if resolver.IsFakeIP() {
		return withFakeIP(resolver.cache, resolver.pool)
	}

	if resolver.hosts != nil {
		return withHost(resolver, withResolver(resolver))
	}

	return withResolver(resolver)
}