package main

import (
	"context"
	"errors"
	"fmt"
	"math/rand"
	"net"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"codeberg.org/miekg/dns"
	"github.com/jedisct1/dlog"
	"github.com/lifenjoiner/dhcpdns"
	netproxy "golang.org/x/net/proxy"
)

type SearchSequenceItemType int

const (
	Explicit SearchSequenceItemType = iota
	Bootstrap
	DHCP
	Resolvconf
)

type SearchSequenceItem struct {
	typ        SearchSequenceItemType
	servers    []string
	resolvconf string
	rcLastFail atomic.Int64 // unix timestamp of last failed resolv.conf read
}

type PluginForwardEntry struct {
	domain   string
	sequence []SearchSequenceItem
	viaProxy bool
}

type PluginForward struct {
	forwardMap         []PluginForwardEntry
	bootstrapResolvers []string
	dhcpdns            []*dhcpdns.Detector
	proxyDialer        netproxy.Dialer

	// Hot-reloading support
	rwLock        sync.RWMutex
	configFile    string
	configWatcher *ConfigWatcher
	stagingMap    []PluginForwardEntry
}

func (plugin *PluginForward) Name() string {
	return "forward"
}

func (plugin *PluginForward) Description() string {
	return "Route queries matching specific domains to a dedicated set of servers"
}

func (plugin *PluginForward) Init(proxy *Proxy) error {
	plugin.configFile = proxy.forwardFile
	dlog.Noticef("Loading the set of forwarding rules from [%s]", plugin.configFile)

	if proxy.xTransport != nil {
		plugin.bootstrapResolvers = proxy.xTransport.bootstrapResolvers
		if proxy.xTransport.proxyDialer != nil {
			plugin.proxyDialer = *proxy.xTransport.proxyDialer
		}
	}

	lines, err := ReadTextFile(plugin.configFile)
	if err != nil {
		return err
	}

	requiresDHCP, forwardMap, err := plugin.parseForwardFile(lines)
	if err != nil {
		return err
	}

	plugin.forwardMap = forwardMap

	if requiresDHCP {
		if len(proxy.userName) > 0 {
			dlog.Warn("DHCP/DNS detection may not work when 'user_name' is set or when starting as a non-root user")
		}
		if proxy.SourceIPv6 {
			dlog.Notice("Starting a DHCP/DNS detector for IPv6")
			d6 := &dhcpdns.Detector{RemoteIPPort: "[2001:DB8::53]:80"}
			go d6.Serve(9, 10)
			plugin.dhcpdns = append(plugin.dhcpdns, d6)
		}
		if proxy.SourceIPv4 {
			dlog.Notice("Starting a DHCP/DNS detector for IPv4")
			d4 := &dhcpdns.Detector{RemoteIPPort: "192.0.2.53:80"}
			go d4.Serve(9, 10)
			plugin.dhcpdns = append(plugin.dhcpdns, d4)
		}
	}
	return nil
}

