package cmd

import (
  "fmt"
  "log"
  "os"
  "os/exec"

  "github.com/spf13/cobra"
)

var (
  size int
  tempfile string
)

var diskCmd = &cobra.Command{
  Use:   "disk",
  Short: "Disk speed test.",
  Long:  "Uses the local 'dd' command to test a disk's speed",
  Run: func(cmd *cobra.Command, args []string) {
    _, err := exec.LookPath("dd")
  	if err != nil {
  		panic(err)
  	}

    // get number of blocks to write (1M * 1024 * number of Gigabytes)
    blocksCount := fmt.Sprintf("count=%d", (size * 1024))
    // get outfile path
    outfile := fmt.Sprintf("of=%s", tempfile)

    ddCmd := exec.Command("dd", "bs=1M", blocksCount, "if=/dev/zero", outfile)
    ddCmd.Stderr = os.Stdout
    err = ddCmd.Run()
    if err != nil {
      log.Fatalf("cmd.Run() failed with %s\n", err)
    }

    fileErr := os.Remove(tempfile)
    if fileErr != nil {
      fmt.Println(fileErr)
    }
  },
}

func init() {
  // link flags to variables
  diskCmd.Flags().IntVarP(&size, "size", "s", 10, "Size of the file to write, in Gigabytes.")
  diskCmd.Flags().StringVarP(&tempfile, "tempfile", "t", "./tempfile", "Where to store the tempfile.")

  // add ping command to the main program
  rootCmd.AddCommand(diskCmd)
}