// SPDX-License-Identifier: BSD-3-Clause

package query

import (
	"fmt"
	"strconv"
	"strings"
	"time"

	"git.froth.zone/sam/awl/util"
	"github.com/dchest/uniuri"
	"github.com/miekg/dns"
	"golang.org/x/net/idna"
)

const (
	tcp = "tcp"
	udp = "udp"
)

// ToString turns the response into something that looks a lot like dig
//
// Much of this is taken from https://github.com/miekg/dns/blob/master/msg.go#L900
func ToString(res util.Response, opts util.Options) (string, error) {
	if res.DNS == nil {
		return "<nil> MsgHdr", errNoMessage
	}

	var (
		s   string
		opt *dns.OPT
	)

	if !opts.Short {
		if opts.Display.Comments {
			s += res.DNS.MsgHdr.String() + " "
			s += "QUERY: " + strconv.Itoa(len(res.DNS.Question)) + ", "
			s += "ANSWER: " + strconv.Itoa(len(res.DNS.Answer)) + ", "
			s += "AUTHORITY: " + strconv.Itoa(len(res.DNS.Ns)) + ", "
			s += "ADDITIONAL: " + strconv.Itoa(len(res.DNS.Extra)) + "\n"
			opt = res.DNS.IsEdns0()

			if opt != nil && opts.Display.Opt {
				// OPT PSEUDOSECTION
				s += opt.String() + "\n"
			}
		}

		if opts.Display.Question {
			if len(res.DNS.Question) > 0 {
				if opts.Display.Comments {
					s += "\n;; QUESTION SECTION:\n"
				}

				for _, r := range res.DNS.Question {
					str, err := stringParse(r.String(), false, opts)
					if err != nil {
						return "", fmt.Errorf("%w", err)
					}

					s += str + "\n"
				}
			}
		}

		if opts.Display.Answer {
			if len(res.DNS.Answer) > 0 {
				if opts.Display.Comments {
					s += "\n;; ANSWER SECTION:\n"
				}

				for _, r := range res.DNS.Answer {
					if r != nil {
						str, err := stringParse(r.String(), true, opts)
						if err != nil {
							return "", fmt.Errorf("%w", err)
						}

						s += str + "\n"
					}
				}
			}
		}

		if opts.Display.Authority {
			if len(res.DNS.Ns) > 0 {
				if opts.Display.Comments {
					s += "\n;; AUTHORITY SECTION:\n"
				}

				for _, r := range res.DNS.Ns {
					if r != nil {
						str, err := stringParse(r.String(), true, opts)
						if err != nil {
							return "", fmt.Errorf("%w", err)
						}

						s += str + "\n"
					}
				}
			}
		}

		if opts.Display.Additional {
			if len(res.DNS.Extra) > 0 && (opt == nil || len(res.DNS.Extra) > 1) {
				if opts.Display.Comments {
					s += "\n;; ADDITIONAL SECTION:\n"
				}

				for _, r := range res.DNS.Extra {
					if r != nil && r.Header().Rrtype != dns.TypeOPT {
						str, err := stringParse(r.String(), true, opts)
						if err != nil {
							return "", fmt.Errorf("%w", err)
						}

						s += str + "\n"
					}
				}
			}
		}

		if opts.Display.Statistics {
			s += "\n;; Query time: " + res.RTT.String()
			// Add extra information to server string
			var extra string

			switch {
			case opts.TCP:
				extra = ":" + strconv.Itoa(opts.Request.Port) + " (TCP)"
			case opts.TLS:
				extra = ":" + strconv.Itoa(opts.Request.Port) + " (TLS)"
			case opts.HTTPS, opts.DNSCrypt:
				extra = ""
			case opts.QUIC:
				extra = ":" + strconv.Itoa(opts.Request.Port) + " (QUIC)"
			default:
				extra = ":" + strconv.Itoa(opts.Request.Port) + " (UDP)"
			}

			s += "\n;; SERVER: " + opts.Request.Server + extra
			s += "\n;; WHEN: " + time.Now().Format(time.RFC1123Z)
			s += "\n;; MSG SIZE  rcvd: " + strconv.Itoa(res.DNS.Len()) + "\n"
		}
	} else {
		// Print just the responses, nothing else
		for i, resp := range res.DNS.Answer {
			temp := strings.Split(resp.String(), "\t")
			s += temp[len(temp)-1]

			if opts.Identify {
				s += " from server " + opts.Request.Server + " in " + res.RTT.String()
			}

			// Don't print newline on last line
			if i != len(res.DNS.Answer)-1 {
				s += "\n"
			}
		}
	}

	return s, nil
}

