package main

import (
    "syscall"
    "fmt"
    "log"
    "github.com/go-telegram-bot-api/telegram-bot-api"
    "time"
    "strings"
    "flag"
    "strconv"
)

var (
    token string
    treshold float64
    interval float64
    mPoints []string
    chatIds = []string{}
)

type DiskStatus struct {
    All uint64
    Used uint64
    Free uint64
}

func DiskUsage(path string) (*DiskStatus, error) {
    fs := syscall.Statfs_t{}
    err := syscall.Statfs(path, &fs)
    if err != nil {
        log.Println("Can't get data for", path, err)
        return nil, err
    }
    var disk DiskStatus
    disk.All = fs.Blocks * uint64(fs.Bsize)
    disk.Free = fs.Bfree * uint64(fs.Bsize)
    disk.Used = disk.All - disk.Free
    return &disk, nil
}

func HumanizeSize(bytes uint64) string {
    points := [6]string{"B", "KB", "MB", "GB", "TB", "PB"}
    size := float64(bytes)
    index := 0
    for size > 1024 && index < len(points) {
        size /= 1024
        index++
    }
    return fmt.Sprintf("%.2f%s", size, points[index])
}

func SendMessage(text string, chatId int64) {
    bot, err := tgbotapi.NewBotAPI(token)
    if err != nil {
        log.Println("Failed initializion telegram API.", err)
        return
    }

    msg := tgbotapi.NewMessage(chatId, text)
    bot.Send(msg)
}

func Notify(text string) {
    if token == "" {
        log.Println("Token is empty.")
        return
    }

    if len(chatIds) < 1 {
        log.Println("Chat ids not found.")
        return
    }

    for _, id := range chatIds {
        chatId, err := strconv.Atoi(id)
        if err != nil {
            log.Println(err)
            continue
        }
        SendMessage(text, int64(chatId))
    }
}

func CheckStatus(mPoints []string) {
    for _, path := range mPoints {
        disk, err := DiskUsage(path)
        if err != nil {
            log.Println(err)
            continue
        }
        percent := float64(disk.Used) / float64(disk.All)
        freePercent := (1 - percent) * 100.0
        if freePercent <= treshold {
            Notify(fmt.Sprintf("Low disk space: '%s' (%s)\n", path, HumanizeSize(disk.Free)))
        }
    }
}

func schedule(duration time.Duration, function func()) {
    for {
        timer := time.NewTimer(duration)
        <-timer.C
        function()
    }
}

func init() {
    var paths string
    var userIds string

    flag.StringVar(&paths, "p", "", "Set mount points for monitoring. Separator: \".\" Ex. /volume1,/volume2")
    flag.StringVar(&token, "t", "", "Telegram bot token.")
    flag.StringVar(&userIds, "c", "", "Users chat ids for send notifications.")
    flag.Float64Var(&interval, "i",600.0, "Disk check interval in seconds. Default: 600 seconds")
    flag.Float64Var(&treshold, "d", 5.0, "Alert disk space treshold in percent. Default: 10%")
    flag.Parse()

    if paths != "" {
        mPoints = strings.Split(paths, ",")
    }

    if userIds != "" {
        chatIds = strings.Split(userIds, ",")
    }
}

func main() {
    if len(mPoints) < 1 {
        log.Fatalln("Mount points empty.")
    }
    schedule(time.Second * time.Duration(interval), func() {
        CheckStatus(mPoints)
    })
}