// Copyright 2023 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// Siso is a Ninja-compatible build system optimized for remote execution.
package main

import (
	"context"
	"flag"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"runtime"
	"runtime/pprof"
	"runtime/trace"
	"strings"
	"sync/atomic"
	"syscall"
	"time"

	log "github.com/golang/glog"
	"github.com/google/subcommands"

	"go.chromium.org/build/siso/auth/cred"
	"go.chromium.org/build/siso/subcmd/auth"
	"go.chromium.org/build/siso/subcmd/collector"
	"go.chromium.org/build/siso/subcmd/fetch"
	"go.chromium.org/build/siso/subcmd/fscmd"
	"go.chromium.org/build/siso/subcmd/isolate"
	"go.chromium.org/build/siso/subcmd/metricscmd"
	"go.chromium.org/build/siso/subcmd/ninja"
	"go.chromium.org/build/siso/subcmd/ninjafrontend"
	"go.chromium.org/build/siso/subcmd/proxy"
	"go.chromium.org/build/siso/subcmd/ps"
	"go.chromium.org/build/siso/subcmd/query"
	"go.chromium.org/build/siso/subcmd/recall"
	"go.chromium.org/build/siso/subcmd/report"
	"go.chromium.org/build/siso/subcmd/sandbox"
	"go.chromium.org/build/siso/subcmd/scandeps"
	"go.chromium.org/build/siso/subcmd/spawnhelper"
	"go.chromium.org/build/siso/subcmd/version"
	"go.chromium.org/build/siso/subcmd/webui"
	"go.chromium.org/build/siso/ui"

	_ "net/http/pprof" // import to let pprof register its HTTP handlers
)

var (
	pprofAddr          string
	cpuprofile         string
	memprofile         string
	memprofileInterval time.Duration
	memprofileGC       bool
	mutexprofile       string
	blockprofile       string
	blockprofRate      int
	mutexprofFrac      int
	traceFile          string
)

const versionID = "v1.5.21"
const versionStr = "siso " + versionID

func main() {
	// Wraps sisoMain() because os.Exit() doesn't wait defers.
	os.Exit(sisoMain())
}

