// Copyright The OpenTelemetry Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package testbed

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"sync"
	"syscall"
	"time"

	"github.com/shirou/gopsutil/cpu"
	"github.com/shirou/gopsutil/process"
	"go.uber.org/atomic"
)

// ResourceSpec is a resource consumption specification.
type ResourceSpec struct {
	// Percentage of one core the process is expected to consume at most.
	// Test is aborted and failed if consumption during
	// ResourceCheckPeriod exceeds this number. If 0 the CPU
	// consumption is not monitored and does not affect the test result.
	ExpectedMaxCPU uint32

	// Maximum RAM in MiB the process is expected to consume.
	// Test is aborted and failed if consumption exceeds this number.
	// If 0 memory consumption is not monitored and does not affect
	// the test result.
	ExpectedMaxRAM uint32

	// Period during which CPU and RAM of the process are measured.
	// Bigger numbers will result in more averaging of short spikes.
	ResourceCheckPeriod time.Duration
}

// isSpecified returns true if any part of ResourceSpec is specified,
// i.e. has non-zero value.
func (rs *ResourceSpec) isSpecified() bool {
	return rs != nil && (rs.ExpectedMaxCPU != 0 || rs.ExpectedMaxRAM != 0)
}

// ChildProcess implements the OtelcolRunner interface as a child process on the same machine executing
//the test. The process can be monitored and the output of which will be written to a log file.
type ChildProcess struct {
	// Descriptive name of the process
	name string

	// Config file name
	configFileName string

	// Command to execute
	cmd *exec.Cmd

	// WaitGroup for copying process output
	outputWG sync.WaitGroup

	// Various starting/stopping flags
	isStarted  bool
	stopOnce   sync.Once
	isStopped  bool
	doneSignal chan struct{}

	// Resource specification that must be monitored for.
	resourceSpec *ResourceSpec

	// Process monitoring data.
	processMon *process.Process

	// Time when process was started.
	startTime time.Time

	// Last tick time we monitored the process.
	lastElapsedTime time.Time

	// Process times that were fetched on last monitoring tick.
	lastProcessTimes *cpu.TimesStat

	// Current RAM RSS in MiBs
	ramMiBCur atomic.Uint32

	// Current CPU percentage times 1000 (we use scaling since we have to use int for atomic operations).
	cpuPercentX1000Cur atomic.Uint32

	// Maximum CPU seen
	cpuPercentMax float64

	// Number of memory measurements
	memProbeCount int

	// Cumulative RAM RSS in MiBs
	ramMiBTotal uint64

	// Maximum RAM seen
	ramMiBMax uint32
}

type StartParams struct {
	name         string
	logFilePath  string
	cmd          string
	cmdArgs      []string
	resourceSpec *ResourceSpec
}

type ResourceConsumption struct {
	CPUPercentAvg float64
	CPUPercentMax float64
	RAMMiBAvg     uint32
	RAMMiBMax     uint32
}

func (cp *ChildProcess) PrepareConfig(configStr string) (configCleanup func(), err error) {
	configCleanup = func() {
		// NoOp
	}
	var file *os.File
	file, err = ioutil.TempFile("", "agent*.yaml")
	if err != nil {
		log.Printf("%s", err)
		return configCleanup, err
	}

	defer func() {
		errClose := file.Close()
		if errClose != nil {
			log.Printf("%s", errClose)
		}
	}()

	if _, err = file.WriteString(configStr); err != nil {
		log.Printf("%s", err)
		return configCleanup, err
	}
	cp.configFileName = file.Name()
	configCleanup = func() {
		os.Remove(cp.configFileName)
	}
	return configCleanup, err
}

