package main

import (
	"context"
	"crypto/rand"
	"embed"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/fs"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"sort"
	"strings"
	"sync"
	"time"
)

//go:embed web/*
var webFS embed.FS

type resolver struct{ Name, Country, City, Region, URL string }
type record struct {
	Value string `json:"value"`
	TTL   int    `json:"ttl"`
}
type result struct {
	Resolver     string   `json:"resolver"`
	Country      string   `json:"country"`
	City         string   `json:"city"`
	Region       string   `json:"region"`
	Status       string   `json:"status"`
	Records      []record `json:"records"`
	ResponseTime int64    `json:"responseTime"`
	Error        string   `json:"error,omitempty"`
}
type query struct {
	ID, Domain, Type string
	Created          time.Time
	Results          chan result
}
type app struct {
	mu      sync.RWMutex
	queries map[string]*query
	client  *http.Client
}

var resolvers = []resolver{
	{"Cloudflare", "全球", "Anycast", "全球", "https://cloudflare-dns.com/dns-query"},
	{"Google", "全球", "Anycast", "全球", "https://dns.google/resolve"},
	{"Quad9", "瑞士", "Zurich", "欧洲", "https://dns.quad9.net/dns-query"},
	{"AdGuard", "塞浦路斯", "Limassol", "欧洲", "https://dns.adguard-dns.com/resolve"},
	{"AliDNS", "中国", "Hangzhou", "亚洲", "https://dns.alidns.com/resolve"},
	{"DNSPod", "中国", "Shanghai", "亚洲", "https://doh.pub/dns-query"},
	{"360 Secure DNS", "中国", "Beijing", "亚洲", "https://doh.360.cn/dns-query"},
	{"Tencent EdgeOne", "新加坡", "Singapore", "亚洲", "https://public.dns.iij.jp/dns-query"},
	{"IIJ", "日本", "Tokyo", "亚洲", "https://public.dns.iij.jp/dns-query"},
	{"Mullvad", "瑞典", "Gothenburg", "欧洲", "https://doh.mullvad.net/dns-query"},
	{"Control D", "加拿大", "Toronto", "北美", "https://freedns.controld.com/p0"},
	{"CleanBrowsing", "美国", "Virginia", "北美", "https://doh.cleanbrowsing.org/doh/security-filter/"},
	{"OpenDNS", "美国", "San Francisco", "北美", "https://doh.opendns.com/dns-query"},
	{"NextDNS", "美国", "Anycast", "北美", "https://dns.nextdns.io"},
	{"LibreDNS", "德国", "Frankfurt", "欧洲", "https://doh.libredns.gr/dns-query"},
	{"Digitale Gesellschaft", "瑞士", "Zurich", "欧洲", "https://dns.digitale-gesellschaft.ch/dns-query"},
	{"Foundation RESTENA", "卢森堡", "Luxembourg", "欧洲", "https://kaitain.restena.lu/dns-query"},
	{"CIRA Canadian Shield", "加拿大", "Ottawa", "北美", "https://private.canadianshield.cira.ca/dns-query"},
	{"Australian DNS", "澳大利亚", "Sydney", "大洋洲", "https://dns0.eu/dns-query"},
	{"DNS0.eu", "法国", "Paris", "欧洲", "https://dns0.eu/dns-query"},
}

var allowed = map[string]bool{"A": true, "AAAA": true, "CNAME": true, "MX": true, "NS": true, "TXT": true, "SOA": true, "CAA": true}

func main() {
	if len(os.Args) > 1 && os.Args[1] == "--healthcheck" {
		resp, err := (&http.Client{Timeout: 2 * time.Second}).Get("http://127.0.0.1:" + env("PORT", "8080") + "/api/health")
		if err != nil || resp.StatusCode != http.StatusOK {
			os.Exit(1)
		}
		resp.Body.Close()
		return
	}
	a := &app{queries: map[string]*query{}, client: &http.Client{Timeout: 6 * time.Second}}
	mux := http.NewServeMux()
	mux.HandleFunc("POST /api/dns/query", a.start)
	mux.HandleFunc("GET /api/dns/query/{id}/events", a.events)
	mux.HandleFunc("GET /api/health", func(w http.ResponseWriter, _ *http.Request) {
		json.NewEncoder(w).Encode(map[string]string{"status": "ok"})
	})
	mux.Handle("/web/", http.StripPrefix("/web", staticHandler()))
	mux.Handle("/", staticHandler())
	addr := ":" + env("PORT", "8080")
	log.Printf("DNS Orbit listening on %s", addr)
	log.Fatal(http.ListenAndServe(addr, securityHeaders(mux)))
}

func staticHandler() http.Handler {
	assets, err := fs.Sub(webFS, "web")
	if err != nil {
		panic(err)
	}
	return http.FileServer(http.FS(assets))
}

func env(k, v string) string {
	if x := os.Getenv(k); x != "" {
		return x
	}
	return v
}

func securityHeaders(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("X-Content-Type-Options", "nosniff")
		w.Header().Set("X-Frame-Options", "DENY")
		w.Header().Set("Referrer-Policy", "no-referrer")
		next.ServeHTTP(w, r)
	})
}

