package main

import (
	"bytes"
	"context"
	"crypto/tls"
	"encoding/base64"
	"encoding/binary"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/ameshkov/dnscrypt/v2"
	"github.com/jedisct1/go-dnsstamps"
	rotatelogs "github.com/lestrrat-go/file-rotatelogs"
	"github.com/miekg/dns"
	quic "github.com/quic-go/quic-go"
	"github.com/quic-go/quic-go/http3"
	"golang.org/x/net/http2"
	"golang.org/x/net/idna"
	"gopkg.in/yaml.v3"
)

type Config struct {
	DNSServers  []string  `yaml:"dns_servers"`
	LogConfig   LogConfig `yaml:"log_config"`
	Port        string    `yaml:"port"`
	UseTLS      bool      `yaml:"use_tls"`
	TLSCertFile string    `yaml:"tls_cert_file"`
	TLSKeyFile  string    `yaml:"tls_key_file"`
}

type LogConfig struct {
	LogToFile    bool   `yaml:"logEnabled"`
	LogFile      string `yaml:"log_file"`
	MaxAgeDays   int    `yaml:"max_age_days"`
	RotationDays int    `yaml:"rotation_days"`
}

var (
	config     Config
	logger     *log.Logger
	logChannel chan string
	wg         sync.WaitGroup // 用于等待所有 goroutine 完成
)

func main() {
	configFile := flag.String("config", "config.yaml", "YAML配置文件路径")
	flag.Parse()

	loadConfig(*configFile)
	initLogger()
	// 启动日志处理goroutine
	go logProcessor()

	proxy := NewModifyRemoteAddrReverseProxy()

	http.HandleFunc("/resolve", proxy.ServeHTTP)
	http.HandleFunc("/dns-query", dnsQueryHandler)
	fmt.Printf("DNS解析器监听端口 :%s...\n", config.Port)

	var err error
	if config.UseTLS {
		err = http.ListenAndServeTLS(":"+config.Port, config.TLSCertFile, config.TLSKeyFile, nil)
	} else {
		err = http.ListenAndServe(":"+config.Port, nil)
	}

	if err != nil {
		fmt.Printf("服务器启动失败: %v\n", err)
		os.Exit(1)
	}
}

type ModifyRemoteAddrReverseProxy struct {
	httputil.ReverseProxy
}

func NewModifyRemoteAddrReverseProxy() *ModifyRemoteAddrReverseProxy {
	return &ModifyRemoteAddrReverseProxy{
		ReverseProxy: httputil.ReverseProxy{},
	}
}

// 辅助函数用于移除重复的 DNS 记录
func removeDuplicates(records []dns.RR) []dns.RR {
	seen := make(map[string]bool) // 创建一个映射，用于跟踪已经添加的记录
	result := []dns.RR{}          // 创建一个切片，用于存放去重后的记录
	for _, record := range records {
		// 使用记录的名称、类型和字符串表示形式生成唯一键
		key := fmt.Sprintf("%s-%d-%s", record.Header().Name, record.Header().Rrtype, record.String())
		if !seen[key] { // 如果记录还没有被处理过
			seen[key] = true                // 标记记录为已处理
			result = append(result, record) // 将记录添加到结果切片中
		}
	}
	return result // 返回去重后的记录切片
}

func (p *ModifyRemoteAddrReverseProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	name := r.URL.Query().Get("name")
	if name == "" {
		writeErrorResponse(w, "参数 'name' 是必需的", http.StatusBadRequest)
		return
	}

	// 将中文域名转换为 ASCII（Punycode）
	asciiName, err := idna.ToASCII(name)
	if err != nil {
		writeErrorResponse(w, "域名编码转换失败", http.StatusBadRequest)
		return
	}

	if !strings.HasSuffix(asciiName, ".") {
		asciiName += "."
	}

	clientIP := getClientIP(r)

	// 创建两个 DNS 查询，一个查询 A 记录，一个查询 AAAA 记录
	queryA := dns.Msg{}
	queryA.SetQuestion(dns.Fqdn(asciiName), dns.TypeA)
	edns0Subnet := getEDNS0SubnetOption(clientIP)
	edns0 := &dns.OPT{Hdr: dns.RR_Header{Name: ".", Rrtype: dns.TypeOPT}, Option: []dns.EDNS0{edns0Subnet}}
	queryA.Extra = append(queryA.Extra, edns0)

	queryAAAA := dns.Msg{}
	queryAAAA.SetQuestion(dns.Fqdn(asciiName), dns.TypeAAAA)
	queryAAAA.Extra = append(queryAAAA.Extra, edns0)

	msgA, err := queryA.Pack()
	if err != nil {
		writeErrorResponse(w, fmt.Sprintf("打包A记录DNS查询失败: %v", err), http.StatusInternalServerError)
		return
	}

	msgAAAA, err := queryAAAA.Pack()
	if err != nil {
		writeErrorResponse(w, fmt.Sprintf("打包AAAA记录DNS查询失败: %v", err), http.StatusInternalServerError)
		return
	}

	b64A := base64.RawURLEncoding.EncodeToString(msgA)
	b64AAAA := base64.RawURLEncoding.EncodeToString(msgAAAA)
	ctx, cancel := context.WithTimeout(r.Context(), 5*time.Second)
	defer cancel()

	resultsA := make(chan *http.Response, len(config.DNSServers))
	resultsAAAA := make(chan *http.Response, len(config.DNSServers))
	var wg sync.WaitGroup

	for _, server := range config.DNSServers {
		wg.Add(2)
		go func(server string) {
			defer wg.Done()
			resp, err := sendDNSQuery(ctx, server, b64A, 3)
			if err != nil {
				log.Printf("A记录DNS查询失败: %v", err)
				return
			}
			resultsA <- resp
		}(server)

		go func(server string) {
			defer wg.Done()
			resp, err := sendDNSQuery(ctx, server, b64AAAA, 3)
			if err != nil {
				log.Printf("AAAA记录DNS查询失败: %v", err)
				return
			}
			resultsAAAA <- resp
		}(server)
	}

	go func() {
		wg.Wait()
		close(resultsA)
		close(resultsAAAA)
	}()

	var responseA, responseAAAA *dns.Msg
	var cnameRecords []dns.CNAME

	// 处理 A 记录
	select {
	case resp, ok := <-resultsA:
		if ok {
			bodyBytes, err := io.ReadAll(resp.Body)
			if err != nil {
				resp.Body.Close()
				writeErrorResponse(w, fmt.Sprintf("读取A记录响应体失败: %v", err), http.StatusInternalServerError)
				return
			}
			defer resp.Body.Close()

			responseA = &dns.Msg{}
			err = responseA.Unpack(bodyBytes)
			if err != nil {
				writeErrorResponse(w, fmt.Sprintf("解包A记录DNS响应失败: %v", err), http.StatusInternalServerError)
				return
			}

			// 从 A 响应中收集 CNAME 记录
			for _, ans := range responseA.Answer {
				if cname, ok := ans.(*dns.CNAME); ok {
					cnameRecords = append(cnameRecords, *cname)
				}
			}
		}
	case <-ctx.Done():
		writeErrorResponse(w, "A记录查询请求超时或被取消", http.StatusGatewayTimeout)
		return
	}

	// 处理 AAAA 记录
	select {
	case resp, ok := <-resultsAAAA:
		if ok {
			bodyBytes, err := io.ReadAll(resp.Body)
			if err != nil {
				resp.Body.Close()
				writeErrorResponse(w, fmt.Sprintf("读取AAAA记录响应体失败: %v", err), http.StatusInternalServerError)
				return
			}
			defer resp.Body.Close()

			responseAAAA = &dns.Msg{}
			err = responseAAAA.Unpack(bodyBytes)
			if err != nil {
				writeErrorResponse(w, fmt.Sprintf("解包AAAA记录DNS响应失败: %v", err), http.StatusInternalServerError)
				return
			}
		}
	case <-ctx.Done():
		writeErrorResponse(w, "AAAA记录查询请求超时或被取消", http.StatusGatewayTimeout)
		return
	}

	// 如果找到 CNAME 记录，则处理 CNAME 查询
	if len(cnameRecords) > 0 {
		for _, cname := range cnameRecords {
			cnameQueryA := dns.Msg{}
			cnameQueryA.SetQuestion(cname.Target, dns.TypeA)

			cnameQueryAAAA := dns.Msg{}
			cnameQueryAAAA.SetQuestion(cname.Target, dns.TypeAAAA)

			// 为 A 类型的查询添加 EDNS0 选项
			edns0Subnet := getEDNS0SubnetOption(clientIP)
			edns0 := &dns.OPT{Hdr: dns.RR_Header{Name: ".", Rrtype: dns.TypeOPT}, Option: []dns.EDNS0{edns0Subnet}}
			cnameQueryA.Extra = append(cnameQueryA.Extra, edns0)
			cnameQueryAAAA.Extra = append(cnameQueryAAAA.Extra, edns0)

			cnameQueryABytes, err := cnameQueryA.Pack()
			if err != nil {
				log.Printf("打包 CNAME 查询 A 记录失败: %v", err)
				continue
			}

			cnameQueryAAAABytes, err := cnameQueryAAAA.Pack()
			if err != nil {
				log.Printf("打包 CNAME 查询 AAAA 记录失败: %v", err)
				continue
			}

			var wgCNAME sync.WaitGroup
			resultsCNAMEA := make(chan *http.Response, len(config.DNSServers))
			resultsCNAMEAAAA := make(chan *http.Response, len(config.DNSServers))
			wgCNAME.Add(2)
			go func() {
				defer wgCNAME.Done()
				for _, server := range config.DNSServers {
					resp, err := sendDNSQuery(ctx, server, base64.RawURLEncoding.EncodeToString(cnameQueryABytes), 3)
					if err != nil {
						log.Printf("CNAME记录A查询失败: %v", err)
						continue
					}
					resultsCNAMEA <- resp
				}
			}()

			go func() {
				defer wgCNAME.Done()
				for _, server := range config.DNSServers {
					resp, err := sendDNSQuery(ctx, server, base64.RawURLEncoding.EncodeToString(cnameQueryAAAABytes), 3)
					if err != nil {
						log.Printf("CNAME记录AAAA查询失败: %v", err)
						continue
					}
					resultsCNAMEAAAA <- resp
				}
			}()

			wgCNAME.Wait()
			close(resultsCNAMEA)
			close(resultsCNAMEAAAA)

			// 从 CNAME 响应中收集 A 和 AAAA 记录
			for resp := range resultsCNAMEA {
				bodyBytes, err := io.ReadAll(resp.Body)
				resp.Body.Close()
				if err != nil {
					log.Printf("读取 CNAME A 记录响应体失败: %v", err)
					continue
				}
				if err := responseA.Unpack(bodyBytes); err != nil {
					log.Printf("解包 CNAME A 记录响应失败: %v", err)
					continue
				}
				// 将结果添加到 cnameRecords
				for _, ans := range responseA.Answer {
					if cname, ok := ans.(*dns.CNAME); ok {
						cnameRecords = append(cnameRecords, *cname)
					}
					// 收集 A 记录
					if a, ok := ans.(*dns.A); ok {
						responseA.Answer = append(responseA.Answer, a)
					}
				}
			}

			for resp := range resultsCNAMEAAAA {
				bodyBytes, err := io.ReadAll(resp.Body)
				resp.Body.Close()
				if err != nil {
					log.Printf("读取 CNAME AAAA 记录响应体失败: %v", err)
					continue
				}
				if err := responseAAAA.Unpack(bodyBytes); err != nil {
					log.Printf("解包 CNAME AAAA 记录响应失败: %v", err)
					continue
				}
				// 收集 AAAA 记录
				for _, ans := range responseAAAA.Answer {
					if aaaa, ok := ans.(*dns.AAAA); ok {
						responseAAAA.Answer = append(responseAAAA.Answer, aaaa)
					}
				}
			}
		}
	}

	// 准备最终的 DNS 响应
	if net.ParseIP(clientIP).To4() != nil {
		mergedResponse := &dns.Msg{}
		if responseA != nil {
			mergedResponse.Answer = append(mergedResponse.Answer, responseA.Answer...)
		}
		// 去重处理
		mergedResponse.Answer = removeDuplicates(mergedResponse.Answer)
		jsonResponse := convertDnsMsgToJSON(*mergedResponse)
		jsonOutput, err := json.MarshalIndent(jsonResponse, "", "  ")
		if err != nil {
			writeErrorResponse(w, fmt.Sprintf("序列化JSON失败: %v", err), http.StatusInternalServerError)
			return
		}
		logDNSQuery(name, config.DNSServers, *mergedResponse, clientIP)
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		w.Write(jsonOutput)
	} else {
		// 合并A和AAAA的结果
		mergedResponse := &dns.Msg{}
		if responseA != nil {
			mergedResponse.Answer = append(mergedResponse.Answer, responseA.Answer...)
		}
		if responseAAAA != nil {
			mergedResponse.Answer = append(mergedResponse.Answer, responseAAAA.Answer...)
		}

		// 去重处理
		mergedResponse.Answer = removeDuplicates(mergedResponse.Answer)

		jsonResponse := convertDnsMsgToJSON(*mergedResponse)
		jsonOutput, err := json.MarshalIndent(jsonResponse, "", "  ")
		if err != nil {
			writeErrorResponse(w, fmt.Sprintf("序列化JSON失败: %v", err), http.StatusInternalServerError)
			return
		}
		logDNSQuery(name, config.DNSServers, *mergedResponse, clientIP)
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		w.Write(jsonOutput)
	}
}