// parseForwardFile parses forward rules from text
func (plugin *PluginForward) parseForwardFile(lines string) (bool, []PluginForwardEntry, error) {
	requiresDHCP := false
	forwardMap := []PluginForwardEntry{}

	for lineNo, line := range strings.Split(lines, "\n") {
		line = TrimAndStripInlineComments(line)
		if len(line) == 0 {
			continue
		}
		domain, serversStr, ok := StringTwoFields(line)
		domain = strings.TrimPrefix(domain, "*.")
		if strings.Contains(domain, "*") {
			ok = false
		}
		if !ok {
			return false, nil, fmt.Errorf(
				"Syntax error for a forwarding rule at line %d. Expected syntax: example.com 9.9.9.9,8.8.8.8",
				1+lineNo,
			)
		}
		domain = strings.ToLower(domain)
		viaProxy := false
		if strings.HasPrefix(serversStr, "$PROXY:") {
			viaProxy = true
			serversStr = strings.TrimSpace(serversStr[len("$PROXY:"):])
			if plugin.proxyDialer == nil {
				return false, nil, fmt.Errorf(
					"Forwarding rule at line %d uses $PROXY:, but no proxy is available. The `proxy` option must be configured in the main configuration file",
					1+lineNo,
				)
			}
		}
		var sequence []SearchSequenceItem
		for server := range strings.SplitSeq(serversStr, ",") {
			server = strings.TrimSpace(server)
			switch server {
			case "$BOOTSTRAP":
				if len(plugin.bootstrapResolvers) == 0 {
					return false, nil, fmt.Errorf(
						"Syntax error for a forwarding rule at line %d. No bootstrap resolvers available",
						1+lineNo,
					)
				}
				if len(sequence) > 0 && sequence[len(sequence)-1].typ == Bootstrap {
					// Ignore repetitions
				} else {
					sequence = append(sequence, SearchSequenceItem{typ: Bootstrap})
					dlog.Infof("Forwarding [%s] to the bootstrap servers", domain)
				}
			case "$DHCP":
				if len(sequence) > 0 && sequence[len(sequence)-1].typ == DHCP {
					// Ignore repetitions
				} else {
					sequence = append(sequence, SearchSequenceItem{typ: DHCP})
					dlog.Infof("Forwarding [%s] to the DHCP servers", domain)
				}
				requiresDHCP = true
			default:
				const resolvconfPrefix = "$RESOLVCONF:"
				if strings.HasPrefix(server, resolvconfPrefix) {
					file := server[len(resolvconfPrefix):]
					if len(file) == 0 {
						dlog.Criticalf(
							"File needs to be specified for $RESOLVCONF in line %d",
							1+lineNo,
						)
						continue
					}
					file = filepath.Clean(file)
					if !filepath.IsAbs(file) {
						dlog.Warnf(
							"$RESOLVCONF path '%s' at line %d is not absolute; "+
								"this may not resolve as expected", file, 1+lineNo,
						)
					}
					sequence = append(sequence, SearchSequenceItem{
						typ:        Resolvconf,
						resolvconf: file,
					})
					dlog.Infof("Forwarding [%s] to the servers specified in '%s'", domain, file)
					continue
				}
				if strings.HasPrefix(server, "$") {
					dlog.Criticalf("Unknown keyword [%s] at line %d", server, 1+lineNo)
					continue
				}
				if server, err := normalizeIPAndOptionalPort(server, "53"); err != nil {
					dlog.Criticalf("Syntax error for a forwarding rule at line %d: %s", 1+lineNo, err)
					continue
				} else {
					idxServers := -1
					for i := range sequence {
						if sequence[i].typ == Explicit {
							idxServers = i
						}
					}
					if idxServers == -1 {
						sequence = append(sequence, SearchSequenceItem{typ: Explicit, servers: []string{server}})
					} else {
						sequence[idxServers].servers = append(sequence[idxServers].servers, server)
					}
					dlog.Infof("Forwarding [%s] to [%s]", domain, server)
				}
			}
		}
		if viaProxy {
			if len(sequence) == 0 {
				return false, nil, fmt.Errorf(
					"Syntax error for a forwarding rule at line %d. $PROXY: requires at least one valid server address",
					1+lineNo,
				)
			}
			for i := range sequence {
				if sequence[i].typ != Explicit {
					return false, nil, fmt.Errorf(
						"Syntax error for a forwarding rule at line %d. $PROXY: can only be combined with explicit server IP addresses",
						1+lineNo,
					)
				}
			}
		}
		forwardMap = append(forwardMap, PluginForwardEntry{
			domain:   domain,
			sequence: sequence,
			viaProxy: viaProxy,
		})
	}

	return requiresDHCP, forwardMap, nil
}

func (plugin *PluginForward) Drop() error {
	if plugin.configWatcher != nil {
		plugin.configWatcher.RemoveFile(plugin.configFile)
	}
	return nil
}

