package resolve

import (
	"fmt"
	"io"
	"net/http"
	"regexp"
	"strconv"
	"sync"

	"github.com/rs/xid"
)

const (
	maxWildcardChecks = 3
)

// ResolutionPool is a pool of resolvers created for resolving urls
// for a given host.
type ResolutionPool struct {
	*Resolver
	Tasks          chan HostEntry
	Results        chan Result
	wg             *sync.WaitGroup
	removeWildcard bool

	wildcardIPs map[string]struct{}
}

// HostEntry defines a host with the source
type HostEntry struct {
	Host   string
	Source string
}

// Result contains the result for a host resolution
type Result struct {
	Type       ResultType
	Host       string
	IP         string
	Error      error
	Source     string
	StatusCode string
	UrlTitle   string
}

// ResultType is the type of result found
type ResultType int

// Types of data result can return
const (
	URL ResultType = iota
	Error
)

// NewResolutionPool creates a pool of resolvers for resolving urls of a given domain
func (r *Resolver) NewResolutionPool(workers int, removeWildcard bool) *ResolutionPool {
	resolutionPool := &ResolutionPool{
		Resolver:       r,
		Tasks:          make(chan HostEntry),
		Results:        make(chan Result),
		wg:             &sync.WaitGroup{},
		removeWildcard: removeWildcard,
		wildcardIPs:    make(map[string]struct{}),
	}

	go func() {
		for i := 0; i < workers; i++ {
			resolutionPool.wg.Add(1)
			go resolutionPool.resolveWorker()
		}
		resolutionPool.wg.Wait()
		close(resolutionPool.Results)
	}()

	return resolutionPool
}

// InitWildcards inits the wildcard ips array
func (r *ResolutionPool) InitWildcards(domain string) error {
	for i := 0; i < maxWildcardChecks; i++ {
		uid := xid.New().String()

		hosts, _ := r.DNSClient.Lookup(uid + "." + domain)
		if len(hosts) == 0 {
			return fmt.Errorf("%s is not a wildcard domain", domain)
		}

		// Append all wildcard ips found for domains
		for _, host := range hosts {
			r.wildcardIPs[host] = struct{}{}
		}
	}
	return nil
}

func GetStatus(domain string) (int, error) {
	resp, err := http.Get(fmt.Sprintf("%s", domain))
	if err != nil {
		return 0, err
	}
	defer resp.Body.Close()

	return resp.StatusCode, nil
}

var TitleRegexp = regexp.MustCompile("(?Uis)<title>(.*)</title>")

func GetTitle(domain string) (string, error) {
	resp, err := http.Get(fmt.Sprintf("%s", domain))
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	titleMatches := TitleRegexp.FindStringSubmatch(string(body))
	if len(titleMatches) > 1 {
		return titleMatches[1], nil
	}

	return "", nil
}

func (r *ResolutionPool) resolveWorker() {
	for task := range r.Tasks {
		//Get urls title
		title, err := GetTitle(task.Host)
		if err != nil {
			r.Results <- Result{Type: Error, Host: task.Host, Source: task.Source, Error: err}
			continue
		}
		//Get Url status code
		status, err := GetStatus(task.Host)
		if err != nil {
			r.Results <- Result{Type: Error, Host: task.Host, Source: task.Source, Error: err}
			continue
		}
		r.Results <- Result{
			Type:       URL,
			Host:       task.Host,
			UrlTitle:   title,
			StatusCode: strconv.Itoa(status),
			Source:     task.Source,
		}
	}
	r.wg.Done()
}