func dnsQueryHandler(w http.ResponseWriter, r *http.Request) {
	var dnsQuery []byte
	var err error

	// 创建一个带超时的 context
	ctx, cancel := context.WithTimeout(r.Context(), 5*time.Second)
	defer cancel()

	switch r.Method {
	case http.MethodGet:
		dnsParam := r.URL.Query().Get("dns")
		if dnsParam == "" {
			writeErrorResponse(w, "缺少 'dns' 查询参数", http.StatusBadRequest)
			return
		}

		dnsQuery, err = base64.RawURLEncoding.DecodeString(dnsParam)
		if err != nil {
			writeErrorResponse(w, "解码 'dns' 参数失败", http.StatusBadRequest)
			return
		}

	case http.MethodPost:
		dnsQuery, err = io.ReadAll(r.Body)
		if err != nil {
			writeErrorResponse(w, "读取请求体失败", http.StatusBadRequest)
			return
		}

	default:
		writeErrorResponse(w, "方法不被允许", http.StatusMethodNotAllowed)
		return
	}

	msg := new(dns.Msg)
	err = msg.Unpack(dnsQuery)
	if err != nil {
		writeErrorResponse(w, "解析 DNS 查询失败", http.StatusBadRequest)
		return
	}

	clientIP := getClientIP(r)
	query := dns.Msg{}
	query.SetQuestion(dns.Fqdn(msg.Question[0].Name), msg.Question[0].Qtype)
	query.Id = msg.Id

	// 仅在查询类型为 A 或 AAAA 时添加 EDNS0 选项
	if msg.Question[0].Qtype == dns.TypeA || msg.Question[0].Qtype == dns.TypeAAAA {
		edns0Subnet := getEDNS0SubnetOption(clientIP)
		edns0 := &dns.OPT{Hdr: dns.RR_Header{Name: ".", Rrtype: dns.TypeOPT}, Option: []dns.EDNS0{edns0Subnet}}
		query.Extra = append(query.Extra, edns0)
	}

	queryBytes, err := query.Pack()
	if err != nil {
		log.Printf("打包查询失败: %v\n", err)
		writeErrorResponse(w, fmt.Sprintf("打包查询失败: %v", err), http.StatusInternalServerError)
		return
	}

	// Helper function to send DNS queries
	sendQuery := func(queryBytes []byte, wg *sync.WaitGroup, results chan<- *http.Response) {
		defer wg.Done()
		for _, server := range config.DNSServers {
			resp, err := sendDNSQuery(ctx, server, base64.RawURLEncoding.EncodeToString(queryBytes), 3)
			if err != nil {
				log.Printf("%s 记录 DNS 查询失败: %v", dns.TypeToString[msg.Question[0].Qtype], err)
				continue
			}
			results <- resp
		}
	}

	var wg sync.WaitGroup
	results := make(chan *http.Response, len(config.DNSServers))
	wg.Add(1)
	go sendQuery(queryBytes, &wg, results)

	wg.Wait()
	close(results)

	var responseMsg dns.Msg
	var cnameRecords []dns.CNAME
	foundARecord := false

	for resp := range results {
		bodyBytes, err := io.ReadAll(resp.Body)
		resp.Body.Close()
		if err != nil {
			log.Printf("读取 %s 记录响应体失败: %v", dns.TypeToString[msg.Question[0].Qtype], err)
			continue
		}
		if err := responseMsg.Unpack(bodyBytes); err != nil {
			log.Printf("解包 %s 记录响应失败: %v", dns.TypeToString[msg.Question[0].Qtype], err)
			continue
		}

		// 处理 CNAME 记录
		for _, ans := range responseMsg.Answer {
			if cname, ok := ans.(*dns.CNAME); ok {
				cnameRecords = append(cnameRecords, *cname)
			}
		}

		// 检查是否有 A 或 AAAA 记录
		for _, ans := range responseMsg.Answer {
			if ans.Header().Rrtype == dns.TypeA || ans.Header().Rrtype == dns.TypeAAAA {
				foundARecord = true
				break
			}
		}

		if foundARecord {
			break
		}
	}

	// 如果找到 CNAME 记录，则进行新的查询
	if len(cnameRecords) > 0 && !foundARecord {
		for _, cname := range cnameRecords {
			cnameQuery := dns.Msg{}
			cnameQuery.SetQuestion(cname.Target, msg.Question[0].Qtype)
			cnameQuery.Id = msg.Id

			// Add EDNS0 option if the query type is A or AAAA
			if msg.Question[0].Qtype == dns.TypeA || msg.Question[0].Qtype == dns.TypeAAAA {
				edns0Subnet := getEDNS0SubnetOption(clientIP)
				edns0 := &dns.OPT{Hdr: dns.RR_Header{Name: ".", Rrtype: dns.TypeOPT}, Option: []dns.EDNS0{edns0Subnet}}
				cnameQuery.Extra = append(cnameQuery.Extra, edns0)
			}

			cnameQueryBytes, err := cnameQuery.Pack()
			if err != nil {
				log.Printf("打包 CNAME 查询失败: %v\n", err)
				continue
			}

			var wgCNAME sync.WaitGroup
			resultsCNAME := make(chan *http.Response, len(config.DNSServers))
			wgCNAME.Add(1)
			go sendQuery(cnameQueryBytes, &wgCNAME, resultsCNAME)

			wgCNAME.Wait()
			close(resultsCNAME)

			for resp := range resultsCNAME {
				bodyBytes, err := io.ReadAll(resp.Body)
				resp.Body.Close()
				if err != nil {
					log.Printf("读取 %s 记录响应体失败: %v", dns.TypeToString[msg.Question[0].Qtype], err)
					continue
				}
				if err := responseMsg.Unpack(bodyBytes); err != nil {
					log.Printf("解包 %s 记录响应失败: %v", dns.TypeToString[msg.Question[0].Qtype], err)
					continue
				}

				// 检查是否有 A 或 AAAA 记录
				for _, ans := range responseMsg.Answer {
					if ans.Header().Rrtype == dns.TypeA || ans.Header().Rrtype == dns.TypeAAAA {
						foundARecord = true
						break
					}
				}

				if foundARecord {
					break
				}
			}

			if foundARecord {
				break
			}
		}
	}

	// 返回响应
	logDNSQuery(msg.Question[0].Name, config.DNSServers, responseMsg, clientIP)
	responseBytes, err := responseMsg.Pack()
	if err != nil {
		writeErrorResponse(w, fmt.Sprintf("打包 %s 记录响应失败: %v", dns.TypeToString[msg.Question[0].Qtype], err), http.StatusInternalServerError)
		return
	}
	w.Header().Set("Content-Type", "application/dns-message")
	w.WriteHeader(http.StatusOK)
	w.Write(responseBytes)
}