// PrepareReload loads new rules into staging structure but doesn't apply them yet
func (plugin *PluginForward) PrepareReload() error {
	// Read the configuration file
	lines, err := SafeReadTextFile(plugin.configFile)
	if err != nil {
		return fmt.Errorf("error reading config file during reload preparation: %w", err)
	}

	// Parse the forward file
	_, stagingMap, err := plugin.parseForwardFile(lines)
	if err != nil {
		return fmt.Errorf("error parsing config during reload preparation: %w", err)
	}

	// Store in staging area
	plugin.rwLock.Lock()
	plugin.stagingMap = stagingMap
	plugin.rwLock.Unlock()

	return nil
}

// ApplyReload atomically replaces the active rules with the staging ones
func (plugin *PluginForward) ApplyReload() error {
	plugin.rwLock.Lock()
	defer plugin.rwLock.Unlock()

	if plugin.stagingMap == nil {
		return errors.New("no staged configuration to apply")
	}

	plugin.forwardMap = plugin.stagingMap
	plugin.stagingMap = nil

	dlog.Noticef("Applied new configuration for plugin [%s]", plugin.Name())
	return nil
}

// CancelReload cleans up any staging resources
func (plugin *PluginForward) CancelReload() {
	plugin.rwLock.Lock()
	plugin.stagingMap = nil
	plugin.rwLock.Unlock()
}

// Reload implements hot-reloading for the plugin
func (plugin *PluginForward) Reload() error {
	dlog.Noticef("Reloading configuration for plugin [%s]", plugin.Name())

	// Prepare the new configuration
	if err := plugin.PrepareReload(); err != nil {
		plugin.CancelReload()
		return err
	}

	// Apply the new configuration
	return plugin.ApplyReload()
}

// GetConfigPath returns the path to the plugin's configuration file
func (plugin *PluginForward) GetConfigPath() string {
	return plugin.configFile
}

// SetConfigWatcher sets the config watcher for this plugin
func (plugin *PluginForward) SetConfigWatcher(watcher *ConfigWatcher) {
	plugin.configWatcher = watcher
}