// start a child process.
//
// Parameters:
// name is the human readable name of the process (e.g. "Agent"), used for logging.
// logFilePath is the file path to write the standard output and standard error of
// the process to.
// cmd is the executable to run.
// cmdArgs is the command line arguments to pass to the process.
func (cp *ChildProcess) Start(params StartParams) (receiverAddr string, err error) {

	cp.name = params.name
	cp.doneSignal = make(chan struct{})
	cp.resourceSpec = params.resourceSpec

	log.Printf("Starting %s (%s)", cp.name, params.cmd)

	// Prepare log file
	var logFile *os.File
	logFile, err = os.Create(params.logFilePath)
	if err != nil {
		return receiverAddr, fmt.Errorf("cannot create %s: %s", params.logFilePath, err.Error())
	}
	log.Printf("Writing %s log to %s", cp.name, params.logFilePath)

	// Prepare to start the process.
	// #nosec
	args := params.cmdArgs
	if !containsConfig(args) {
		if cp.configFileName == "" {
			configFile := path.Join("testdata", "agent-config.yaml")
			cp.configFileName, err = filepath.Abs(configFile)
			if err != nil {
				return receiverAddr, err
			}
		}
		args = append(args, "--config")
		args = append(args, cp.configFileName)
	}
	cp.cmd = exec.Command(params.cmd, args...)

	// Capture standard output and standard error.
	stdoutIn, err := cp.cmd.StdoutPipe()
	if err != nil {
		return receiverAddr, fmt.Errorf("cannot capture stdout of %s: %s", params.cmd, err.Error())
	}
	stderrIn, err := cp.cmd.StderrPipe()
	if err != nil {
		return receiverAddr, fmt.Errorf("cannot capture stderr of %s: %s", params.cmd, err.Error())
	}

	// Start the process.
	if err = cp.cmd.Start(); err != nil {
		return receiverAddr, fmt.Errorf("cannot start executable at %s: %s", params.cmd, err.Error())
	}

	cp.startTime = time.Now()
	cp.isStarted = true

	log.Printf("%s running, pid=%d", cp.name, cp.cmd.Process.Pid)

	// Create a WaitGroup that waits for both outputs to be finished copying.
	cp.outputWG.Add(2)

	// Begin copying outputs.
	go func() {
		_, _ = io.Copy(logFile, stdoutIn)
		cp.outputWG.Done()
	}()
	go func() {
		_, _ = io.Copy(logFile, stderrIn)
		cp.outputWG.Done()
	}()

	receiverAddr = fmt.Sprintf("%s:%d", DefaultHost, 0)
	return receiverAddr, err
}

func (cp *ChildProcess) Stop() (stopped bool, err error) {
	if !cp.isStarted || cp.isStopped {
		return false, nil
	}
	cp.stopOnce.Do(func() {

		if !cp.isStarted {
			// Process wasn't started, nothing to stop.
			return
		}

		cp.isStopped = true

		log.Printf("Gracefully terminating %s pid=%d, sending SIGTEM...", cp.name, cp.cmd.Process.Pid)

		// Notify resource monitor to stop.
		close(cp.doneSignal)

		// Gracefully signal process to stop.
		if err = cp.cmd.Process.Signal(syscall.SIGTERM); err != nil {
			log.Printf("Cannot send SIGTEM: %s", err.Error())
		}

		finished := make(chan struct{})

		// Setup a goroutine to wait a while for process to finish and send kill signal
		// to the process if it doesn't finish.
		go func() {
			// Wait 10 seconds.
			t := time.After(10 * time.Second)
			select {
			case <-t:
				// Time is out. Kill the process.
				log.Printf("%s pid=%d is not responding to SIGTERM. Sending SIGKILL to kill forcedly.",
					cp.name, cp.cmd.Process.Pid)
				if err = cp.cmd.Process.Signal(syscall.SIGKILL); err != nil {
					log.Printf("Cannot send SIGKILL: %s", err.Error())
				}
			case <-finished:
				// Process is successfully finished.
			}
		}()

		// Wait for output to be fully copied.
		cp.outputWG.Wait()

		// Wait for process to terminate
		err = cp.cmd.Wait()

		// Let goroutine know process is finished.
		close(finished)

		// Set resource consumption stats to 0
		cp.ramMiBCur.Store(0)
		cp.cpuPercentX1000Cur.Store(0)

		log.Printf("%s process stopped, exit code=%d", cp.name, cp.cmd.ProcessState.ExitCode())

		if err != nil {
			log.Printf("%s execution failed: %s", cp.name, err.Error())
		}
	})
	stopped = true
	return stopped, err
}

func (cp *ChildProcess) WatchResourceConsumption() error {
	if !cp.resourceSpec.isSpecified() {
		// Resource monitoring is not enabled.
		return nil
	}

	var err error
	cp.processMon, err = process.NewProcess(int32(cp.cmd.Process.Pid))
	if err != nil {
		return fmt.Errorf("cannot monitor process %d: %s",
			cp.cmd.Process.Pid, err.Error())
	}

	cp.fetchRAMUsage()

	// Begin measuring elapsed and process CPU times.
	cp.lastElapsedTime = time.Now()
	cp.lastProcessTimes, err = cp.processMon.Times()
	if err != nil {
		return fmt.Errorf("cannot get process times for %d: %s",
			cp.cmd.Process.Pid, err.Error())
	}

	// Measure every ResourceCheckPeriod.
	ticker := time.NewTicker(cp.resourceSpec.ResourceCheckPeriod)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			cp.fetchRAMUsage()
			cp.fetchCPUUsage()

			if err := cp.checkAllowedResourceUsage(); err != nil {
				cp.Stop()
				return err
			}

		case <-cp.doneSignal:
			log.Printf("Stopping process monitor.")
			return nil
		}
	}
}