func getEDNS0SubnetOption(ip string) *dns.EDNS0_SUBNET {
	subnet := &dns.EDNS0_SUBNET{
		Code:        dns.EDNS0SUBNET,
		SourceScope: 0,
		Address:     net.ParseIP(ip),
	}
	if subnet.Address.To4() != nil {
		subnet.Family = 1
		subnet.SourceNetmask = 32
	} else if subnet.Address.To16() != nil {
		subnet.Family = 2
		subnet.SourceNetmask = 128
	}
	return subnet
}

// writeErrorResponse 使用setJSONResponse来发送包含错误信息的JSON响应。
func headerWritten(w http.ResponseWriter) bool {
	h := w.Header()
	if len(h) == 0 {
		return false
	}
	for _, v := range h {
		if len(v) > 0 {
			return true
		}
	}
	return false
}

func setJSONResponse(w http.ResponseWriter, statusCode int, data interface{}) {
	// 设置响应头
	if w.Header().Get("Content-Type") == "" {
		w.Header().Set("Content-Type", "application/json")
	}

	// 如果响应头尚未写入，则写入状态码
	if !headerWritten(w) {
		w.WriteHeader(statusCode)
	}

	// 尝试编码数据为 JSON 并写入响应
	err := json.NewEncoder(w).Encode(data)
	if err != nil {
		// logger.Printf("Failed to encode JSON response: %v", err)
		// 如果编码失败，写入错误信息
		if !headerWritten(w) {
			http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		}
	}
}

