package runner

import (
	"time"

	"git.supernets.org/perp/blink/v1/pkg/dns"
	mdns "github.com/miekg/dns"
	"github.com/panjf2000/ants/v2"
)

// Bruteforce runner
type Runner struct {
	options *Options         // CLI options
	client  *mdns.Client     // DNS client
	pool    *ants.Pool       // Goroutine pool
	results chan *dns.Result // Results channel
}

// Return a new Runner
func New(options *Options) *Runner {
	// Create client
	client := &mdns.Client{
		Net:     "tcp",
		Timeout: time.Second * time.Duration(options.Timeout),
	}

	// Check UDP
	if options.UDP {
		client.Net = "udp"
	}

	// Create pool
	pool, err := ants.NewPool(options.Threads)
	if err != nil {
		panic(err)
	}

	// Store task count
	tasks := len(options.Wordlist) * len(options.Domains)

	// Create channel
	results := make(chan *dns.Result, tasks)

	return &Runner{
		options: options,
		client:  client,
		pool:    pool,
		results: results,
	}
}

// Detect wildcard
func (r *Runner) Wildcard(domain string) bool {
	return dns.Wildcard(r.client, r.options.Resolvers, domain)
}

// Submit tasks into pool
func (r *Runner) Submit() {
	// Go through wordlist
	for _, word := range r.options.Wordlist {
		// Go through domains
		for _, domain := range r.options.Domains {
			// Create query
			query := &dns.Query{
				Client:    r.client,
				IPv6:      r.options.IPv6,
				Resolvers: r.options.Resolvers,
				Domain:    domain,
				Subdomain: word + "." + domain,
				Results:   r.results,
			}

			// Submit query
			r.pool.Submit(func() {
				query.Lookup()
			})
		}
	}
}

// Receive tasks from pool
func (r *Runner) Receive() {
	// Store task count
	tasks := len(r.options.Wordlist) * len(r.options.Domains)

	// Go through tasks
	for range tasks {
		select {
		case result := <-r.results:
			// Send result
			r.options.OnResult(result)
		}
	}
}

// Start runner
func (r *Runner) Start() {
	go r.Submit()
	r.Receive()
}