package app

import (
	"context"
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"
	"time"

	"balancedns/internal/cache"
	"balancedns/internal/config"
	"balancedns/internal/hosts"
	"balancedns/internal/logx"
	"balancedns/internal/metrics"
	"balancedns/internal/plugin"
	"balancedns/internal/router"
	control "balancedns/internal/runtime"

	"github.com/miekg/dns"
)

type Server struct {
	cfg      *config.Config
	logger   *logx.Logger
	metrics  *metrics.Provider
	cache    *cache.Cache
	plugins  *plugin.Engine
	resolver *router.Resolver
	hosts    *hosts.Table
	acl      []*net.IPNet

	chain     []string
	blacklist []blacklistRule

	supervisor *control.Supervisor
}

type blacklistRule struct {
	suffix bool
	value  string
}

func New(cfg *config.Config) (*Server, error) {
	m := metrics.New()
	logger := logx.New(cfg.Logging.Level, cfg.Logging.LogQueries)

	resolver, err := router.NewResolver(cfg.Upstreams, m)
	if err != nil {
		return nil, err
	}

	var c *cache.Cache
	if cfg.Cache.Enabled {
		c = cache.New(cfg.Cache.Capacity, cfg.Cache.MinTTLSeconds, cfg.Cache.MaxTTLSeconds)
	}

	var engine *plugin.Engine
	if cfg.Plugins.Enabled && len(cfg.Plugins.Entries) > 0 {
		engine, err = plugin.NewEngine(cfg.Plugins.Entries, time.Duration(cfg.Plugins.TimeoutMS)*time.Millisecond)
		if err != nil {
			return nil, err
		}
	}

	var hostTable *hosts.Table
	if cfg.Hosts.File != "" {
		hostTable, err = hosts.Load(cfg.Hosts.File, cfg.Hosts.TTL)
		if err != nil {
			return nil, err
		}
	}

	acl, err := parseACL(cfg.ACL)
	if err != nil {
		return nil, err
	}

	s := &Server{
		cfg:       cfg,
		logger:    logger,
		metrics:   m,
		cache:     c,
		plugins:   engine,
		resolver:  resolver,
		hosts:     hostTable,
		acl:       acl,
		chain:     normalizeChain(cfg.Routing.Chain),
		blacklist: parseBlacklist(cfg.Blacklist.Domains),
	}

	return s, nil
}

func (s *Server) Run(ctx context.Context) error {
	dnsMux := dns.NewServeMux()
	dnsMux.HandleFunc(".", s.handleDNS)

	components := []control.ComponentConfig{
		{
			Name:     "dns-udp",
			Required: true,
			Start:    s.runDNSComponent("udp", s.cfg.Listen.DNS, dnsMux),
		},
		{
			Name:     "dns-tcp",
			Required: true,
			Start:    s.runDNSComponent("tcp", s.cfg.Listen.DNS, dnsMux),
		},
		{
			Name:     "metrics-http",
			Required: true,
			Start:    s.runMetricsComponent(),
		},
	}

	if s.cfg.Listen.DoT != "" {
		components = append(components, control.ComponentConfig{
			Name:     "dns-dot",
			Required: true,
			Start:    s.runDoTComponent(dnsMux),
		})
	}
	if s.cfg.Listen.DoH != "" {
		components = append(components, control.ComponentConfig{
			Name:     "dns-doh",
			Required: true,
			Start:    s.runDoHComponent(),
		})
	}

	s.supervisor = control.New(s.logger, s.metrics, components, control.Options{
		RestartBackoff:      time.Duration(s.cfg.Control.RestartBackoffMS) * time.Millisecond,
		RestartMaxBackoff:   time.Duration(s.cfg.Control.RestartMaxBackoffMS) * time.Millisecond,
		MaxConsecutiveFails: s.cfg.Control.MaxConsecutiveFailure,
		MinStableRun:        time.Duration(s.cfg.Control.MinStableRunMS) * time.Millisecond,
	})

	s.logger.Infof("control plane started: components=%d", len(components))
	err := s.supervisor.Run(ctx)
	if err != nil {
		return err
	}
	s.logger.Infof("graceful shutdown completed")
	return nil
}