func writeErrorResponse(w http.ResponseWriter, message string, statusCode int) {
	errorResponse := map[string]string{"error": message}
	setJSONResponse(w, statusCode, errorResponse)
}

// 获取客户端 IP 地址
func getClientIP(r *http.Request) string {
	if ip := r.Header.Get("X-Forwarded-For"); ip != "" {
		ips := strings.Split(ip, ",")
		if len(ips) > 0 {
			return strings.TrimSpace(ips[0])
		}
	}
	if ip := r.Header.Get("X-Real-Ip"); ip != "" {
		return ip
	}
	if ip, _, err := net.SplitHostPort(r.RemoteAddr); err == nil {
		return ip
	}
	// 如果 SplitHostPort 失败，可能是因为 RemoteAddr 不包含端口号，直接返回
	return r.RemoteAddr
}

func sendDNSQuery(ctx context.Context, dnsServer, b64 string, retries int) (*http.Response, error) {
	var resp *http.Response
	var err error

	for attempt := 0; attempt < retries; attempt++ {
		resp, err = executeDNSQuery(ctx, dnsServer, b64)
		if err == nil {
			return resp, nil
		}
		// logger.Printf("DNS 查询失败 (尝试 %d/%d): %v\n", attempt+1, retries, err)
		time.Sleep(10 * time.Second) // 等待一段时间后重试
	}

	return nil, fmt.Errorf("所有重试均失败: %v", err)
}

func executeDNSQuery(ctx context.Context, dnsServer, b64 string) (*http.Response, error) {
	parsedURL, err := url.Parse(dnsServer)
	if err != nil {
		return nil, err
	}

	switch parsedURL.Scheme {
	case "https":
		return sendDoHQuery(ctx, dnsServer, b64)
	case "tls":
		return sendDoTQuery(ctx, dnsServer, b64)
	case "tcp":
		return sendTCPQuery(ctx, dnsServer, b64)
	case "udp":
		return sendUDPQuery(ctx, dnsServer, b64)
	case "quic":
		return sendDoQQuery(ctx, dnsServer, b64)
	case "h3":
		return sendH3Query(ctx, dnsServer, b64)
	case "sdns":
		return sendCryptQuery(ctx, dnsServer, b64)
	default:
		return nil, fmt.Errorf("不支持的协议: %s", parsedURL.Scheme)
	}
}

func sendDoHQuery(ctx context.Context, dnsServer, b64 string) (*http.Response, error) {
	// 公共的传输配置
	tlsConfig := &tls.Config{InsecureSkipVerify: true}
	timeout := 10 * time.Second
	expectContinueTimeout := 1 * time.Second
	maxIdleConns := 200
	maxIdleConnsPerHost := 10

	// HTTP/3 客户端
	http3Transport := &http3.Transport{
		TLSClientConfig: tlsConfig,
	}
	client := &http.Client{
		Timeout:   timeout,
		Transport: http3Transport,
	}

	// 构建请求 URL
	reqURL := dnsServer + "?dns=" + b64
	req, err := http.NewRequestWithContext(ctx, "GET", reqURL, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Accept", "application/dns-message")

	// 尝试发送 HTTP/3 请求
	resp, err := client.Do(req)
	if err == nil && resp.StatusCode == http.StatusOK {
		return resp, nil
	}

	// HTTP/2 Transport 配置
	http2Transport := &http.Transport{
		TLSClientConfig:       tlsConfig,
		TLSHandshakeTimeout:   timeout,
		ExpectContinueTimeout: expectContinueTimeout,
		MaxIdleConns:          maxIdleConns,
		MaxIdleConnsPerHost:   maxIdleConnsPerHost,
	}
	if err := http2.ConfigureTransport(http2Transport); err != nil {
		return nil, fmt.Errorf("无法配置 HTTP/2: %v", err)
	}
	client.Transport = http2Transport

	resp, err = client.Do(req)
	if err == nil && resp.StatusCode == http.StatusOK {
		return resp, nil
	}

	// HTTP/1.1 Transport 配置
	http1Transport := &http.Transport{
		TLSClientConfig:       tlsConfig,
		TLSHandshakeTimeout:   timeout,
		ExpectContinueTimeout: expectContinueTimeout,
		MaxIdleConns:          maxIdleConns,
		MaxIdleConnsPerHost:   maxIdleConnsPerHost,
		ForceAttemptHTTP2:     false, // 禁用 HTTP/2 以使用 HTTP/1.1
	}
	client.Transport = http1Transport

	resp, err = client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("DoH查询失败: %v", err)
	}
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("DoH查询失败, 状态码: %d", resp.StatusCode)
	}

	return resp, nil
}

