package main

import (
	"bufio"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"strings"
)

type SmartDev struct {
	DeviceFile string
	Name       string
	Type       string
}

// Lower case for zabbix-expected Json
type SmartDiscoEntry struct {
	Name       string  `json:"name"`
	Type       string  `json:"type"`
	Model      *string `json:"model"`
	SN         *string `json:"sn"`
	Rotations  *int32  `json:"rotations"`
	Lbs        *int32  `json:"lbs"`
	IsExternal int32   `json:"isExternal"`
	IsSSD      int32   `json:"isSSD"`
}

type SmartInfoEntry struct {
	Name            string
	DeviceFile      string
	Health          string
	PowerOnHours    int32
	PowerCycleCount int32
	Temperature     float32
}

type smartctlJson struct {
	ModelName string `json:"model_name"`
	SN        string `json:"serial_number"`
	Rotations int32  `json:"rotation_rate"`
	Lbs       int32  `json:"logical_block_size"`
}

func ExecSmartCtl(resAsSingleString bool, args ...string) ([]string, error) {

	path, err := exec.LookPath("smartctl")
	if err != nil {
		return nil, fmt.Errorf("coudn't find smartctl %v", err)
	}

	out, err := exec.Command(path, args...).Output()
	if err != nil {
		var exitErr *exec.ExitError
		if errors.As(err, &exitErr) && exitErr.ExitCode() == 2 {
			return nil, fmt.Errorf("smartctl failed, not running root?: %v", exitErr)
		} else if errors.As(err, &exitErr) && exitErr.ExitCode() == 4 {
			// ignore, no or partial smart support
		} else {
			return nil, fmt.Errorf("smartctl failed: %v", exitErr)
		}
	}

	if !resAsSingleString {
		var lns []string
		s := bufio.NewScanner(strings.NewReader(string(out)))
		for s.Scan() {
			lns = append(lns, s.Text())
		}
		return lns, nil
	} else {
		var lns []string
		lns = append(lns, string(out))
		return lns, nil
	}
}

func IsUSBDrive(dev string) (bool, error) {
	path, err := exec.LookPath("lsblk")
	if err != nil {
		return false, fmt.Errorf("couldn't find lsblk: %v", err)
	}

	out, err := exec.Command(path, "-d", dev, "-o", "SUBSYSTEMS").Output()
	if err != nil {
		return false, fmt.Errorf("calling lsblk failed: %v", err)
	}

	return strings.Index(string(out), "usb") >= 0, nil
}

func GetSmartDevsFromScan() ([]SmartDev, error) {
	var r []SmartDev

	lns, err := ExecSmartCtl(false, "-n", "standby", "--scan")
	if err != nil {
		return r, fmt.Errorf("Scanning for devices: %w", err)
	}

	for _, ln := range lns {
		if len(ln) > 0 {
			//fmt.Println("x" + ln)
			splt := strings.Split(ln, " ")
			var sd SmartDev
			sd.DeviceFile = splt[0]
			sd.Name = strings.Split(splt[0], "/")[2]
			sd.Type = strings.Split(strings.Split(ln, ",")[1], " ")[1]
			r = append(r, sd)
		}
	}

	return r, nil
}

func GetSmartDisco() ([]SmartDiscoEntry, error) {
	var r []SmartDiscoEntry

	devs, err := GetSmartDevsFromScan()
	if err != nil {
		return r, err
	}

	for _, dev := range devs {
		isUSB, err := IsUSBDrive(dev.DeviceFile)
		if err != nil {
			return r, fmt.Errorf("failed get disco: %v", err)
		}

		var entry SmartDiscoEntry

		if !isUSB {
			//fmt.Println(dev.DeviceFile)
			lns, err := ExecSmartCtl(true, "-n", "standby", "-a", dev.DeviceFile, "--json")
			if err != nil {
				fmt.Fprintf(os.Stderr, "? skipped device %v: %v\n", dev.Name, err)
				continue
			}
			//j := lns[0]
			var sc smartctlJson
			// Prevent collision with bash-call
			err = json.Unmarshal([]byte(strings.ReplaceAll(lns[0], "'", "")), &sc)
			if err != nil {
				return r, fmt.Errorf("failed parsing json: %v", err)
			}
			entry.Model = &sc.ModelName
			entry.SN = &sc.SN
			entry.Rotations = &sc.Rotations
			entry.Lbs = &sc.Lbs

		} else {
			entry.Model = nil
			entry.SN = nil
			entry.Rotations = nil
			entry.Lbs = nil
		}
		entry.Name = dev.Name
		entry.Type = dev.Type
		entry.IsExternal = func() int32 {
			if isUSB {
				return 1
			}
			return 0
		}()
		entry.IsSSD = 0 // TODO
		r = append(r, entry)

	}

	return r, nil
}

func SendSmartDiscoToZabbix(testMode bool) error {
	disco, err := GetSmartDisco()
	if err != nil {
		return fmt.Errorf("SendSmartDiscoToZabbix failed: %w", err)
	}

	j, err := json.Marshal(disco)
	if err != nil {
		return fmt.Errorf("SendSmartDiscoToZabbix failed: %v", err)
	}

	if testMode {
		fmt.Printf("Testmode: Would send %v\n", string(j))

	} else {
		path, err := exec.LookPath("bash")
		if err != nil {
			return fmt.Errorf("coudn't find bash %v", err)
		}

		_, err = exec.Command(path, "-c", ""+
			"zabbix_sender -vv -k 8o_smartcheck.disco_devs -o '"+
			string(j)+
			"' -c /etc/zabbix/zabbix_agentd.conf"+
			" |  logger -t zbs_smart").Output()
		if err != nil {
			return fmt.Errorf("zabbix_sender exec failed: %v", err)
		}

	}
	return nil
}

func main() {

	doDiscoPtr := flag.Bool("smart_disco", false, "Do discovery")
	doCheckPtr := flag.Bool("smart_check", false, "Do check")
	testModePtr := flag.Bool("test", false, "TestMode")

	flag.Parse()

	var err error

	if *testModePtr {
		fmt.Println("----- Testing GetSmartDevsFromScan")
		s, _ := GetSmartDevsFromScan()
		j, _ := json.Marshal(s)
		fmt.Println(string(j))
		fmt.Println()

		fmt.Println("----- Testing IsUsbDrive")
		for _, dev := range s {
			isUSB, err := IsUSBDrive(dev.DeviceFile)
			fmt.Printf("- %s %v %+v\n", dev.Name, isUSB, err)
		}

		fmt.Println("----- SendSmartDiscoToZabbix")
		err = SendSmartDiscoToZabbix(*testModePtr)
	}

	if *doDiscoPtr {
		fmt.Println("Doing SMART device discovery")
		err = SendSmartDiscoToZabbix(*testModePtr)
	}

	if err == nil && *doCheckPtr {
		fmt.Println("Doig SMART checks")
		// TODO
	}

	if err != nil {
		fmt.Fprintf(os.Stderr, "! %+v\n", err)
	}

	fmt.Println("Done")

}