func normalizeDomain(raw string) (string, error) {
	raw = strings.TrimSpace(raw)
	if raw == "" {
		return "", errors.New("请输入域名")
	}
	if !strings.Contains(raw, "://") {
		raw = "https://" + raw
	}
	u, err := url.Parse(raw)
	if err != nil || u.Hostname() == "" {
		return "", errors.New("域名格式无效")
	}
	d := strings.TrimSuffix(strings.ToLower(u.Hostname()), ".")
	if len(d) > 253 || strings.Contains(d, "_") || net.ParseIP(d) != nil {
		return "", errors.New("仅支持公共域名")
	}
	parts := strings.Split(d, ".")
	if len(parts) < 2 {
		return "", errors.New("请输入完整公共域名")
	}
	for _, p := range parts {
		if p == "" || len(p) > 63 || p[0] == '-' || p[len(p)-1] == '-' {
			return "", errors.New("域名格式无效")
		}
		for _, c := range p {
			if !(c >= 'a' && c <= 'z' || c >= '0' && c <= '9' || c == '-') {
				return "", errors.New("域名格式无效")
			}
		}
	}
	return d, nil
}

func (a *app) start(w http.ResponseWriter, r *http.Request) {
	var in struct{ Domain, Type string }
	if json.NewDecoder(io.LimitReader(r.Body, 4096)).Decode(&in) != nil {
		http.Error(w, "invalid json", 400)
		return
	}
	d, err := normalizeDomain(in.Domain)
	t := strings.ToUpper(in.Type)
	if err != nil || !allowed[t] {
		if err == nil {
			err = errors.New("不支持该记录类型")
		}
		http.Error(w, err.Error(), 400)
		return
	}
	b := make([]byte, 6)
	rand.Read(b)
	q := &query{ID: "dns_" + hex.EncodeToString(b), Domain: d, Type: t, Created: time.Now(), Results: make(chan result, len(resolvers))}
	a.mu.Lock()
	a.queries[q.ID] = q
	a.mu.Unlock()
	go a.run(q)
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]string{"queryId": q.ID, "domain": d, "type": t, "status": "running"})
}

func (a *app) run(q *query) {
	var wg sync.WaitGroup
	for _, res := range resolvers {
		wg.Add(1)
		go func(rs resolver) { defer wg.Done(); q.Results <- a.lookup(q, rs) }(res)
	}
	wg.Wait()
	close(q.Results)
	time.AfterFunc(5*time.Minute, func() { a.mu.Lock(); delete(a.queries, q.ID); a.mu.Unlock() })
}

func (a *app) lookup(q *query, rs resolver) result {
	started := time.Now()
	out := result{Resolver: rs.Name, Country: rs.Country, City: rs.City, Region: rs.Region, Status: "error", Records: []record{}}
	u, err := url.Parse(rs.URL)
	if err != nil {
		out.Error = err.Error()
		return out
	}
	params := u.Query()
	params.Set("name", q.Domain)
	params.Set("type", q.Type)
	u.RawQuery = params.Encode()
	ctx, cancel := context.WithTimeout(context.Background(), 6*time.Second)
	defer cancel()
	req, _ := http.NewRequestWithContext(ctx, "GET", u.String(), nil)
	req.Header.Set("Accept", "application/dns-json")
	resp, err := a.client.Do(req)
	out.ResponseTime = time.Since(started).Milliseconds()
	if err != nil {
		if errors.Is(ctx.Err(), context.DeadlineExceeded) {
			out.Status = "timeout"
		}
		out.Error = "查询失败"
		return out
	}
	defer resp.Body.Close()
	if resp.StatusCode/100 != 2 {
		out.Error = fmt.Sprintf("HTTP %d", resp.StatusCode)
		return out
	}
	var body struct {
		Status int `json:"Status"`
		Answer []struct {
			Name string `json:"name"`
			Type int    `json:"type"`
			TTL  int    `json:"TTL"`
			Data string `json:"data"`
		} `json:"Answer"`
	}
	if json.NewDecoder(io.LimitReader(resp.Body, 1<<20)).Decode(&body) != nil {
		out.Error = "响应格式无效"
		return out
	}
	if body.Status != 0 {
		out.Status = "error"
		out.Error = fmt.Sprintf("DNS RCODE %d", body.Status)
		return out
	}
	if len(body.Answer) == 0 {
		out.Status = "empty"
		return out
	}
	for _, v := range body.Answer {
		out.Records = append(out.Records, record{strings.Trim(v.Data, "\""), v.TTL})
	}
	sort.Slice(out.Records, func(i, j int) bool { return out.Records[i].Value < out.Records[j].Value })
	out.Status = "success"
	return out
}

func (a *app) events(w http.ResponseWriter, r *http.Request) {
	a.mu.RLock()
	q := a.queries[r.PathValue("id")]
	a.mu.RUnlock()
	if q == nil {
		http.Error(w, "query not found", 404)
		return
	}
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	flusher, ok := w.(http.Flusher)
	if !ok {
		http.Error(w, "stream unsupported", 500)
		return
	}
	for x := range q.Results {
		b, _ := json.Marshal(x)
		fmt.Fprintf(w, "event: result\ndata: %s\n\n", b)
		flusher.Flush()
	}
	fmt.Fprint(w, "event: complete\ndata: {}\n\n")
	flusher.Flush()
}