/*
  Copyright 2019 Adrien Waksberg
  
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at
  
      https://www.apache.org/licenses/LICENSE-2.0
  
  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.
*/

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", -1, "limit to return a warning alert")
  CRITICAL = flag.Int("critical", -1, "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 == -1 || *CRITICAL == -1 {
    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)
}