func (plugin *PluginForward) Eval(pluginsState *PluginsState, msg *dns.Msg) error {
	qName := pluginsState.qName
	qNameLen := len(qName)

	// Use read lock for thread-safe access to forwardMap
	plugin.rwLock.RLock()
	var sequence []SearchSequenceItem
	viaProxy := false
	for _, candidate := range plugin.forwardMap {
		candidateLen := len(candidate.domain)
		if candidateLen > qNameLen {
			continue
		}
		if (qName[qNameLen-candidateLen:] == candidate.domain &&
			(candidateLen == qNameLen || (qName[qNameLen-candidateLen-1] == '.'))) ||
			(candidate.domain == ".") {
			sequence = candidate.sequence
			viaProxy = candidate.viaProxy
			break
		}
	}
	plugin.rwLock.RUnlock()

	if len(sequence) == 0 {
		return nil
	}
	var err error
	var respMsg *dns.Msg
	tries := 4
	const resolvconfRetryInterval int64 = 30 // seconds

	for i := range sequence {
		var server string
		switch sequence[i].typ {
		case Explicit:
			server = sequence[i].servers[rand.Intn(len(sequence[i].servers))]
		case Bootstrap:
			server = plugin.bootstrapResolvers[rand.Intn(len(plugin.bootstrapResolvers))]
		case DHCP:
			const maxInconsistency = 9
			for _, dhcpdns := range plugin.dhcpdns {
				inconsistency, ip, dhcpDNS, err := dhcpdns.Status()
				if err != nil && ip != "" && inconsistency > maxInconsistency {
					dlog.Infof("No response from the DHCP server while resolving [%s]: %v", qName, err)
					continue
				}
				if len(dhcpDNS) > 0 {
					server = net.JoinHostPort(dhcpDNS[rand.Intn(len(dhcpDNS))].String(), "53")
					break
				}
			}
			if len(server) == 0 {
				dlog.Infof("DHCP didn't provide any DNS server to forward [%s]", qName)
				continue
			}
		case Resolvconf:
			if lastFail := sequence[i].rcLastFail.Load(); lastFail != 0 &&
				time.Now().Unix()-lastFail < resolvconfRetryInterval {
				continue
			}
			servers, warnings, err := parseResolvConf(sequence[i].resolvconf)
			if err != nil {
				dlog.Warnf(
					"Failed to open '%s' while resolving [%s]: %v",
					sequence[i].resolvconf, qName, err,
				)
				sequence[i].rcLastFail.Store(time.Now().Unix())
				continue
			}
			if len(servers) == 0 {
				for _, w := range warnings {
					dlog.Warn(w)
				}
				dlog.Warnf(
					"No valid nameservers in '%s' while resolving [%s]",
					sequence[i].resolvconf, qName,
				)
				sequence[i].rcLastFail.Store(time.Now().Unix())
				continue
			}
			sequence[i].rcLastFail.Store(0) // clear failure state on successful read
			nameserver := servers[rand.Intn(len(servers))]
			server, err = normalizeIPAndOptionalPort(nameserver, "53")
			if err != nil {
				dlog.Warnf(
					"Syntax error in address '%s' while resolving [%s]: %v",
					nameserver, qName, err,
				)
				continue
			}
		}
		pluginsState.serverName = server
		if len(server) == 0 {
			continue
		}

		if tries == 0 {
			break
		}
		tries--

		// Create a clean copy of the message without Extra section for forwarding
		forwardMsg := msg.Copy()
		forwardMsg.Extra = nil
		forwardMsg.Data = nil // Clear packed data so Exchange will re-pack without Extra

		if viaProxy {
			dlog.Debugf("Forwarding [%s] to [%s] using DNS-over-TCP through the configured proxy", qName, server)
			respMsg, err = plugin.exchangeViaProxy(forwardMsg, server, pluginsState.timeout)
		} else {
			dlog.Debugf("Forwarding [%s] to [%s]", qName, server)
			respMsg, err = plugin.exchangeDirect(forwardMsg, pluginsState.serverProto, server, pluginsState.timeout)
		}
		if err != nil {
			continue
		}
		if err := validateResponseQuestion(forwardMsg, respMsg); err != nil {
			continue
		}
		if !respMsg.Security {
			respMsg.AuthenticatedData = false
		}
		respMsg.ID = msg.ID
		pluginsState.synthResponse = respMsg
		pluginsState.action = PluginsActionSynth
		pluginsState.returnCode = PluginsReturnCodeForward
		if len(sequence) > 0 {
			switch respMsg.Rcode {
			case dns.RcodeNameError, dns.RcodeRefused, dns.RcodeNotAuth:
				continue
			}
		}
		return nil
	}
	return err
}

// exchangeDirect sends a query to the selected upstream server using the current
// protocol, retrying a truncated UDP response over TCP.
func (plugin *PluginForward) exchangeDirect(
	forwardMsg *dns.Msg,
	serverProto string,
	server string,
	timeout time.Duration,
) (*dns.Msg, error) {
	client := dns.Client{}
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	respMsg, _, err := client.Exchange(ctx, forwardMsg, serverProto, server)
	cancel()
	if err != nil && (respMsg == nil || !respMsg.Truncated) {
		return nil, err
	}
	if respMsg != nil && respMsg.Truncated {
		ctx, cancel = context.WithTimeout(context.Background(), timeout)
		respMsg, _, err = client.Exchange(ctx, forwardMsg, "tcp", server)
		cancel()
		if err != nil {
			return nil, err
		}
	}
	return respMsg, nil
}

// deadlineCappedConn caps every deadline set on the underlying connection, so helpers
// that reset deadlines cannot extend I/O beyond the query timeout.
type deadlineCappedConn struct {
	net.Conn
	deadline time.Time
}

func (c *deadlineCappedConn) capped(t time.Time) time.Time {
	if t.IsZero() || t.After(c.deadline) {
		return c.deadline
	}
	return t
}

func (c *deadlineCappedConn) SetDeadline(t time.Time) error {
	return c.Conn.SetDeadline(c.capped(t))
}

