package shodan

import (
	"fmt"
	"os"

	"github.com/panjf2000/ants/v2"

	"git.tcp.direct/perp/shogo/internal/shodan/alert"
	"git.tcp.direct/perp/shogo/internal/shodan/scan"
	"git.tcp.direct/perp/shogo/internal/utils"
)

// Manage a goroutine pool
func Pool(function string, flags *Flags, alert *alert.Alert, scan *scan.Scan) {
	// Results channel
	results := make(chan string, len(flags.Args)-1)

	// Thread count
	var threads int

	// Get thread count
	if flags.Threads == 0 {
		threads = 1
	} else {
		threads = flags.Threads
	}

	// Goroutine pool
	pool, err := ants.NewPool(threads)
	if err != nil {
		fmt.Printf("%s: %s\n", utils.Red("Error"), err.Error())
		os.Exit(1)
	}
	defer pool.Release()

	// Go through flags
	for a := 0; a <= len(flags.Args)-1; a++ {
		// Shodan manager
		shodan := &Shodan{
			Flags:   flags,
			Results: results,
		}
		shodan.Flags.Query = flags.Args[a] // Bug

		// Scan manager
		if scan != nil {
			shodan.Scan = scan
			shodan.Scan.Results = results
		}

		// Alert manager
		if alert != nil {
			shodan.Alert = alert
			shodan.Alert.ID = flags.Args[a] // Bug
			shodan.Alert.Results = results
		}

		// Submit job
		pool.Submit(func() {
			switch function {
			case "alert clear":
				go shodan.Alert.Clear()
			case "alert create":
				go shodan.Alert.Create()
			case "alert delete":
				go shodan.Alert.Delete()
			case "alert info":
				go shodan.Alert.Info()
			case "alert list":
				go shodan.Alert.List()

			case "scan internet":
				go shodan.Scan.Internet()
			case "scan protocols":
				go shodan.Scan.Protocols()
			case "scan submit":
				go shodan.Scan.Submit()

			case "count":
				go shodan.Count()
			case "host":
				go shodan.Host()
			case "info":
				go shodan.Info()
			case "init":
				go shodan.Init()
			case "myip":
				go shodan.MyIp()
			case "search":
				go shodan.Search()
			case "stats":
				go shodan.Stats()
			}
		})
	}

	// Get result
	for j := 0; j <= len(flags.Args)-1; j++ {
		result := <-results
		// Multiple results needs newline
		if j != len(flags.Args)-1 {
			fmt.Println(result + "\n")
		} else {
			fmt.Println(result)
		}
	}
	defer close(results)
}