func sendDoTQuery(ctx context.Context, dnsServer, b64 string) (*http.Response, error) {
	parsedURL, err := url.Parse(dnsServer)
	if err != nil {
		return nil, err
	}

	host := parsedURL.Host
	if !strings.Contains(host, ":") {
		host += ":853" // 默认DoT端口
	}

	// 使用 DialContext 进行连接
	dialer := &net.Dialer{}
	conn, err := dialer.DialContext(ctx, "tcp", host)
	if err != nil {
		return nil, err
	}
	defer conn.Close()

	// 创建 TLS 客户端连接
	tlsConn := tls.Client(conn, &tls.Config{
		InsecureSkipVerify: true,
	})
	defer tlsConn.Close()

	query, err := base64.RawURLEncoding.DecodeString(b64)
	if err != nil {
		return nil, err
	}

	lengthField := make([]byte, 2)
	lengthField[0] = byte(len(query) >> 8)
	lengthField[1] = byte(len(query) & 0xff)

	_, err = tlsConn.Write(lengthField)
	if err != nil {
		return nil, err
	}
	_, err = tlsConn.Write(query)
	if err != nil {
		return nil, err
	}

	respLength := make([]byte, 2)
	_, err = io.ReadFull(tlsConn, respLength)
	if err != nil {
		return nil, err
	}
	respLen := int(respLength[0])<<8 | int(respLength[1])

	resp := make([]byte, respLen)
	_, err = io.ReadFull(tlsConn, resp)
	if err != nil {
		return nil, err
	}

	tempResp := &http.Response{
		StatusCode: http.StatusOK,
		Body:       io.NopCloser(strings.NewReader(string(resp))),
	}

	return tempResp, nil
}

func sendTCPQuery(ctx context.Context, dnsServer, b64 string) (*http.Response, error) {
	parsedURL, err := url.Parse(dnsServer)
	if err != nil {
		return nil, err
	}

	host := parsedURL.Host
	if !strings.Contains(host, ":") {
		host += ":53" // 默认TCP端口
	}

	dialer := &net.Dialer{}
	conn, err := dialer.DialContext(ctx, "tcp", host)
	if err != nil {
		return nil, err
	}
	defer conn.Close()

	query, err := base64.RawURLEncoding.DecodeString(b64)
	if err != nil {
		return nil, err
	}

	lengthField := make([]byte, 2)
	lengthField[0] = byte(len(query) >> 8)
	lengthField[1] = byte(len(query) & 0xff)

	_, err = conn.Write(lengthField)
	if err != nil {
		return nil, err
	}
	_, err = conn.Write(query)
	if err != nil {
		return nil, err
	}

	respLength := make([]byte, 2)
	_, err = io.ReadFull(conn, respLength)
	if err != nil {
		return nil, err
	}
	respLen := int(respLength[0])<<8 | int(respLength[1])

	resp := make([]byte, respLen)
	_, err = io.ReadFull(conn, resp)
	if err != nil {
		return nil, err
	}

	tempResp := &http.Response{
		StatusCode: http.StatusOK,
		Body:       io.NopCloser(strings.NewReader(string(resp))),
	}
	return tempResp, nil
}

func sendUDPQuery(ctx context.Context, dnsServer, b64 string) (*http.Response, error) {
	parsedURL, err := url.Parse(dnsServer)
	if err != nil {
		return nil, err
	}

	host := parsedURL.Host
	if !strings.Contains(host, ":") {
		host += ":53" // 默认UDP端口
	}

	// 设置 UDP 连接的上下文超时
	dialer := &net.Dialer{}
	conn, err := dialer.DialContext(ctx, "udp", host)
	if err != nil {
		return nil, err
	}
	defer conn.Close()

	// 解码 base64 查询内容
	query, err := base64.RawURLEncoding.DecodeString(b64)
	if err != nil {
		return nil, err
	}

	// 发送查询请求
	_, err = conn.Write(query)
	if err != nil {
		return nil, err
	}

	// 设置读取超时，防止长时间阻塞
	conn.SetReadDeadline(time.Now().Add(2 * time.Second))

	resp := make([]byte, 512) // DNS 响应通常小于 512 字节
	n, err := conn.Read(resp)
	if err != nil {
		return nil, err
	}

	// 返回 http.Response 对象
	tempResp := &http.Response{
		StatusCode: http.StatusOK,
		Body:       io.NopCloser(bytes.NewReader(resp[:n])),
	}

	return tempResp, nil
}

// QUIC 是用于通过 QUIC 协议进行 DNS 查询的结构体
type QUIC struct {
	Server    string
	TLSConfig *tls.Config
	PMTUD     bool
	ReuseConn bool
	conn      quic.Connection // 接口类型，而不是指针
}

// setServerName 设置 TLS 配置中的服务器名称
func (q *QUIC) setServerName() {
	host, _, err := net.SplitHostPort(q.Server)
	if err != nil {
		fmt.Printf("无效的 QUIC 服务器地址: %s", err)
	}
	q.TLSConfig.ServerName = host
}

