package main

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

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

type resolver struct{ Name, Region, Country, City, URL string }
type record struct { Value string `json:"value"`; TTL uint32 `json:"ttl"` }
type result struct { Node resolverView `json:"node"`; Status string `json:"status"`; Records []record `json:"records"`; ResponseTime int64 `json:"responseTime"`; Error string `json:"error,omitempty"` }
type resolverView struct { Name string `json:"name"`; Region string `json:"region"`; Country string `json:"country"`; City string `json:"city"` }
type query struct { ID string `json:"queryId"`; Domain string `json:"domain"`; Type string `json:"type"`; Status string `json:"status"`; CreatedAt time.Time `json:"createdAt"`; Results []result `json:"results,omitempty"`; Subscribers map[chan result]struct{} `json:"-"`; mu sync.RWMutex }
type store struct { mu sync.RWMutex; queries map[string]*query }

var nodes = []resolver{
	{"Cloudflare", "North America", "United States", "Global Anycast", "https://cloudflare-dns.com/dns-query"},
	{"Google", "North America", "United States", "Global Anycast", "https://dns.google/resolve"},
	{"AdGuard", "Europe", "Cyprus", "Global Anycast", "https://dns.adguard-dns.com/resolve"},
	{"AliDNS", "Asia", "China", "Global Anycast", "https://dns.alidns.com/resolve"},
}

func main() {
	s := &store{queries: make(map[string]*query)}
	mux := http.NewServeMux()
	mux.HandleFunc("GET /healthz", func(w http.ResponseWriter, _ *http.Request) { w.Header().Set("Content-Type", "application/json"); io.WriteString(w, `{"status":"ok"}`) })
	mux.HandleFunc("POST /api/dns/query", s.create)
	mux.HandleFunc("GET /api/dns/query/{id}", s.get)
	mux.HandleFunc("GET /api/dns/query/{id}/events", s.events)

	server := &http.Server{Addr: env("ADDR", ":8080"), Handler: security(rateLimit(mux)), ReadHeaderTimeout: 5 * time.Second}
	go s.cleanup()
	go func() { slog.Info("server started", "addr", server.Addr); if err := server.ListenAndServe(); !errors.Is(err, http.ErrServerClosed) { slog.Error("server stopped", "error", err); os.Exit(1) } }()
	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM); defer stop(); <-ctx.Done()
	shutdown, cancel := context.WithTimeout(context.Background(), 10*time.Second); defer cancel(); _ = server.Shutdown(shutdown)
}

func (s *store) create(w http.ResponseWriter, r *http.Request) {
	var body struct { Domain string `json:"domain"`; Type string `json:"type"` }
	if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 4096)).Decode(&body); err != nil { problem(w, 400, "请求格式无效"); return }
	domain, err := normalizeDomain(body.Domain); kind := strings.ToUpper(strings.TrimSpace(body.Type))
	if err != nil { problem(w, 400, err.Error()); return }; if !allowedTypes[kind] { problem(w, 400, "不支持的记录类型"); return }
	q := &query{ID: newID(), Domain: domain, Type: kind, Status: "running", CreatedAt: time.Now().UTC(), Subscribers: make(map[chan result]struct{})}
	s.mu.Lock(); s.queries[q.ID] = q; s.mu.Unlock(); go s.run(q)
	writeJSON(w, 202, q)
}

func (s *store) get(w http.ResponseWriter, r *http.Request) { q := s.find(r.PathValue("id")); if q == nil { problem(w, 404, "查询不存在或已过期"); return }; q.mu.RLock(); defer q.mu.RUnlock(); writeJSON(w, 200, q) }

func (s *store) events(w http.ResponseWriter, r *http.Request) {
	q := s.find(r.PathValue("id")); if q == nil { problem(w, 404, "查询不存在或已过期"); return }
	flusher, ok := w.(http.Flusher); if !ok { problem(w, 500, "不支持事件流"); return }
	w.Header().Set("Content-Type", "text/event-stream"); w.Header().Set("Cache-Control", "no-cache"); w.Header().Set("X-Accel-Buffering", "no")
	ch := make(chan result, len(nodes)); q.mu.Lock(); existing := append([]result(nil), q.Results...); done := q.Status == "completed"; if !done { q.Subscribers[ch] = struct{}{} }; q.mu.Unlock()
	defer func() { q.mu.Lock(); delete(q.Subscribers, ch); q.mu.Unlock() }()
	for _, item := range existing { sendEvent(w, "result", item) }; if done { sendEvent(w, "complete", map[string]any{"count": len(existing)}); flusher.Flush(); return }; flusher.Flush()
	for { select { case item, open := <-ch: if !open { sendEvent(w, "complete", map[string]any{"count": len(nodes)}); flusher.Flush(); return }; sendEvent(w, "result", item); flusher.Flush(); case <-r.Context().Done(): return } }
}

