package scan

import (
	"strconv"
	"sync"
	"sync/atomic"
	"time"
)

type Result struct {
	Text string
}

type Progress struct {
	Done  int
	Total int
}

type Dispatcher struct {
	IPs      []string
	Ports    []int
	Timeout  time.Duration
	Workers  int
	PingOnly bool

	OnResult   func(Result)
	OnProgress func(Progress)
	OnFinish   func()

	stop    chan struct{}
	done    sync.WaitGroup
	counter int64
}

func (d *Dispatcher) Start() {
	d.stop = make(chan struct{})
	tasks := d.buildTasks()
	total := len(tasks)
	if total == 0 {
		if d.OnFinish != nil {
			d.OnFinish()
		}
		return
	}

	workers := d.Workers
	if workers <= 0 {
		workers = 1
	}

	taskCh := make(chan task)
	for i := 0; i < workers; i++ {
		d.done.Add(1)
		go d.worker(taskCh, total)
	}

feed:
	for _, t := range tasks {
		select {
		case <-d.stop:
			break feed
		case taskCh <- t:
		}
	}
	close(taskCh)
	d.done.Wait()
	if d.OnFinish != nil {
		d.OnFinish()
	}
}

func (d *Dispatcher) worker(taskCh <-chan task, total int) {
	defer d.done.Done()

	var pinger *Pinger
	if d.PingOnly {
		pinger, _ = NewPinger()
		defer pinger.Close()
	}

	for task := range taskCh {
		select {
		case <-d.stop:
			return
		default:
		}

		ok := false
		if task.ping {
			if pinger != nil {
				ok = pinger.Ping(task.ip, d.Timeout)
			}
		} else {
			ok = TCPProbe(task.ip, task.port, d.Timeout)
		}

		if ok && d.OnResult != nil {
			d.OnResult(Result{Text: task.label})
		}

		done := int(atomic.AddInt64(&d.counter, 1))
		if d.OnProgress != nil {
			d.OnProgress(Progress{Done: done, Total: total})
		}
	}
}

func (d *Dispatcher) Stop() {
	select {
	case <-d.stop:
	default:
		close(d.stop)
	}
}

type task struct {
	ip    string
	port  int
	label string
	ping  bool
}

func (d *Dispatcher) buildTasks() []task {
	var tasks []task
	if d.PingOnly {
		for _, ip := range d.IPs {
			tasks = append(tasks, task{ip: ip, label: ip, ping: true})
		}
	} else {
		for _, ip := range d.IPs {
			for _, port := range d.Ports {
				tasks = append(tasks, task{
					ip:    ip,
					port:  port,
					label: ip + ":" + strconv.Itoa(port),
				})
			}
		}
	}
	return tasks
}