// sendDoQQuery 将 base64 编码的 DNS 查询发送到 QUIC 服务器并返回响应
func sendDoQQuery(ctx context.Context, dnsServer, b64 string) (*http.Response, error) {
	parsedURL, err := url.Parse(dnsServer)
	if err != nil {
		return nil, err
	}

	host := parsedURL.Host
	if !strings.Contains(host, ":") {
		host += ":853" // 默认 QUIC 端口
	}

	// 解码 base64 查询内容
	query, err := base64.RawURLEncoding.DecodeString(b64)
	if err != nil {
		return nil, err
	}

	// 创建 QUIC 连接和 TLS 配置
	tlsConfig := &tls.Config{
		ServerName:             host,
		InsecureSkipVerify:     true, // 测试时可以使用 true，生产时应关闭
		NextProtos:             []string{"doq"},
		SessionTicketsDisabled: false,
	}

	quicConfig := &quic.Config{
		DisablePathMTUDiscovery: true,             // 默认启用 PMTU
		KeepAlivePeriod:         time.Second * 20, // 默认值为 30s
		TokenStore:              newQUICTokenStore(),
	}

	// 设置 DNS over QUIC 客户端
	quicClient := &QUIC{
		Server:    host,
		TLSConfig: tlsConfig,
	}

	// 建立 QUIC 连接
	if quicClient.conn == nil || !quicClient.ReuseConn {
		quicClient.setServerName()

		conn, err := quic.DialAddr(
			ctx,
			host,
			tlsConfig,
			quicConfig,
		)
		if err != nil {
			return nil, fmt.Errorf("打开到 %s 的 QUIC 会话失败: %v", host, err)
		}
		quicClient.conn = conn
	}

	// 打开 QUIC 流并发送查询
	stream, err := quicClient.conn.OpenStream()
	if err != nil {
		return nil, fmt.Errorf("打开 QUIC 流失败: %v", err)
	}

	// 发送查询
	_, err = stream.Write(addPrefix(query))
	if err != nil {
		return nil, fmt.Errorf("无法写入 QUIC 流: %w", err)
	}

	// 关闭流并读取响应
	_ = stream.Close()

	respBuf, err := io.ReadAll(stream)
	if err != nil {
		return nil, fmt.Errorf("从 QUIC 流读取响应失败: %v", err)
	}
	if len(respBuf) == 0 {
		return nil, fmt.Errorf("从 %s 收到空响应", host)
	}

	// 解包 DNS 响应
	m := new(dns.Msg)
	err = m.Unpack(respBuf[2:]) // 跳过前 2 个字节的长度前缀
	if err != nil {
		return nil, fmt.Errorf("DNS 响应解包失败: %w", err)
	}

	// 重新打包 DNS 响应
	packedResp, err := m.Pack()
	if err != nil {
		return nil, fmt.Errorf("DNS 响应打包失败: %w", err)
	}

	// 构造 HTTP 响应
	tempResp := &http.Response{
		StatusCode: http.StatusOK,
		Body:       io.NopCloser(bytes.NewReader(packedResp)),
	}

	return tempResp, nil
}

func addPrefix(b []byte) (m []byte) {
	m = make([]byte, 2+len(b))
	binary.BigEndian.PutUint16(m, uint16(len(b)))
	copy(m[2:], b)

	return m
}

func newQUICTokenStore() (s quic.TokenStore) {
	return quic.NewLRUTokenStore(1, 10)
}