func (s *Server) runDNSComponent(network, addr string, handler dns.Handler) func(context.Context) error {
	return func(ctx context.Context) error {
		srv := &dns.Server{
			Addr:         addr,
			Net:          network,
			Handler:      handler,
			ReusePort:    s.cfg.Listen.ReusePort,
			ReuseAddr:    s.cfg.Listen.ReuseAddr,
			UDPSize:      s.cfg.Listen.UDPSize,
			ReadTimeout:  time.Duration(s.cfg.Listen.ReadTimeoutMS) * time.Millisecond,
			WriteTimeout: time.Duration(s.cfg.Listen.WriteTimeoutMS) * time.Millisecond,
		}

		go func() {
			<-ctx.Done()
			_ = srv.Shutdown()
		}()

		s.logger.Infof("%s component started on %s", network, addr)
		err := srv.ListenAndServe()
		if ctx.Err() != nil {
			return nil
		}
		return fmt.Errorf("%s listener failed: %w", network, err)
	}
}

func (s *Server) runDoTComponent(handler dns.Handler) func(context.Context) error {
	return func(ctx context.Context) error {
		cert, err := tls.LoadX509KeyPair(s.cfg.Listen.TLSCertFile, s.cfg.Listen.TLSKeyFile)
		if err != nil {
			return fmt.Errorf("load dot certificate: %w", err)
		}

		srv := &dns.Server{
			Addr:         s.cfg.Listen.DoT,
			Net:          "tcp-tls",
			Handler:      handler,
			ReadTimeout:  time.Duration(s.cfg.Listen.ReadTimeoutMS) * time.Millisecond,
			WriteTimeout: time.Duration(s.cfg.Listen.WriteTimeoutMS) * time.Millisecond,
			TLSConfig: &tls.Config{
				MinVersion:   tls.VersionTLS12,
				Certificates: []tls.Certificate{cert},
			},
		}

		go func() {
			<-ctx.Done()
			_ = srv.Shutdown()
		}()

		s.logger.Infof("dot component started on %s", s.cfg.Listen.DoT)
		err = srv.ListenAndServe()
		if ctx.Err() != nil {
			return nil
		}
		return fmt.Errorf("dot listener failed: %w", err)
	}
}

func (s *Server) runDoHComponent() func(context.Context) error {
	return func(ctx context.Context) error {
		mux := http.NewServeMux()
		mux.HandleFunc(s.cfg.Listen.DoHPath, s.handleDoH)

		server := &http.Server{
			Addr:         s.cfg.Listen.DoH,
			Handler:      mux,
			ReadTimeout:  time.Duration(s.cfg.Listen.ReadTimeoutMS) * time.Millisecond,
			WriteTimeout: time.Duration(s.cfg.Listen.WriteTimeoutMS) * time.Millisecond,
			TLSConfig: &tls.Config{
				MinVersion: tls.VersionTLS12,
			},
		}

		go func() {
			<-ctx.Done()
			shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			defer cancel()
			_ = server.Shutdown(shutdownCtx)
		}()

		s.logger.Infof("doh component started on %s%s", s.cfg.Listen.DoH, s.cfg.Listen.DoHPath)
		err := server.ListenAndServeTLS(s.cfg.Listen.TLSCertFile, s.cfg.Listen.TLSKeyFile)
		if ctx.Err() != nil || errors.Is(err, http.ErrServerClosed) {
			return nil
		}
		return fmt.Errorf("doh listener failed: %w", err)
	}
}

