package main

import (
	"fmt"
	"flag"
	"os"

	"database/sql"
	_ "github.com/go-sql-driver/mysql"
)

// Option to parse
var (
  HELP     = flag.Bool("help", false, "print the help message")
  HOSTNAME = flag.String("hostname", "localhost", "mysql hostname")
  PORT     = flag.Int("port", 3306, "mysql port")
  USERNAME = flag.String("username", "", "mysql username")
  PASSWORD = flag.String("password", "", "mysql password")
  DATABASE = flag.String("database", "", "mysql database")
  QUERY    = flag.String("query", "", "mysql query")
  WARNING  = flag.Int("warning", 0, "limit to return a warning alert")
  CRITICAL = flag.Int("critical", 0, "limit to return a critical alert")
)

func init() {
  flag.Parse()

  if *HELP {
    flag.PrintDefaults()
    os.Exit(3)
  }

  if *QUERY == "" {
    Critical("Please set query option")
  }

  if *WARNING == 0 || *CRITICAL == 0 {
    Critical("Please set warning and critical options")
  }
}

// Critical print a message and exit with code 2
func Critical(msg string) {
  fmt.Println(fmt.Sprintf("CRITICAL - %s", msg))
  os.Exit(2)
}

// Warning print a message and exit with code 1
func Warning(msg string) {
  fmt.Println(fmt.Sprintf("WARNING - %s", msg))
  os.Exit(1)
}

// Ok print a message and exit with code 0
func Ok(msg string) {
  fmt.Println(fmt.Sprintf("OK - %s", msg))
  os.Exit(0)
}

// Count execute query and return the result
func Count() (int, error) {
  var count int

  conn := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s", *USERNAME, *PASSWORD, *HOSTNAME, *PORT, *DATABASE)
  db, err := sql.Open("mysql", conn)
  if err != nil {
    return 0, err
  }
  defer db.Close()

  err = db.QueryRow(*QUERY).Scan(&count)
  if err != nil {
    return 0, err
  }

  return count, nil
}

func main() {
  count, err := Count()
  if err != nil {
    Critical(fmt.Sprintf("%v", err))
  }

  msg := fmt.Sprintf("there are %d rows", count)
  if count >= *CRITICAL {
    Critical(msg)
  }

  if count >= *WARNING {
    Warning(msg)
  }

  Ok(msg)
}