package haproxy

import (
        "encoding/csv"
        "fmt"
        "github.com/influxdata/telegraf"
        "github.com/influxdata/telegraf/plugins/inputs"
        "io"
        "net"
        "net/http"
        "net/url"
        "strconv"
        "strings"
        "sync"
        "time"
)

type haproxy struct {
        Servers []string

        client *http.Client
}

var sampleConfig = `
  ## An array of address to gather stats about. Specify an ip on hostname
  ## with optional port. ie localhost, 10.10.3.33:1936, etc.
  ## If no servers are specified, then default to 127.0.0.1:1936
  servers = ["http://myhaproxy.com:1936", "http://anotherhaproxy.com:1936"]
  ## Or you can also use local socket
  ## servers = ["socket:/run/haproxy/admin.sock"]
`

func (r *haproxy) SampleConfig() string {
        return sampleConfig
}

func (r *haproxy) Description() string {
        return "Read metrics of haproxy, via socket or csv stats page"
}

// Reads stats from all configured servers accumulates stats.
// Returns one of the errors encountered while gather stats (if any).
func (g *haproxy) Gather(acc telegraf.Accumulator) error {
        if len(g.Servers) == 0 {
                return g.gatherServer("http://127.0.0.1:1936", acc)
        }

        var wg sync.WaitGroup

        var outerr error

        for _, serv := range g.Servers {
                wg.Add(1)
                go func(serv string) {
                        defer wg.Done()
                        outerr = g.gatherServer(serv, acc)
                }(serv)
        }

        wg.Wait()

        return outerr
}

func (g *haproxy) gatherServerSocket(addr string, acc telegraf.Accumulator) error {
        var socketPath string
        socketAddr := strings.Split(addr, ":")

        if len(socketAddr) >= 2 {
                socketPath = socketAddr[1]
        } else {
                socketPath = socketAddr[0]
        }

        c, err := net.Dial("unix", socketPath)

        if err != nil {
                return fmt.Errorf("Could not connect to socket '%s': %s", addr, err)
        }

        _, errw := c.Write([]byte("show stat\n"))

        if errw != nil {
                return fmt.Errorf("Could not write to socket '%s': %s", addr, errw)
        }

        return importCsvResult(c, acc, socketPath)
}

func (g *haproxy) gatherServer(addr string, acc telegraf.Accumulator) error {
        if !strings.HasPrefix(addr, "http") {
                return g.gatherServerSocket(addr, acc)
        }

        if g.client == nil {
                tr := &http.Transport{ResponseHeaderTimeout: time.Duration(3 * time.Second)}
                client := &http.Client{
                        Transport: tr,
                        Timeout:   time.Duration(4 * time.Second),
                }
                g.client = client
        }

        u, err := url.Parse(addr)
        if err != nil {
                return fmt.Errorf("Unable parse server address '%s': %s", addr, err)
        }

        req, err := http.NewRequest("GET", fmt.Sprintf("%s://%s%s/;csv", u.Scheme, u.Host, u.Path), nil)
        if u.User != nil {
                p, _ := u.User.Password()
                req.SetBasicAuth(u.User.Username(), p)
        }

        res, err := g.client.Do(req)
        if err != nil {
                return fmt.Errorf("Unable to connect to haproxy server '%s': %s", addr, err)
        }

        if res.StatusCode != 200 {
                return fmt.Errorf("Unable to get valid stat result from '%s': %s", addr, err)
        }

        return importCsvResult(res.Body, acc, u.Host)
}

func importCsvResult(r io.Reader, acc telegraf.Accumulator, host string) error {
        csv := csv.NewReader(r)
        result, err := csv.ReadAll()
        now := time.Now()

        var keys []string
        px_sv_status := make(map[string]map[string]int)
        non_stat_fields := map[string]bool{
                "pxname":       true,
                "svname":       true,
                "status":       true,
                "tracked":      true,
                "check_status": true,
                "last_chk":     true,
                "pid":          true,
                "iid":          true,
                "sid":          true,
                "lastchg":      true,
                "type":         true,
                "check_code":   true,
        }

        for i := range result {
                if i == 0 {
                        keys = result[i]
                        keys[0] = strings.Replace(keys[0], "# ", "", -1)
                } else {

                        row := make(map[string]string, len(result[i]))

                        for f, v := range result[i] {
                                row[keys[f]] = v
                        }

                        tags := map[string]string{
                                "server": host,
                                "proxy":  row["pxname"],
                                "sv":     row["svname"],
                        }

                        if row["svname"] != "BACKEND" && row["svname"] != "FRONTEND" {
                                if len(px_sv_status[row["pxname"]]) == 0 {
                                        px_sv_status[row["pxname"]] = map[string]int{
                                                "_status_act":       0,
                                                "_status_act_up":    0,
                                                "_status_act_down":  0,
                                                "_status_act_maint": 0,
                                                "_status_act_drain": 0,
                                                "_status_act_other": 0,
                                                "_status_bck":       0,
                                                "_status_bck_up":    0,
                                                "_status_bck_down":  0,
                                                "_status_bck_maint": 0,
                                                "_status_bck_drain": 0,
                                                "_status_bck_other": 0,
                                        }
                                }

                                if row["act"] == "1" {
                                        px_sv_status[row["pxname"]]["_status_act"]++
                                        switch row["status"] {
                                        case "UP":
                                                px_sv_status[row["pxname"]]["_status_act_up"]++
                                        case "DOWN":
                                                px_sv_status[row["pxname"]]["_status_act_down"]++
                                        case "MAINT":
                                                px_sv_status[row["pxname"]]["_status_act_maint"]++
                                        case "DRAIN":
                                                px_sv_status[row["pxname"]]["_status_act_drain"]++
                                        default:
                                                px_sv_status[row["pxname"]]["_status_act_other"]++
                                        }

                                } else {
                                        px_sv_status[row["pxname"]]["_status_bck"]++
                                        switch row["status"] {
                                        case "UP":
                                                px_sv_status[row["pxname"]]["_status_bck_up"]++
                                        case "DOWN":
                                                px_sv_status[row["pxname"]]["_status_bck_down"]++
                                        case "MAINT":
                                                px_sv_status[row["pxname"]]["_status_bck_maint"]++
                                        case "DRAIN":
                                                px_sv_status[row["pxname"]]["_status_bck_drain"]++
                                        default:
                                                px_sv_status[row["pxname"]]["_status_bck_other"]++
                                        }
                                }
                        }

                        if row["svname"] == "BACKEND" && len(px_sv_status[row["pxname"]]) > 0 {
                                for s, c := range px_sv_status[row["pxname"]] {
                                        row[s] = strconv.Itoa(c)
                                }
                        }

                        fields := make(map[string]interface{})
                        for field, v := range row {
                                if non_stat_fields[field] {
                                        continue
                                }

                                ival, err := strconv.ParseUint(v, 10, 64)
                                if err == nil {
                                        fields[field] = ival
                                }
                        }
                        acc.AddFields("haproxy", fields, tags, now)
                }
        }
        return err
}

func init() {
        inputs.Add("haproxy", func() telegraf.Input {
                return &haproxy{}
        })
}