func (s *Server) runMetricsComponent() func(context.Context) error {
	return func(ctx context.Context) error {
		mux := http.NewServeMux()
		mux.Handle("/metrics", s.metrics.Handler())
		mux.HandleFunc("/healthz", s.handleHealthz)
		mux.HandleFunc("/readyz", s.handleReadyz)
		mux.HandleFunc("/statusz", s.handleStatusz)

		server := &http.Server{
			Addr:         s.cfg.Listen.Metrics,
			Handler:      mux,
			ReadTimeout:  time.Duration(s.cfg.Listen.ReadTimeoutMS) * time.Millisecond,
			WriteTimeout: time.Duration(s.cfg.Listen.WriteTimeoutMS) * time.Millisecond,
		}

		go func() {
			<-ctx.Done()
			shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			defer cancel()
			_ = server.Shutdown(shutdownCtx)
		}()

		s.logger.Infof("metrics component started on %s", s.cfg.Listen.Metrics)
		err := server.ListenAndServe()
		if ctx.Err() != nil || errors.Is(err, http.ErrServerClosed) {
			return nil
		}
		return fmt.Errorf("metrics listener failed: %w", err)
	}
}

func (s *Server) handleHealthz(w http.ResponseWriter, _ *http.Request) {
	s.writeSupervisorStatus(w, s.supervisor.Healthy(true))
}

func (s *Server) handleReadyz(w http.ResponseWriter, _ *http.Request) {
	s.writeSupervisorStatus(w, s.supervisor.Healthy(true))
}

func (s *Server) handleStatusz(w http.ResponseWriter, _ *http.Request) {
	s.writeSupervisorStatus(w, true)
}

func (s *Server) writeSupervisorStatus(w http.ResponseWriter, healthy bool) {
	if s.supervisor == nil {
		http.Error(w, "supervisor not initialized", http.StatusServiceUnavailable)
		return
	}

	resp := struct {
		Status     string          `json:"status"`
		Timestamp  time.Time       `json:"timestamp"`
		Components []control.State `json:"components"`
	}{
		Status:     "ok",
		Timestamp:  time.Now().UTC(),
		Components: s.supervisor.Snapshot(),
	}
	w.Header().Set("Content-Type", "application/json")
	if !healthy {
		resp.Status = "degraded"
		w.WriteHeader(http.StatusServiceUnavailable)
	} else {
		w.WriteHeader(http.StatusOK)
	}
	_ = json.NewEncoder(w).Encode(resp)
}

func (s *Server) handleDNS(w dns.ResponseWriter, req *dns.Msg) {
	resp := s.resolveDNS(req, w.RemoteAddr())
	if err := w.WriteMsg(resp); err != nil {
		s.logger.Errorf("write dns response: %v", err)
	}
}

