package dataset

import (
	"strings"

	"github.com/miekg/dns"
)

type DomainTree struct {
	root *domainTreeNode
}

type domainTreeNode struct {
	leaf  map[string]*domainTreeNode
	isEnd bool
}

func NewDomainList(domains ...string) *DomainTree {
	tree := &DomainTree{
		root: &domainTreeNode{leaf: make(map[string]*domainTreeNode)},
	}
	for _, domain := range domains {
		tree.Add(domain)
	}
	return tree
}

func (tree *DomainTree) Add(domain string) {
	domain = normalizeDomain(domain)
	if domain == "" {
		return
	}

	labels := dns.SplitDomainName(domain)
	if len(labels) == 0 {
		return
	}

	node := tree.root
	for i := len(labels) - 1; i >= 0; i-- {
		label := labels[i]
		if label == "" {
			continue
		}
		if node.leaf == nil {
			node.leaf = make(map[string]*domainTreeNode)
		}
		if node.leaf[label] == nil {
			node.leaf[label] = &domainTreeNode{}
		}
		node = node.leaf[label]
	}
	node.isEnd = true
}

func (tree *DomainTree) Contains(domain string) bool {
	domain = normalizeDomain(domain)
	if domain == "" {
		return false
	}

	labels := dns.SplitDomainName(domain)
	if len(labels) == 0 {
		return false
	}

	node := tree.root
	for i := len(labels) - 1; i >= 0; i-- {
		if node.isEnd {
			return true
		}

		if node.leaf == nil {
			return false
		}

		label := labels[i]
		if node = node.leaf[label]; node == nil {
			return false
		}
	}
	return node.isEnd
}

func normalizeDomain(domain string) string {
	domain = strings.ToLower(strings.TrimSpace(domain))
	if domain == "" {
		return ""
	}

	// Remove trailing dot if present, dns.Fqdn will add it back properly
	domain = strings.TrimSuffix(domain, ".")

	if domain == "" {
		return ""
	}

	return dns.Fqdn(domain)
}