func (cp *ChildProcess) GetProcessMon() *process.Process {
	return cp.processMon
}

func (cp *ChildProcess) fetchRAMUsage() {
	// Get process memory and CPU times
	mi, err := cp.processMon.MemoryInfo()
	if err != nil {
		log.Printf("cannot get process memory for %d: %s",
			cp.cmd.Process.Pid, err.Error())
		return
	}

	// Calculate RSS in MiBs.
	ramMiBCur := uint32(mi.RSS / mibibyte)

	// Calculate aggregates.
	cp.memProbeCount++
	cp.ramMiBTotal += uint64(ramMiBCur)
	if ramMiBCur > cp.ramMiBMax {
		cp.ramMiBMax = ramMiBCur
	}

	// Store current usage.
	cp.ramMiBCur.Store(ramMiBCur)
}

func (cp *ChildProcess) fetchCPUUsage() {
	times, err := cp.processMon.Times()
	if err != nil {
		log.Printf("cannot get process times for %d: %s",
			cp.cmd.Process.Pid, err.Error())
		return
	}

	now := time.Now()

	// Calculate elapsed and process CPU time deltas in seconds
	deltaElapsedTime := now.Sub(cp.lastElapsedTime).Seconds()
	deltaCPUTime := times.Total() - cp.lastProcessTimes.Total()

	cp.lastProcessTimes = times
	cp.lastElapsedTime = now

	// Calculate CPU usage percentage in elapsed period.
	cpuPercent := deltaCPUTime * 100 / deltaElapsedTime
	if cpuPercent > cp.cpuPercentMax {
		cp.cpuPercentMax = cpuPercent
	}

	curCPUPercentageX1000 := uint32(cpuPercent * 1000)

	// Store current usage.
	cp.cpuPercentX1000Cur.Store(curCPUPercentageX1000)
}

func (cp *ChildProcess) checkAllowedResourceUsage() error {
	// Check if current CPU usage exceeds expected.
	var errMsg string
	if cp.resourceSpec.ExpectedMaxCPU != 0 && cp.cpuPercentX1000Cur.Load()/1000 > cp.resourceSpec.ExpectedMaxCPU {
		errMsg = fmt.Sprintf("CPU consumption is %.1f%%, max expected is %d%%",
			float64(cp.cpuPercentX1000Cur.Load())/1000.0, cp.resourceSpec.ExpectedMaxCPU)
	}

	// Check if current RAM usage exceeds expected.
	if cp.resourceSpec.ExpectedMaxRAM != 0 && cp.ramMiBCur.Load() > cp.resourceSpec.ExpectedMaxRAM {
		errMsg = fmt.Sprintf("RAM consumption is %d MiB, max expected is %d MiB",
			cp.ramMiBCur, cp.resourceSpec.ExpectedMaxRAM)
	}

	if errMsg == "" {
		return nil
	}

	log.Printf("Performance error: %s", errMsg)

	return errors.New(errMsg)
}

// GetResourceConsumption returns resource consumption as a string
func (cp *ChildProcess) GetResourceConsumption() string {
	if !cp.resourceSpec.isSpecified() {
		// Monitoring is not enabled.
		return ""
	}

	curRSSMib := cp.ramMiBCur.Load()
	curCPUPercentageX1000 := cp.cpuPercentX1000Cur.Load()

	return fmt.Sprintf("%s RAM (RES):%4d MiB, CPU:%4.1f%%", cp.name,
		curRSSMib, float64(curCPUPercentageX1000)/1000.0)
}

// GetTotalConsumption returns total resource consumption since start of process
func (cp *ChildProcess) GetTotalConsumption() *ResourceConsumption {
	rc := &ResourceConsumption{}

	if cp.processMon != nil {
		// Get total elapsed time since process start
		elapsedDuration := cp.lastElapsedTime.Sub(cp.startTime).Seconds()

		if elapsedDuration > 0 {
			// Calculate average CPU usage since start of process
			rc.CPUPercentAvg = cp.lastProcessTimes.Total() / elapsedDuration * 100.0
		}
		rc.CPUPercentMax = cp.cpuPercentMax

		if cp.memProbeCount > 0 {
			// Calculate average RAM usage by averaging all RAM measurements
			rc.RAMMiBAvg = uint32(cp.ramMiBTotal / uint64(cp.memProbeCount))
		}
		rc.RAMMiBMax = cp.ramMiBMax
	}

	return rc
}

func containsConfig(s []string) bool {
	for _, a := range s {
		if a == "--config" {
			return true
		}
	}
	return false
}
