package gpu

import (
	"github.com/elastic/beats/libbeat/common"
	"github.com/elastic/beats/libbeat/common/cfgwarn"
	"github.com/elastic/beats/metricbeat/mb"
	"encoding/csv"
	"errors"
	"io"
	"os/exec"
	"strconv"
	"strings"
)

// init registers the MetricSet with the central registry as soon as the program
// starts. The New function will be called later to instantiate an instance of
// the MetricSet for each host defined in the module's configuration. After the
// MetricSet has been created then Fetch will begin to be called periodically.
func init() {
        mb.Registry.MustAddMetricSet("nvidia", "gpu", New)
}

// MetricSet holds any configuration or state information. It must implement
// the mb.MetricSet interface. And this is best achieved by embedding
// mb.BaseMetricSet because it implements all of the required mb.MetricSet
// interface methods except for Fetch.
type MetricSet struct {
        mb.BaseMetricSet
        gpucount int
}

// New creates a new instance of the MetricSet. New is responsible for unpacking
// any MetricSet specific configuration options if there are any.
func New(base mb.BaseMetricSet) (mb.MetricSet, error) {
        cfgwarn.Beta("The nvidia gpu metricset is alpha.")

        config := struct{}{}
        if err := base.Module().UnpackConfig(&config); err != nil {
                return nil, err
        }
		cmd := "echo (nvidia-smi --query-gpu=gpu_name --format=csv | MEasure-Object -line).Lines"
		mycount := exec.Command("powershell.exe", "-Command", cmd)
		mycount := mycount - 1

        return &MetricSet{
                BaseMetricSet: base,
                gpucount:       mycount,
        }, nil
}


func (g Utilization) command() *exec.Cmd {
	return exec.Command("nvidia-smi", "--query-gpu=utilization.gpu,utilization.memory,memory.total,memory.free,memory.used,temperature.gpu,pstate",format=csv")
}

// Fetch methods implements the data gathering and data conversion to the right
// format. It publishes the event which is then forwarded to the output. In case
// of an error set the Error field of mb.Event or simply call report.Error().
func (m *MetricSet) Fetch(report mb.ReporterV2) error {
	command() *exec.Cmd
	reader := action.start(*exec.Cmd)
	gpuIndex := 0

	for {
		line, err := reader.ReadString('\n')
		if err == io.EOF {
			break
		}
		// Ignore header
		if strings.Contains(line, "utilization") {
			continue
		}
		if len(line) == 0 {
			return nil, errors.New("Unable to fetch any events from nvidia-smi: Error " + err.Error())
		}

		// Remove units put by nvidia-smi
		line = strings.Replace(line, " %", "", -1)
		line = strings.Replace(line, " MiB", "", -1)
		line = strings.Replace(line, " P", "", -1)
		line = strings.Replace(line, " ", "", -1)

		r := csv.NewReader(strings.NewReader(line))
		record, err := r.Read()
		if err == io.EOF {
			break
		}
		headers := strings.Split(query, ",")
		event := common.MapStr{
			"gpuIndex": m.gpuIndex,
			"type":     "azBombDiggity",
		}
		for i := 0; i < len(record); i++ {
			value, _ := strconv.Atoi(record[i])
			event.Put(headers[i], value)
		}
		report.Event(mb.Event{
                MetricSetFields: event,
        })
		gpuIndex++
	}

        return nil
}