// Copyright 2022 The Prometheus Authors
// 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
//
// http://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 vultr

import (
	"context"
	"fmt"
	"net"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/go-kit/log"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/config"
	"github.com/prometheus/common/model"
	"github.com/prometheus/common/version"
	"github.com/vultr/govultr/v2"

	"github.com/prometheus/prometheus/discovery"
	"github.com/prometheus/prometheus/discovery/refresh"
	"github.com/prometheus/prometheus/discovery/targetgroup"
)

const (
	vultrInstanceLabel             = model.MetaLabelPrefix + "vultr_instance_"
	vultrInstanceLabelID           = vultrInstanceLabel + "id"
	vultrInstanceLabelLabel        = vultrInstanceLabel + "label"
	vultrInstanceLabelOS           = vultrInstanceLabel + "os"
	vultrInstanceLabelOSID         = vultrInstanceLabel + "os_id"
	vultrInstanceLabelRegion       = vultrInstanceLabel + "region"
	vultrInstanceLabelPlan         = vultrInstanceLabel + "plan"
	vultrInstanceLabelMainIP       = vultrInstanceLabel + "main_ip"
	vultrInstanceLabelMainIPV6     = vultrInstanceLabel + "main_ipv6"
	vultrInstanceLabelInternalIP   = vultrInstanceLabel + "internal_ip"
	vultrInstanceLabelFeatures     = vultrInstanceLabel + "features"
	vultrInstanceLabelTags         = vultrInstanceLabel + "tags"
	vultrInstanceLabelHostname     = vultrInstanceLabel + "hostname"
	vultrInstanceLabelServerStatus = vultrInstanceLabel + "server_status"
	vultrInstanceLabelVCPU         = vultrInstanceLabel + "vcpu_count"
	vultrInstanceLabelMemory       = vultrInstanceLabel + "ram_mb"
	vultrInstanceLabelBandwidth    = vultrInstanceLabel + "allowed_bandwidth_gb"
	vultrInstanceLabelDisk         = vultrInstanceLabel + "disk_gb"
	separator                      = ","
)

// DefaultSDConfig is the default Vultr SD configuration.
var DefaultSDConfig = SDConfig{
	Port:             80,
	RefreshInterval:  model.Duration(60 * time.Second),
	HTTPClientConfig: config.DefaultHTTPClientConfig,
}

func init() {
	discovery.RegisterConfig(&SDConfig{})
}

type SDConfig struct {
	HTTPClientConfig config.HTTPClientConfig `yaml:",inline"`

	RefreshInterval model.Duration `yaml:"refresh_interval"`
	Port            int            `yaml:"port"`
}

// Name returns the name of the Config.
func (*SDConfig) Name() string { return "vultr" }

// NewDiscoverer returns a Discoverer for the Config.
func (c *SDConfig) NewDiscoverer(opts discovery.DiscovererOptions) (discovery.Discoverer, error) {
	return NewDiscovery(c, opts.Logger, opts.Registerer)
}

// SetDirectory joins any relative file paths with dir.
func (c *SDConfig) SetDirectory(dir string) {
	c.HTTPClientConfig.SetDirectory(dir)
}

// UnmarshalYAML implements the yaml.Unmarshaler interface.
func (c *SDConfig) UnmarshalYAML(unmarshal func(interface{}) error) error {
	*c = DefaultSDConfig
	type plain SDConfig
	if err := unmarshal((*plain)(c)); err != nil {
		return err
	}

	return c.HTTPClientConfig.Validate()
}

// Discovery periodically performs Vultr requests. It implements
// the Discoverer interface.
type Discovery struct {
	*refresh.Discovery
	client *govultr.Client
	port   int
}