func (c *deadlineCappedConn) SetReadDeadline(t time.Time) error {
	return c.Conn.SetReadDeadline(c.capped(t))
}

func (c *deadlineCappedConn) SetWriteDeadline(t time.Time) error {
	return c.Conn.SetWriteDeadline(c.capped(t))
}

// exchangeViaProxy performs a DNS-over-TCP exchange through the proxy configured with the
// top-level `proxy` option. The connection is created per query and always closed here.
func (plugin *PluginForward) exchangeViaProxy(
	forwardMsg *dns.Msg,
	server string,
	timeout time.Duration,
) (*dns.Msg, error) {
	if plugin.proxyDialer == nil {
		return nil, errors.New("proxied forwarding requires the `proxy` option to be configured")
	}
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	var conn net.Conn
	var err error
	if contextDialer, ok := plugin.proxyDialer.(netproxy.ContextDialer); ok {
		conn, err = contextDialer.DialContext(ctx, "tcp", server)
	} else {
		conn, err = plugin.proxyDialer.Dial("tcp", server)
	}
	if err != nil {
		return nil, fmt.Errorf("unable to reach [%s] through the proxy: %w", server, err)
	}
	defer conn.Close()
	deadline, _ := ctx.Deadline()
	remaining := time.Until(deadline)
	if remaining <= 0 {
		return nil, fmt.Errorf("timeout while connecting to [%s] through the proxy", server)
	}
	// ExchangeWithConn dereferences the client transport and replaces the connection's
	// read and write deadlines with its own timeouts, so the transport must be non-nil
	// and the wrapper is what keeps the query timeout an end-to-end bound.
	conn = &deadlineCappedConn{Conn: conn, deadline: deadline}
	if err := conn.SetDeadline(deadline); err != nil {
		return nil, err
	}
	client := dns.Client{Transport: &dns.Transport{ReadTimeout: remaining, WriteTimeout: remaining}}
	respMsg, _, err := client.ExchangeWithConn(ctx, forwardMsg, conn)
	if err != nil {
		return nil, fmt.Errorf("proxied exchange with [%s] failed: %w", server, err)
	}
	return respMsg, nil
}

func parseResolvConf(filename string) (servers []string, warnings []string, err error) {
	data, err := os.ReadFile(filename)
	if err != nil {
		return nil, nil, err
	}
	for line := range strings.SplitSeq(string(data), "\n") {
		line = strings.TrimSpace(line)
		if !strings.HasPrefix(line, "nameserver") {
			continue
		}
		fields := strings.Fields(line)
		if len(fields) < 2 {
			continue
		}
		addr := fields[1]
		host := addr
		if h, _, err := net.SplitHostPort(addr); err == nil {
			host = h
		}
		if net.ParseIP(host) == nil {
			warnings = append(warnings, fmt.Sprintf(
				"Ignoring invalid nameserver address '%s' in [%s]", addr, filename,
			))
			continue
		}
		servers = append(servers, addr)
	}
	return
}

func normalizeIPAndOptionalPort(addr string, defaultPort string) (string, error) {
	var host, port string
	var err error

	if strings.HasPrefix(addr, "[") {
		if !strings.Contains(addr, "]:") {
			if addr[len(addr)-1] != ']' {
				return "", fmt.Errorf("invalid IPv6 format: missing closing ']'")
			}
			host = addr[1 : len(addr)-1]
			port = defaultPort
		} else {
			host, port, err = net.SplitHostPort(addr)
			if err != nil {
				return "", err
			}
		}
	} else {
		host, port, err = net.SplitHostPort(addr)
		if err != nil {
			host = addr
			port = defaultPort
		}
	}
	ip := net.ParseIP(host)
	if ip == nil {
		return "", fmt.Errorf("invalid IP address: [%s]", host)
	}
	if ip.To4() != nil {
		return fmt.Sprintf("%s:%s", ip.String(), port), nil
	}
	return fmt.Sprintf("[%s]:%s", ip.String(), port), nil
}
