// Copyright 2019 Google LLC.
// Use of this source code is governed by a BSD-style license that can be found in the LICENSE file.
package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"log"
	"math/rand"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

// Too many GPU processes and we'll start to overwhelm your GPU,
// even hanging your machine in the worst case.  Here's a reasonable default.
func defaultGpuLimit() int {
	limit := 8
	if n := runtime.NumCPU(); n < limit {
		return n
	}
	return limit
}

var script = flag.String("script", "", "A file with jobs to run, one per line. - for stdin.")
var random = flag.Bool("random", true, "Assign sources into job batches randomly?")
var quiet = flag.Bool("quiet", false, "Print only failures?")
var exact = flag.Bool("exact", false, "Match GM names only exactly.")
var cpuLimit = flag.Int("cpuLimit", runtime.NumCPU(),
	"Maximum number of concurrent processes for CPU-bound work.")
var gpuLimit = flag.Int("gpuLimit", defaultGpuLimit(),
	"Maximum number of concurrent processes for GPU-bound work.")

func init() {
	flag.StringVar(script, "s", *script, "Alias for --script.")
	flag.BoolVar(random, "r", *random, "Alias for --random.")
	flag.BoolVar(quiet, "q", *quiet, "Alias for --quiet.")
	flag.BoolVar(exact, "e", *exact, "Alias for --exact.")
	flag.IntVar(cpuLimit, "c", *cpuLimit, "Alias for --cpuLimit.")
	flag.IntVar(gpuLimit, "g", *gpuLimit, "Alias for --gpuLimit.")
}

func listAllGMs(fm string) (gms []string, err error) {
	// Query fm binary for list of all available GMs by running with no arguments.
	cmd := exec.Command(fm)
	stdout, err := cmd.Output()
	if err != nil {
		return
	}
	// GM names are listed line-by-line.
	scanner := bufio.NewScanner(bytes.NewReader(stdout))
	for scanner.Scan() {
		gms = append(gms, scanner.Text())
	}
	err = scanner.Err()
	return
}

type work struct {
	Sources []string
	Flags   []string
}

func parseWork(args []string, gms []string) (*work, error) {
	w := &work{}
	for _, arg := range args {
		// I wish we could parse flags here too, but it's too late.
		if strings.HasPrefix(arg, "-") {
			msg := "Is '%s' an fm flag? If so please pass it using flag=value syntax."
			if flag.Lookup(arg[1:]) != nil {
				msg = "Please pass fm_bot flags like '%s' on the command line before the FM binary."
			}
			return nil, fmt.Errorf(msg, arg)
		}

		// Everything after a # is a comment.
		if strings.HasPrefix(arg, "#") {
			break
		}

		// Treat "gm" or "gms" as a shortcut for all known GMs.
		if arg == "gm" || arg == "gms" {
			w.Sources = append(w.Sources, gms...)
			continue
		}

		// Is this an option to pass through to fm?
		if parts := strings.Split(arg, "="); len(parts) == 2 {
			f := "-"
			if len(parts[0]) > 1 {
				f += "-"
			}
			f += parts[0]

			w.Flags = append(w.Flags, f, parts[1])
			continue
		}

		// Is this argument naming a GM?
		matchedAnyGM := false
		for _, gm := range gms {
			if (*exact && gm == arg) || (!*exact && strings.Contains(gm, arg)) {
				w.Sources = append(w.Sources, gm)
				matchedAnyGM = true
			}
		}
		if matchedAnyGM {
			continue
		}

		// Anything left ought to be on the file system: a file, a directory, or a glob.
		// Not all shells expand globs, so we'll do it here just in case.
		matches, err := filepath.Glob(arg)
		if err != nil {
			return nil, err
		}
		if len(matches) == 0 {
			return nil, fmt.Errorf("Don't understand '%s'.", arg)
		}

		for _, match := range matches {
			err := filepath.Walk(match, func(path string, info os.FileInfo, err error) error {
				if !info.IsDir() {
					w.Sources = append(w.Sources, path)
				}
				return err
			})
			if err != nil {
				return nil, err
			}
		}
	}
	return w, nil
}

func main() {
	flag.Parse()

	if flag.NArg() < 1 {
		log.Fatal("Please pass an fm binary as the first argument.")
	}
	fm := flag.Args()[0]

	gms, err := listAllGMs(fm)
	if err != nil {
		log.Fatalln("Could not query", fm, "for GMs:", err)
	}

	// One job can comes right on the command line,
	// and any number can come one per line from -script.
	jobs := [][]string{flag.Args()[1:]}
	if *script != "" {
		file := os.Stdin
		if *script != "-" {
			file, err = os.Open(*script)
			if err != nil {
				log.Fatal(err)
			}
			defer file.Close()
		}

		scanner := bufio.NewScanner(file)
		for scanner.Scan() {
			jobs = append(jobs, strings.Fields(scanner.Text()))
		}
		if err = scanner.Err(); err != nil {
			log.Fatal(err)
		}
	}

	wg := &sync.WaitGroup{}
	var failures int32 = 0

	worker := func(queue chan work) {
		for w := range queue {
			start := time.Now()

			args := w.Flags[:]
			args = append(args, "-s")
			args = append(args, w.Sources...)

			cmd := exec.Command(fm, args...)
			output, err := cmd.CombinedOutput()

			status := "#done"
			if err != nil {
				status = fmt.Sprintf("#failed (%v)", err)

				if len(w.Sources) == 1 {
					// If a source ran alone and failed, that's just a failure.
					atomic.AddInt32(&failures, 1)
				} else {
					// If a batch of sources ran and failed, split them up and try again.
					for _, source := range w.Sources {
						wg.Add(1)
						queue <- work{[]string{source}, w.Flags}
					}
				}
			}

			if !*quiet || (err != nil && len(w.Sources) == 1) {
				log.Printf("\n%v %v in %v:\n%s",
					strings.Join(cmd.Args, " "), status, time.Since(start), output)
			}

			wg.Done()
		}
	}

	cpu := make(chan work, 1<<20)
	for i := 0; i < *cpuLimit; i++ {
		go worker(cpu)
	}

	gpu := make(chan work, 1<<20)
	for i := 0; i < *gpuLimit; i++ {
		go worker(gpu)
	}

	for _, job := range jobs {
		// Skip blank lines, empty command lines.
		if len(job) == 0 {
			continue
		}

		w, err := parseWork(job, gms)
		if err != nil {
			log.Fatal(err)
		}

		// Determine if this is CPU-bound or GPU-bound work, conservatively assuming GPU.
		queue, limit := gpu, *gpuLimit
		backend := ""
		for i, flag := range w.Flags {
			if flag == "-b" || flag == "--backend" {
				backend = w.Flags[i+1]
			}
		}
		whitelisted := map[string]bool{
			"cpu": true,
			"skp": true,
			"pdf": true,
		}
		if whitelisted[backend] {
			queue, limit = cpu, *cpuLimit
		}

		if *random {
			rand.Shuffle(len(w.Sources), func(i, j int) {
				w.Sources[i], w.Sources[j] = w.Sources[j], w.Sources[i]
			})
		}

		// Round up so there's at least one source per batch.
		sourcesPerBatch := (len(w.Sources) + limit - 1) / limit

		for i := 0; i < len(w.Sources); i += sourcesPerBatch {
			end := i + sourcesPerBatch
			if end > len(w.Sources) {
				end = len(w.Sources)
			}
			batch := w.Sources[i:end]

			wg.Add(1)
			queue <- work{batch, w.Flags}
		}
	}

	wg.Wait()

	if failures > 0 {
		log.Fatalln(failures, "failures after retries")
	}
}