func sisoMain() int {
	flag.CommandLine.Usage = func() {
		w := flag.CommandLine.Output()
		fmt.Fprint(w, `
Usage: siso [flags] [command] [arguments]

e.g.
 $ siso ninja -C out/Default

`)
		fmt.Fprintf(w, "important flags of %s:\n", os.Args[0])

		f := flag.Lookup("credential_helper")
		fmt.Fprintf(w, `  -credential_helper path
    %s
    (default %q)
`, f.Usage, f.DefValue)
		f = flag.Lookup("version")
		fmt.Fprintf(w, `  -version
   %s
`, f.Usage)
		fmt.Fprintf(flag.CommandLine.Output(), `
Use "siso help" to display commands.
Use "siso help [command]" for more information about a command.
Use "siso flags" to display all flags.
`)
	}

	flag.StringVar(&pprofAddr, "pprof_addr", "", `listen address for "go tool pprof". e.g. "localhost:6060"`)
	flag.StringVar(&cpuprofile, "cpuprofile", "", "write cpu profile to this file")
	flag.StringVar(&memprofile, "memprofile", "", "write memory profile to this file")
	flag.DurationVar(&memprofileInterval, "memprofile_interval", 0, `if non-zero, dump heap profiles at this interval to rolling files derived from -memprofile (e.g. "siso-memory.0001.pprof"); the bare -memprofile path is not written.`)
	flag.BoolVar(&memprofileGC, "memprofile_gc", true, "run runtime.GC() before each heap profile dump for a cleaner inuse view")
	flag.StringVar(&mutexprofile, "mutexprofile", "", "write mutex profile to this file")
	flag.StringVar(&blockprofile, "blockprofile", "", "write block profile to this file")
	flag.IntVar(&blockprofRate, "blockprof_rate", 0, "block profile rate")
	flag.IntVar(&mutexprofFrac, "mutexprof_frac", 0, "mutex profile fraction")
	flag.StringVar(&traceFile, "trace", "", `go trace output for "go tool trace"`)

	credHelper := cred.DefaultCredentialHelper()
	flag.StringVar(&credHelper, "credential_helper", credHelper, `path to a credential helper.
    see https://github.com/EngFlow/credential-helper-spec/blob/main/spec.md
    "luci-auth" uses luci-auth.
    "gcloud" uses gcloud.
    "google-application-default" or "" uses Google Application Default Credentials.
    "mTLS" disables per RPC credentials.
    environment variable SISO_CREDENTIAL_HELPER sets default value.`)

	var printVersion bool
	flag.BoolVar(&printVersion, "version", false, "print version")

	// set stderrthreshold to FATAL by default
	// but could be overridden by command line by flag.Parse later.
	if err := flag.Set("stderrthreshold", "FATAL"); err != nil {
		log.Exitf("failed to set --stderrthreshold=FATAL: %v\n", err)
	}
	flag.Parse()

	ctx := context.Background()
	// Flush the log on exit to not lose any messages.
	defer log.Flush()

	// Print a stack trace when a panic occurs.
	defer func() {
		if r := recover(); r != nil {
			const size = 64 << 10
			buf := make([]byte, size)
			buf = buf[:runtime.Stack(buf, false)]
			log.Fatalf("panic: %v\n%s", r, buf)
		}
	}()

	authOpts := cred.AuthOpts(credHelper)
	if printVersion {
		return int(version.Cmd(versionStr, authOpts).Execute(ctx, flag.CommandLine))
	}
	if blockprofile != "" && blockprofRate == 0 {
		blockprofRate = 1
	}
	if mutexprofile != "" && mutexprofFrac == 0 {
		mutexprofFrac = 1
	}
	if blockprofRate > 0 {
		runtime.SetBlockProfileRate(blockprofRate)
	}
	if mutexprofFrac > 0 {
		runtime.SetMutexProfileFraction(mutexprofFrac)
	}

	// Start an HTTP server that can be used to profile Siso during runtime.
	if pprofAddr != "" {
		// https://pkg.go.dev/net/http/pprof
		fmt.Fprintf(os.Stderr, "pprof is enabled, listening at http://%s/debug/pprof/\n", pprofAddr)
		go func() {
			log.Infof("pprof http listener: %v", http.ListenAndServe(pprofAddr, nil))
		}()
		defer func() {
			fmt.Fprintf(os.Stderr, "pprof is still listening at http://%s/debug/pprof/\n", pprofAddr)
			fmt.Fprintln(os.Stderr, "Press Ctrl-C to terminate the process")
			sigch := make(chan os.Signal, 1)
			signal.Notify(sigch, os.Interrupt, syscall.SIGTERM)
			<-sigch
		}()
	}

	// Save a CPU profile to disk on exit.
	if cpuprofile != "" {
		f, err := os.Create(cpuprofile)
		if err != nil {
			log.Fatalf("failed to create cpuprofile file: %v", err)
		}
		defer func() {
			if err := f.Close(); err != nil {
				log.Errorf("failed to close cpuprofile file: %v", err)
			}
		}()
		err = pprof.StartCPUProfile(f)
		if err != nil {
			log.Fatalf("failed to start CPU profiler: %v", err)
		}
		defer pprof.StopCPUProfile()
	}

	// Save a heap profile to disk on exit, or periodically to rolling files.
	if memprofile != "" {
		if memprofileInterval > 0 {
			stop := startPeriodicMemprofile(memprofile, memprofileInterval, memprofileGC)
			defer stop()
		} else {
			f, err := os.Create(memprofile)
			if err != nil {
				log.Fatalf("failed to create memprofile file: %v", err)
			}
			defer func() {
				if memprofileGC {
					runtime.GC()
				}
				if err := pprof.WriteHeapProfile(f); err != nil {
					log.Errorf("failed to write heap profile: %v", err)
				}
				if err := f.Close(); err != nil {
					log.Errorf("failed to close memprofile file: %v", err)
				}
			}()
		}
	}

	// Save a mutex profile to disk on exit.
	if mutexprofile != "" {
		f, err := os.Create(mutexprofile)
		if err != nil {
			log.Fatalf("failed to create mutexprofile file: %v", err)
		}
		defer func() {
			if err := pprof.Lookup("mutex").WriteTo(f, 0); err != nil {
				log.Errorf("failed to write mutex profile: %v", err)
			}
			if err := f.Close(); err != nil {
				log.Errorf("failed to close mutexprofile file: %v", err)
			}
		}()
	}

	// Save a block profile to disk on exit.
	if blockprofile != "" {
		f, err := os.Create(blockprofile)
		if err != nil {
			log.Fatalf("failed to create blockprofile file: %v", err)
		}
		defer func() {
			if err := pprof.Lookup("block").WriteTo(f, 0); err != nil {
				log.Errorf("failed to write block profile: %v", err)
			}
			if err := f.Close(); err != nil {
				log.Errorf("failed to close blockprofile file: %v", err)
			}
		}()
	}

	// Save a go trace to disk during execution.
	if traceFile != "" {
		fmt.Fprintf(os.Stderr, "enable go trace in %q\n", traceFile)
		f, err := os.Create(traceFile)
		if err != nil {
			log.Fatalf("Failed to create go trace output file: %v", err)
		}
		defer func() {
			fmt.Fprintf(os.Stderr, "go trace: go tool trace %s\n", traceFile)
			if err := f.Close(); err != nil {
				log.Errorf("Failed to close go trace output file: %v", err)
			}
		}()
		if err := trace.Start(f); err != nil {
			log.Fatalf("Failed to start go trace: %v", err)
		}
		defer trace.Stop()
	}

	// Initialize the UI and ensure we restore the state of the terminal upon exit.
	ui.Init()
	defer ui.Restore()

	subcommands.Register(ninja.Cmd(authOpts, versionID), "")

	subcommands.Register(recall.Cmd(authOpts), "reapi")
	subcommands.Register(fetch.Cmd(authOpts), "reapi")
	subcommands.Register(isolate.Cmd(authOpts), "reapi")
	subcommands.Register(proxy.Cmd(authOpts), "reapi")

	subcommands.Register(fscmd.Cmd(authOpts), "investigation")
	subcommands.Register(metricscmd.Cmd(), "investigation")
	subcommands.Register(ps.Cmd(), "investigation")
	subcommands.Register(query.Cmd(), "investigation")
	subcommands.Register(report.Cmd(), "investigation")
	subcommands.Register(webui.Cmd(versionID), "investigation")

	subcommands.Register(auth.CheckCmd(authOpts), "auth")
	subcommands.Register(auth.LoginCmd(authOpts), "auth")
	subcommands.Register(auth.LogoutCmd(authOpts), "auth")

	subcommands.Register(ninjafrontend.Cmd(), "debugging")
	subcommands.Register(sandbox.Cmd(), "debugging")
	subcommands.Register(scandeps.Cmd(), "debugging")

	subcommands.Register(collector.Cmd(authOpts, versionID), "internal-helper")
	subcommands.Register(spawnhelper.Cmd(), "internal-helper")

	subcommands.Register(subcommands.FlagsCommand(), "command-help")
	subcommands.Register(subcommands.HelpCommand(), "command-help")
	subcommands.Register(version.Cmd(versionStr, authOpts), "command-help")

	return int(subcommands.Execute(ctx))
}