func (s *store) run(q *query) { var wg sync.WaitGroup; for _, n := range nodes { wg.Add(1); go func() { defer wg.Done(); item := resolve(q.Domain, q.Type, n); q.mu.Lock(); q.Results = append(q.Results, item); for ch := range q.Subscribers { ch <- item }; q.mu.Unlock() }() }; wg.Wait(); q.mu.Lock(); q.Status = "completed"; for ch := range q.Subscribers { close(ch); delete(q.Subscribers, ch) }; q.mu.Unlock() }

func resolve(domain, kind string, n resolver) result {
	view := resolverView{n.Name, n.Region, n.Country, n.City}; started := time.Now(); ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second); defer cancel()
	u, _ := url.Parse(n.URL); params := u.Query(); params.Set("name", domain); params.Set("type", kind); u.RawQuery = params.Encode()
	req, _ := http.NewRequestWithContext(ctx, http.MethodGet, u.String(), nil); req.Header.Set("Accept", "application/dns-json"); resp, err := http.DefaultClient.Do(req); elapsed := time.Since(started).Milliseconds()
	if err != nil { status := "failed"; if errors.Is(ctx.Err(), context.DeadlineExceeded) { status = "timeout" }; return result{Node: view, Status: status, ResponseTime: elapsed, Error: "解析器请求失败"} }; defer resp.Body.Close()
	if resp.StatusCode != 200 { return result{Node: view, Status: "failed", ResponseTime: elapsed, Error: fmt.Sprintf("HTTP %d", resp.StatusCode)} }
	var payload struct { Status int `json:"Status"`; Answer []struct { Data string `json:"data"`; TTL uint32 `json:"TTL"` } `json:"Answer"` }
	if json.NewDecoder(io.LimitReader(resp.Body, 1<<20)).Decode(&payload) != nil || payload.Status != 0 { return result{Node: view, Status: "failed", ResponseTime: elapsed, Error: fmt.Sprintf("DNS 状态 %d", payload.Status)} }
	if len(payload.Answer) == 0 { return result{Node: view, Status: "no_record", ResponseTime: elapsed, Records: []record{}} }
	records := make([]record, 0, len(payload.Answer)); for _, a := range payload.Answer { records = append(records, record{strings.TrimSuffix(a.Data, "."), a.TTL}) }; sort.Slice(records, func(i,j int) bool{return records[i].Value<records[j].Value})
	return result{Node: view, Status: "success", Records: records, ResponseTime: elapsed}
}

func normalizeDomain(input string) (string, error) { input = strings.TrimSpace(input); if !strings.Contains(input, "://") { input = "//" + input }; u, err := url.Parse(input); if err != nil || u.Hostname() == "" { return "", errors.New("请输入合法域名") }; host := strings.TrimSuffix(strings.ToLower(u.Hostname()), "."); if len(host)>253 || net.ParseIP(host)!=nil || host=="localhost" || !validLabels(host) { return "", errors.New("请输入合法的公网域名") }; return host,nil }
func validLabels(host string) bool { labels:=strings.Split(host,"."); if len(labels)<2{return false}; for _,l:=range labels { if len(l)<1||len(l)>63||l[0]=='-'||l[len(l)-1]=='-'{return false}; for _,c:=range l {if !(c>='a'&&c<='z'||c>='0'&&c<='9'||c=='-'){return false}} }; return true }
func (s *store) find(id string)*query{s.mu.RLock();defer s.mu.RUnlock();return s.queries[id]}
func (s *store) cleanup(){ticker:=time.NewTicker(time.Minute);defer ticker.Stop();for range ticker.C{s.mu.Lock();for id,q:=range s.queries{if time.Since(q.CreatedAt)>5*time.Minute{delete(s.queries,id)}};s.mu.Unlock()}}
func sendEvent(w io.Writer,name string,data any){b,_:=json.Marshal(data);fmt.Fprintf(w,"event: %s\ndata: %s\n\n",name,b)}
func writeJSON(w http.ResponseWriter,status int,v any){w.Header().Set("Content-Type","application/json");w.WriteHeader(status);_ = json.NewEncoder(w).Encode(v)}
func problem(w http.ResponseWriter,status int,msg string){writeJSON(w,status,map[string]string{"error":msg})}
func newID()string{b:=make([]byte,6);_,_=rand.Read(b);return "dns_"+hex.EncodeToString(b)}
func env(k,d string)string{if v:=os.Getenv(k);v!=""{return v};return d}

func security(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","same-origin");next.ServeHTTP(w,r)})}
var visitors sync.Map
func rateLimit(next http.Handler) http.Handler{return http.HandlerFunc(func(w http.ResponseWriter,r *http.Request){if r.Method=="POST" {ip,_,_:=net.SplitHostPort(r.RemoteAddr);now:=time.Now();raw,_:=visitors.LoadOrStore(ip,&struct{sync.Mutex;start time.Time;n int}{start:now});v:=raw.(*struct{sync.Mutex;start time.Time;n int});v.Lock();if now.Sub(v.start)>time.Minute{v.start=now;v.n=0};v.n++;blocked:=v.n>30;v.Unlock();if blocked{problem(w,429,"请求过于频繁，请稍后重试");return}};next.ServeHTTP(w,r)})}