package dns

import (
	"context"
	"fmt"
	"net"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	miekgDNS "github.com/miekg/dns"
	"github.com/spf13/viper"
)

type ResponseWriter = miekgDNS.ResponseWriter
type Msg = miekgDNS.Msg
type HandlerFunc = miekgDNS.HandlerFunc

// DNS struct represents the DNS server with its configuration, server, mux, records, and resolver.
type DNS struct {
	Config   Config
	Server   Server
	Mux      Mux
	Records  Records
	Resolver Resolver
}

// Config struct holds the configuration settings.
type Config struct {
	*viper.Viper
}

// Server struct wraps the miekgDNS.Server.
type Server struct {
	*miekgDNS.Server
}

// Mux struct wraps the miekgDNS.ServeMux.
type Mux struct {
	*miekgDNS.ServeMux
}

// Records struct holds the DNS records.
type Records struct {
	TypeA  map[string]string
	TypeNS map[string]string // Add NS records
}

// Resolver struct wraps the net.Resolver.
type Resolver struct {
	*net.Resolver
}

// New initializes a new DNS instance with the given configuration file path.
func New(filePath string) *DNS {
	app := &DNS{}
	app.ConfigInit(filePath)                                                              // Initialize configuration
	app.Records.TypeA = app.Config.GetStringMapString("records.type_a")                   // Load Type A records
	alternate_resolver_ip := app.Config.GetString("alternate_resolver.ip")                // Get alternate resolver IP
	alternate_resolver_port := strconv.Itoa(app.Config.GetInt("alternate_resolver.port")) // Get alternate resolver port
	app.ServerInit()                                                                      // Initialize server
	app.MuxInit()                                                                         // Initialize multiplexer
	app.Resolver = NewResolver(alternate_resolver_ip + ":" + alternate_resolver_port)     // Initialize resolver
	app.Server.Handler = app.Mux                                                          // Set handler

	return app
}

// Run starts the DNS server.
func (a *DNS) Run() {
	err := a.Server.ListenAndServe() // Start server
	if err != nil {
		fmt.Println(err.Error()) // Print error if any
	}
}

// ConfigInit initializes the configuration from the given file path.
func (a *DNS) ConfigInit(filePath string) {
	dir, file := filepath.Split(filePath)  // Split file path
	fe := filepath.Ext(file)               // Get file extension
	ext, _ := strings.CutPrefix(fe, ".")   // Remove leading dot from extension
	name, _ := strings.CutSuffix(file, fe) // Get file name without extension

	a.Config.Viper = viper.New() // Create new Viper instance
	a.Config.AddConfigPath(dir)  // Add config path
	a.Config.SetConfigName(name) // Set config name
	a.Config.SetConfigType(ext)  // Set config type
	a.Config.ReadInConfig()      // Read config
}

// ServerInit initializes the DNS server with the configuration settings.
func (a *DNS) ServerInit() {
	port := a.Config.GetInt("port") // Get port from config
	a.Server.Server = &miekgDNS.Server{
		Addr:    ":" + strconv.Itoa(port),      // Set server address
		Net:     a.Config.GetString("network"), // Set network type
		Handler: nil,                           // Handler will be set later
	}
}

// MuxInit initializes the DNS request multiplexer.
func (a *DNS) MuxInit() {
	a.Mux.ServeMux = miekgDNS.NewServeMux() // Create new ServeMux
	a.Mux.HandleFunc(".", a.HandleTypeA)    // Handle all requests with HandleTypeA
	a.Mux.HandleFunc(".", a.HandleTypeNS)   // Handle all requests with HandleTypeNS
}

// NewResolver creates a new DNS resolver with the given DNS server address.
func NewResolver(DNSserverAddr string) Resolver {
	return Resolver{
		&net.Resolver{
			PreferGo: true, // Prefer Go's resolver
			Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
				d := net.Dialer{Timeout: 5 * time.Second}       // Set dial timeout
				return d.DialContext(ctx, "udp", DNSserverAddr) // Dial DNS server
			},
		},
	}
}

// Lookup performs a DNS lookup for the given address using the resolver.
func (r *Resolver) Lookup(lookupAddr string) string {
	var resp []string
	var err error

	ctx := context.Background()               // Create background context
	resp, err = r.LookupHost(ctx, lookupAddr) // Perform lookup

	if err != nil {
		fmt.Println(err.Error()) // Print error if any
	}

	return resp[0] // Return first result
}