func sendH3Query(ctx context.Context, dnsServer, b64 string) (*http.Response, error) {
	tlsConfig := &tls.Config{InsecureSkipVerify: true}
	timeout := 10 * time.Second

	if strings.HasPrefix(dnsServer, "h3://") {
		dnsServer = "https://" + dnsServer[5:]
	}

	// HTTP/3 客户端
	http3Transport := &http3.Transport{
		TLSClientConfig: tlsConfig,
	}
	client := &http.Client{
		Timeout:   timeout,
		Transport: http3Transport,
	}

	// 构建请求 URL
	reqURL := dnsServer + "?dns=" + b64
	req, err := http.NewRequestWithContext(ctx, "GET", reqURL, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Accept", "application/dns-message")

	// 尝试发送 HTTP/3 请求
	resp, err := client.Do(req)
	if err == nil && resp.StatusCode == http.StatusOK {
		// return resp, nil
	}
	return resp, nil
}

func sendCryptQuery(ctx context.Context, dnsServer, b64 string) (*http.Response, error) {
	// 解码 base64 编码的 DNS 查询
	queryBytes, err := base64.RawURLEncoding.DecodeString(b64)
	if err != nil {
		return nil, fmt.Errorf("base64 解码查询失败: %w", err)
	}

	// 初始化 DNSCrypt 客户端
	client := &dnscrypt.Client{
		Timeout: 5 * time.Second, // 设置超时时间为 5 秒
		Net:     "udp",           // 默认使用 UDP
	}

	// 解析 DNSCrypt stamp
	stamp, err := dnsstamps.NewServerStampFromString(dnsServer)
	if err != nil {
		return nil, fmt.Errorf("解析 DNSCrypt stamp 失败: %w", err)
	}

	// 定义函数用于切换协议
	dialWithFallback := func() (*dnscrypt.ResolverInfo, error) {
		// 尝试 UDP 连接
		resolver, err := client.Dial(stamp.String())
		if err == nil {
			return resolver, nil
		}
		// 切换到 TCP
		client.Net = "tcp"
		resolver, err = client.Dial(stamp.String())
		if err == nil {
			return resolver, nil
		}
		return nil, fmt.Errorf("UDP 和 TCP 均连接失败: %w", err)
	}

	// 连接 DNSCrypt 服务器，自动切换 UDP -> TCP
	resolver, err := dialWithFallback()
	if err != nil {
		return nil, fmt.Errorf("连接 DNSCrypt 服务器失败（UDP 和 TCP 均失败）: %w", err)
	}

	// 将 queryBytes 转换为 *dns.Msg
	msg := &dns.Msg{}
	if err = msg.Unpack(queryBytes); err != nil {
		return nil, fmt.Errorf("解包 DNS 查询失败: %w", err)
	}

	// 使用 DNSCrypt 客户端发送查询
	respMsg, err := client.Exchange(msg, resolver)
	if err != nil {
		return nil, fmt.Errorf("DNSCrypt 查询失败: %w", err)
	}

	// 使用 ctx 超时控制，保证在查询响应过程中不中断
	select {
	case <-ctx.Done():
		// 如果 ctx 被取消或超时，返回错误
		return nil, fmt.Errorf("请求被取消或超时: %w", ctx.Err())
	default:
		// 否则继续处理响应
	}

	// 将响应转换为字节切片
	packedResp, err := respMsg.Pack()
	if err != nil {
		return nil, fmt.Errorf("打包 DNS 响应失败: %w", err)
	}

	// 构造 HTTP 响应并返回
	tempResp := &http.Response{
		StatusCode: http.StatusOK,
		Body:       io.NopCloser(bytes.NewReader(packedResp)),
	}

	return tempResp, nil
}

func convertDnsMsgToJSON(msg dns.Msg) map[string]interface{} {
	// 初始化返回的JSON结构
	result := map[string]interface{}{
		"Status":   msg.Rcode,                  // DNS响应码
		"TC":       msg.Truncated,              // 是否被截断
		"RD":       msg.RecursionDesired,       // 是否需要递归
		"RA":       msg.RecursionAvailable,     // 服务器是否支持递归
		"AD":       msg.AuthenticatedData,      // 数据是否被验证
		"CD":       msg.CheckingDisabled,       // 是否禁用检查
		"Question": map[string]interface{}{},   // 问题部分
		"Answer":   []map[string]interface{}{}, // 回答部分
	}

	// 处理问题部分，如果存在的话
	if len(msg.Question) > 0 {
		q := msg.Question[0]
		result["Question"] = map[string]interface{}{
			"name": q.Name,  // 域名
			"type": q.Qtype, // 查询类型
		}
	}

	// 遍历回答部分，处理每一条记录
	for _, a := range msg.Answer {
		var data string
		switch a.Header().Rrtype {
		case dns.TypeA:
			data = a.(*dns.A).A.String() // 处理A记录
		case dns.TypeAAAA:
			data = a.(*dns.AAAA).AAAA.String() // 处理AAAA记录
		case dns.TypeCNAME:
			data = a.(*dns.CNAME).Target // 处理CNAME记录
		default:
			data = a.String() // 其他记录类型
		}

		// 处理多余的`\t`字符
		dataParts := strings.Split(data, "\t")
		if len(dataParts) > 1 {
			data = dataParts[len(dataParts)-1]
		}

		// 添加记录到结果中
		result["Answer"] = append(result["Answer"].([]map[string]interface{}), map[string]interface{}{
			"name": a.Header().Name,   // 记录名
			"TTL":  a.Header().Ttl,    // 生存时间
			"type": a.Header().Rrtype, // 记录类型
			"data": data,              // 记录内容
		})
	}

	return result // 返回结果
}

func logDNSQuery(name string, servers []string, response dns.Msg, clientIP string) {
	if config.LogConfig.LogToFile {
		if config.LogConfig.LogFile != "" {
			logEntry := fmt.Sprintf("查询名称: %s, 解析服务器: %v, 客户端IP: %s, 响应: %v", name, servers, clientIP, response)
			select {
			case logChannel <- logEntry:
			default:
				fmt.Println("日志通道已满，无法记录日志条目")
			}
		} else {
			logEntry := fmt.Sprintf("查询名称: %s, 解析服务器: %v, 客户端IP: %s, 响应: %v", name, servers, clientIP, response)
			select {
			case logChannel <- logEntry:
			default:
				fmt.Println("日志通道已满，无法记录日志条目")
			}
		}
	} else {
		log.SetOutput(io.Discard)
	}
}

func loadConfig(file string) {
	yamlFile, err := os.ReadFile(file)
	if err != nil {
		fmt.Printf("读取YAML文件失败: %v\n", err)
		os.Exit(1)
	}

	err = yaml.Unmarshal(yamlFile, &config)
	if err != nil {
		fmt.Printf("解析YAML失败: %v\n", err)
		os.Exit(1)
	}
}

func initLogger() {
	logChannel = make(chan string, 10000)
	if config.LogConfig.LogToFile {
		if config.LogConfig.LogFile == "" {
			logger = log.New(os.Stdout, "", log.LstdFlags)
		} else {
			logPath := config.LogConfig.LogFile
			logRotate, err := rotatelogs.New(
				logPath+".%Y%m%d%H%M",
				rotatelogs.WithMaxAge(time.Duration(config.LogConfig.MaxAgeDays)*24*time.Hour),
				rotatelogs.WithRotationTime(time.Duration(config.LogConfig.RotationDays)*24*time.Hour),
			)
			if err != nil {
				fmt.Printf("创建日志切割器失败: %v\n", err)
				logger = log.New(os.Stdout, "", log.LstdFlags)
			} else {
				logger = log.New(logRotate, "", log.LstdFlags)
			}
		}
	} else {
		// 如果 LogToFile 是 false，禁用日志记录
		logger = log.New(io.Discard, "", log.LstdFlags)
	}
}

func logProcessor() {
	defer wg.Done()

	for logEntry := range logChannel {
		logger.Println(logEntry)
	}
}