// startPeriodicMemprofile launches a goroutine that writes a heap profile to a
// rolling file every interval. The base path's extension (if any) is preserved
// after a 4-digit sequence number, so "/tmp/m.pprof" becomes
// "/tmp/m.0001.pprof", "/tmp/m.0002.pprof", and so on. The returned function
// stops the goroutine and writes one final dump.
func startPeriodicMemprofile(base string, interval time.Duration, gc bool) func() {
	dir, file := filepath.Split(base)
	ext := filepath.Ext(file)
	stem := strings.TrimSuffix(file, ext)
	var seq atomic.Uint64
	dump := func() {
		n := seq.Add(1)
		path := filepath.Join(dir, fmt.Sprintf("%s.%04d%s", stem, n, ext))
		f, err := os.Create(path)
		if err != nil {
			log.Errorf("failed to create memprofile file %q: %v", path, err)
			return
		}
		defer func() {
			if err := f.Close(); err != nil {
				log.Errorf("failed to close memprofile file %q: %v", path, err)
			}
		}()
		if gc {
			runtime.GC()
		}
		if err := pprof.WriteHeapProfile(f); err != nil {
			log.Errorf("failed to write heap profile to %q: %v", path, err)
		}
	}
	stopCh := make(chan struct{})
	doneCh := make(chan struct{})
	go func() {
		defer close(doneCh)
		t := time.NewTicker(interval)
		defer t.Stop()
		for {
			select {
			case <-stopCh:
				return
			case <-t.C:
				dump()
			}
		}
	}()
	return func() {
		close(stopCh)
		<-doneCh
		dump()
	}
}