// stringParse edits the raw responses to user requests.
func stringParse(str string, isAns bool, opts util.Options) (string, error) {
	split := strings.Split(str, "\t")

	// Make edits if so requested

	// TODO: make less ew?
	// This exists because the question section should be left alone EXCEPT for punycode.

	if isAns {
		if !opts.Display.TTL {
			// Remove from existence
			split = append(split[:1], split[2:]...)
		}

		if !opts.Display.ShowClass {
			// Position depends on if the TTL is there or not.
			if opts.Display.TTL {
				split = append(split[:2], split[3:]...)
			} else {
				split = append(split[:1], split[2:]...)
			}
		}

		if opts.Display.TTL && opts.Display.HumanTTL {
			ttl, _ := strconv.Atoi(split[1])
			split[1] = (time.Duration(ttl) * time.Second).String()
		}
	}

	if opts.Display.UcodeTranslate {
		var (
			err  error
			semi string
		)

		if strings.HasPrefix(split[0], ";") {
			split[0] = strings.TrimPrefix(split[0], ";")
			semi = ";"
		}

		split[0], err = idna.ToUnicode(split[0])
		if err != nil {
			return "", fmt.Errorf("punycode: %w", err)
		}

		split[0] = semi + split[0]
	}

	return strings.Join(split, "\t"), nil
}

// CreateQuery creates a DNS query from the options given.
// It sets query flags and EDNS flags from the respective options.
func CreateQuery(opts util.Options) (util.Response, error) {
	req := new(dns.Msg)
	req.SetQuestion(opts.Request.Name, opts.Request.Type)
	req.Question[0].Qclass = opts.Request.Class

	// Set standard flags
	req.MsgHdr.Response = opts.QR
	req.MsgHdr.Authoritative = opts.AA
	req.MsgHdr.Truncated = opts.TC
	req.MsgHdr.RecursionDesired = opts.RD
	req.MsgHdr.RecursionAvailable = opts.RA
	req.MsgHdr.Zero = opts.Z
	req.MsgHdr.AuthenticatedData = opts.AD
	req.MsgHdr.CheckingDisabled = opts.CD

	// EDNS time :)
	if opts.EDNS.EnableEDNS {
		o := new(dns.OPT)
		o.Hdr.Name = "."
		o.Hdr.Rrtype = dns.TypeOPT

		o.SetVersion(opts.EDNS.Version)

		if opts.EDNS.Cookie {
			e := new(dns.EDNS0_COOKIE)
			e.Code = dns.EDNS0COOKIE
			e.Cookie = uniuri.NewLenChars(8, []byte("1234567890abcdef"))
			o.Option = append(o.Option, e)

			opts.Logger.Info("Setting EDNS cookie to", e.Cookie)
		}

		if opts.EDNS.Expire {
			o.Option = append(o.Option, new(dns.EDNS0_EXPIRE))

			opts.Logger.Info("Setting EDNS Expire option")
		}

		if opts.EDNS.KeepOpen {
			o.Option = append(o.Option, new(dns.EDNS0_TCP_KEEPALIVE))

			opts.Logger.Info("Setting EDNS TCP Keepalive option")
		}

		if opts.EDNS.Nsid {
			o.Option = append(o.Option, new(dns.EDNS0_NSID))

			opts.Logger.Info("Setting EDNS NSID option")
		}

		if opts.EDNS.Padding {
			o.Option = append(o.Option, new(dns.EDNS0_PADDING))

			opts.Logger.Info("Setting EDNS padding")
		}

		o.SetUDPSize(opts.EDNS.BufSize)

		opts.Logger.Info("EDNS UDP buffer set to", opts.EDNS.BufSize)

		o.SetZ(opts.EDNS.ZFlag)

		opts.Logger.Info("EDNS Z flag set to", opts.EDNS.ZFlag)

		if opts.EDNS.DNSSEC {
			o.SetDo()

			opts.Logger.Info("EDNS DNSSEC OK set")
		}

		if opts.EDNS.Subnet.Address != nil {
			o.Option = append(o.Option, &opts.EDNS.Subnet)
		}

		req.Extra = append(req.Extra, o)
	} else if opts.EDNS.DNSSEC {
		req.SetEdns0(1232, true)
		opts.Logger.Warn("DNSSEC implies EDNS, EDNS enabled")
		opts.Logger.Info("DNSSEC enabled, UDP buffer set to 1232")
	}

	opts.Logger.Debug(req)

	if !opts.Short {
		if opts.Display.ShowQuery {
			opts.Logger.Info("Printing constructed query")

			var (
				str string
				err error
			)

			if opts.JSON || opts.XML || opts.YAML {
				str, err = PrintSpecial(req, opts)
				if err != nil {
					return util.Response{}, err
				}
			} else {
				temp := opts.Display.Statistics
				opts.Display.Statistics = false
				str, err = ToString(
					util.Response{
						DNS: req,
						RTT: 0,
					}, opts)
				if err != nil {
					return util.Response{}, err
				}

				opts.Display.Statistics = temp
				str += "\n;; QUERY SIZE: " + strconv.Itoa(req.Len()) + "\n"
			}

			fmt.Println(str)

			opts.Display.ShowQuery = false
		}
	}

	resolver, err := LoadResolver(opts)
	if err != nil {
		return util.Response{}, err
	}

	opts.Logger.Info("Query successfully loaded")

	//nolint:wrapcheck // Error wrapping not needed here
	return resolver.LookUp(req)
}