// NewDiscovery returns a new Discovery which periodically refreshes its targets.
func NewDiscovery(conf *SDConfig, logger log.Logger, reg prometheus.Registerer) (*Discovery, error) {
	d := &Discovery{
		port: conf.Port,
	}

	rt, err := config.NewRoundTripperFromConfig(conf.HTTPClientConfig, "vultr_sd")
	if err != nil {
		return nil, err
	}

	d.client = govultr.NewClient(&http.Client{
		Transport: rt,
		Timeout:   time.Duration(conf.RefreshInterval),
	})

	d.client.SetUserAgent(fmt.Sprintf("Prometheus/%s", version.Version))

	if err != nil {
		return nil, fmt.Errorf("error setting up vultr agent: %w", err)
	}

	d.Discovery = refresh.NewDiscovery(
		refresh.Options{
			Logger:   logger,
			Mech:     "vultr",
			Interval: time.Duration(conf.RefreshInterval),
			RefreshF: d.refresh,
			Registry: reg,
		},
	)
	return d, nil
}

func (d *Discovery) refresh(ctx context.Context) ([]*targetgroup.Group, error) {
	tg := &targetgroup.Group{
		Source: "Vultr",
	}

	instances, err := d.listInstances(ctx)
	if err != nil {
		return nil, err
	}

	for _, instance := range instances {
		labels := model.LabelSet{
			vultrInstanceLabelID:           model.LabelValue(instance.ID),
			vultrInstanceLabelLabel:        model.LabelValue(instance.Label),
			vultrInstanceLabelOS:           model.LabelValue(instance.Os),
			vultrInstanceLabelOSID:         model.LabelValue(strconv.Itoa(instance.OsID)),
			vultrInstanceLabelRegion:       model.LabelValue(instance.Region),
			vultrInstanceLabelPlan:         model.LabelValue(instance.Plan),
			vultrInstanceLabelVCPU:         model.LabelValue(strconv.Itoa(instance.VCPUCount)),
			vultrInstanceLabelMemory:       model.LabelValue(strconv.Itoa(instance.RAM)),
			vultrInstanceLabelBandwidth:    model.LabelValue(strconv.Itoa(instance.AllowedBandwidth)),
			vultrInstanceLabelDisk:         model.LabelValue(strconv.Itoa(instance.Disk)),
			vultrInstanceLabelMainIP:       model.LabelValue(instance.MainIP),
			vultrInstanceLabelMainIPV6:     model.LabelValue(instance.V6MainIP),
			vultrInstanceLabelInternalIP:   model.LabelValue(instance.InternalIP),
			vultrInstanceLabelHostname:     model.LabelValue(instance.Hostname),
			vultrInstanceLabelServerStatus: model.LabelValue(instance.ServerStatus),
		}

		addr := net.JoinHostPort(instance.MainIP, strconv.FormatUint(uint64(d.port), 10))
		labels[model.AddressLabel] = model.LabelValue(addr)

		// We surround the separated list with the separator as well. This way regular expressions
		// in relabeling rules don't have to consider feature positions.
		if len(instance.Features) > 0 {
			features := separator + strings.Join(instance.Features, separator) + separator
			labels[vultrInstanceLabelFeatures] = model.LabelValue(features)
		}

		if len(instance.Tags) > 0 {
			tags := separator + strings.Join(instance.Tags, separator) + separator
			labels[vultrInstanceLabelTags] = model.LabelValue(tags)
		}

		tg.Targets = append(tg.Targets, labels)
	}

	return []*targetgroup.Group{tg}, nil
}

func (d *Discovery) listInstances(ctx context.Context) ([]govultr.Instance, error) {
	var instances []govultr.Instance

	listOptions := &govultr.ListOptions{
		PerPage: 100,
	}

	for {
		pagedInstances, meta, err := d.client.Instance.List(ctx, listOptions)
		if err != nil {
			return nil, err
		}
		instances = append(instances, pagedInstances...)

		if meta.Links.Next == "" {
			break
		}
		listOptions.Cursor = meta.Links.Next
	}

	return instances, nil
}