func (s *Server) handleDoH(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet && r.Method != http.MethodPost {
		w.Header().Set("Allow", "GET, POST")
		http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
		return
	}

	remoteIP := remoteIPFromString(r.RemoteAddr)
	if !s.allowedRemoteIP(remoteIP) {
		http.Error(w, "forbidden", http.StatusForbidden)
		return
	}

	wire, err := readDoHWireMessage(r)
	if err != nil {
		http.Error(w, err.Error(), http.StatusBadRequest)
		return
	}

	req := new(dns.Msg)
	if err := req.Unpack(wire); err != nil {
		http.Error(w, "invalid dns message", http.StatusBadRequest)
		return
	}

	resp := s.resolveDNS(req, &net.TCPAddr{IP: remoteIP})
	payload, err := resp.Pack()
	if err != nil {
		http.Error(w, "encode dns response", http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/dns-message")
	w.WriteHeader(http.StatusOK)
	_, _ = w.Write(payload)
}

func readDoHWireMessage(r *http.Request) ([]byte, error) {
	switch r.Method {
	case http.MethodGet:
		encoded := strings.TrimSpace(r.URL.Query().Get("dns"))
		if encoded == "" {
			return nil, errors.New("missing dns query parameter")
		}
		data, err := base64.RawURLEncoding.DecodeString(encoded)
		if err != nil {
			return nil, errors.New("invalid dns query parameter")
		}
		return data, nil
	case http.MethodPost:
		defer r.Body.Close()
		data, err := io.ReadAll(io.LimitReader(r.Body, 65535))
		if err != nil {
			return nil, errors.New("failed to read request body")
		}
		if len(data) == 0 {
			return nil, errors.New("empty request body")
		}
		return data, nil
	default:
		return nil, errors.New("unsupported method")
	}
}

func (s *Server) resolveDNS(req *dns.Msg, remoteAddr net.Addr) *dns.Msg {
	s.metrics.IncQueries()
	if len(req.Question) == 0 {
		return s.rcodeResponse(req, dns.RcodeFormatError)
	}

	if !s.allowedRemoteIP(remoteIPFromNetAddr(remoteAddr)) {
		return s.rcodeResponse(req, dns.RcodeRefused)
	}

	current := normalizeQuestion(req.Question[0])
	remote := "<unknown>"
	if remoteAddr != nil {
		remote = remoteAddr.String()
	}
	s.logger.Queryf("query id=%d remote=%s domain=%s type=%s", req.Id, remote, current.Name, dns.TypeToString[current.Qtype])

	for _, stage := range s.chain {
		switch stage {
		case "blacklist":
			if s.isBlocked(current.Name) {
				s.logger.Debugf("blocked domain %s", current.Name)
				return s.rcodeResponse(req, dns.RcodeRefused)
			}

		case "hosts":
			if s.hosts == nil {
				continue
			}
			if ans, ok := s.hosts.Lookup(current.Name, current.Qtype); ok {
				return s.localDataResponse(req, current, plugin.LocalData{IPs: ans.IPs, TTL: ans.TTL})
			}

		case "cache":
			if s.cache == nil {
				continue
			}
			if cached, ok := s.cache.Get(current); ok {
				s.metrics.IncCacheHits()
				cached.Id = req.Id
				cached.Question = []dns.Question{current}
				return cached
			}

		case "lua_policy", "plugin", "plugins", "lua":
			if s.plugins == nil {
				continue
			}
			decision, err := s.plugins.Decide(current)
			if err != nil {
				s.metrics.IncPluginErrors()
				s.logger.Errorf("plugin execution error for %s: %v", current.Name, err)
				continue
			}
			switch decision.Action {
			case plugin.ActionBlock:
				return s.rcodeResponse(req, dns.RcodeRefused)
			case plugin.ActionLocalData:
				return s.localDataResponse(req, decision.Question, decision.Local)
			case plugin.ActionRewrite, plugin.ActionForward:
				current = normalizeQuestion(decision.Question)
			}

		case "upstream":
			resp, up, err := s.resolver.Forward(context.Background(), req, current)
			if err != nil {
				s.logger.Errorf("upstream forward failed for %s: %v", current.Name, err)
				return s.rcodeResponse(req, dns.RcodeServerFailure)
			}
			if s.cache != nil && resp.Rcode == dns.RcodeSuccess {
				s.cache.Set(current, resp)
			}
			s.logger.Debugf("upstream=%s served domain=%s type=%s", up.Name, current.Name, dns.TypeToString[current.Qtype])
			return resp
		default:
			s.logger.Debugf("unknown chain stage: %s", stage)
		}
	}

	return s.rcodeResponse(req, dns.RcodeServerFailure)
}

func (s *Server) localDataResponse(req *dns.Msg, q dns.Question, local plugin.LocalData) *dns.Msg {
	resp := new(dns.Msg)
	resp.SetReply(req)
	resp.Question = []dns.Question{q}
	resp.Authoritative = true

	ttl := local.TTL
	if ttl == 0 {
		ttl = 60
	}

	for _, ip := range local.IPs {
		switch {
		case q.Qtype == dns.TypeA && ip.To4() != nil:
			resp.Answer = append(resp.Answer, &dns.A{Hdr: dns.RR_Header{Name: q.Name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: ttl}, A: ip.To4()})
		case q.Qtype == dns.TypeAAAA && ip.To16() != nil && ip.To4() == nil:
			resp.Answer = append(resp.Answer, &dns.AAAA{Hdr: dns.RR_Header{Name: q.Name, Rrtype: dns.TypeAAAA, Class: dns.ClassINET, Ttl: ttl}, AAAA: ip.To16()})
		case q.Qtype == dns.TypeANY:
			if ip.To4() != nil {
				resp.Answer = append(resp.Answer, &dns.A{Hdr: dns.RR_Header{Name: q.Name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: ttl}, A: ip.To4()})
			} else if ip.To16() != nil {
				resp.Answer = append(resp.Answer, &dns.AAAA{Hdr: dns.RR_Header{Name: q.Name, Rrtype: dns.TypeAAAA, Class: dns.ClassINET, Ttl: ttl}, AAAA: ip.To16()})
			}
		}
	}

	return resp
}

func (s *Server) isBlocked(name string) bool {
	normalized := normalizeDomain(name)
	for _, rule := range s.blacklist {
		if rule.suffix {
			if strings.HasSuffix(normalized, rule.value) {
				return true
			}
			continue
		}
		if normalized == rule.value {
			return true
		}
	}
	return false
}

func (s *Server) rcodeResponse(req *dns.Msg, rcode int) *dns.Msg {
	msg := new(dns.Msg)
	msg.SetRcode(req, rcode)
	return msg
}

func parseBlacklist(domains []string) []blacklistRule {
	rules := make([]blacklistRule, 0, len(domains))
	for _, d := range domains {
		d = strings.TrimSpace(strings.ToLower(d))
		if d == "" {
			continue
		}
		if strings.HasPrefix(d, "*.") {
			rules = append(rules, blacklistRule{suffix: true, value: normalizeDomain(strings.TrimPrefix(d, "*"))})
			continue
		}
		if strings.HasPrefix(d, ".") {
			rules = append(rules, blacklistRule{suffix: true, value: normalizeDomain(d)})
			continue
		}
		rules = append(rules, blacklistRule{value: normalizeDomain(d)})
	}
	return rules
}

func normalizeChain(chain []string) []string {
	out := make([]string, 0, len(chain))
	for _, stage := range chain {
		n := strings.TrimSpace(strings.ToLower(stage))
		if n != "" {
			out = append(out, n)
		}
	}
	if len(out) == 0 {
		return []string{"blacklist", "hosts", "cache", "lua_policy", "upstream"}
	}
	return out
}

func normalizeQuestion(q dns.Question) dns.Question {
	q.Name = normalizeDomain(q.Name)
	if q.Qclass == 0 {
		q.Qclass = dns.ClassINET
	}
	return q
}

func normalizeDomain(name string) string {
	return strings.ToLower(dns.Fqdn(strings.TrimSpace(name)))
}

func parseACL(values []string) ([]*net.IPNet, error) {
	if len(values) == 0 {
		return nil, nil
	}
	out := make([]*net.IPNet, 0, len(values))
	for i, value := range values {
		netmask, err := parseCIDROrIP(value)
		if err != nil {
			return nil, fmt.Errorf("acl[%d]: %w", i, err)
		}
		out = append(out, netmask)
	}
	return out, nil
}

func parseCIDROrIP(value string) (*net.IPNet, error) {
	v := strings.TrimSpace(value)
	if v == "" {
		return nil, errors.New("empty ACL value")
	}
	if _, ipnet, err := net.ParseCIDR(v); err == nil {
		return ipnet, nil
	}
	ip := net.ParseIP(v)
	if ip == nil {
		return nil, fmt.Errorf("invalid CIDR/IP %q", value)
	}
	if ip.To4() != nil {
		_, ipnet, _ := net.ParseCIDR(ip.String() + "/32")
		return ipnet, nil
	}
	_, ipnet, _ := net.ParseCIDR(ip.String() + "/128")
	return ipnet, nil
}

func (s *Server) allowedRemoteIP(ip net.IP) bool {
	if len(s.acl) == 0 {
		return true
	}
	if ip == nil {
		return false
	}
	for _, netmask := range s.acl {
		if netmask.Contains(ip) {
			return true
		}
	}
	return false
}

func remoteIPFromNetAddr(addr net.Addr) net.IP {
	if addr == nil {
		return nil
	}
	switch v := addr.(type) {
	case *net.TCPAddr:
		return v.IP
	case *net.UDPAddr:
		return v.IP
	default:
		return remoteIPFromString(addr.String())
	}
}

func remoteIPFromString(raw string) net.IP {
	host := strings.TrimSpace(raw)
	if host == "" {
		return nil
	}
	if h, _, err := net.SplitHostPort(host); err == nil {
		host = h
	}
	return net.ParseIP(host)
}