/*
** Copyright (C) 2001-2025 Zabbix SIA
**
** This program is free software: you can redistribute it and/or modify it under the terms of
** the GNU Affero General Public License as published by the Free Software Foundation, version 3.
**
** This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;
** without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
** See the GNU Affero General Public License for more details.
**
** You should have received a copy of the GNU Affero General Public License along with this program.
** If not, see <https://www.gnu.org/licenses/>.
**/

package systemd

import (
	"encoding/json"
	"fmt"
	"path/filepath"
	"reflect"
	"strings"
	"sync"

	"github.com/godbus/dbus/v5"
	"golang.zabbix.com/sdk/errs"
	"golang.zabbix.com/sdk/plugin"
)

// Plugin -
type Plugin struct {
	plugin.Base
	connections []*dbus.Conn
	mutex       sync.Mutex
}

var impl Plugin

type unit struct {
	Name        string
	Description string
	LoadState   string
	ActiveState string
	SubState    string
	Followed    string
	Path        string
	JobID       uint32
	JobType     string
	JobPath     string
}

type unitFile struct {
	Name            string
	EnablementState string
}

type unitJson struct {
	Name          string `json:"{#UNIT.NAME}"`
	Description   string `json:"{#UNIT.DESCRIPTION}"`
	LoadState     string `json:"{#UNIT.LOADSTATE}"`
	ActiveState   string `json:"{#UNIT.ACTIVESTATE}"`
	SubState      string `json:"{#UNIT.SUBSTATE}"`
	Followed      string `json:"{#UNIT.FOLLOWED}"`
	Path          string `json:"{#UNIT.PATH}"`
	JobID         uint32 `json:"{#UNIT.JOBID}"`
	JobType       string `json:"{#UNIT.JOBTYPE}"`
	JobPath       string `json:"{#UNIT.JOBPATH}"`
	UnitFileState string `json:"{#UNIT.UNITFILESTATE}"`
	ServiceType   string `json:"{#UNIT.SERVICETYPE}"` //nolint:tagliatelle
}

type state struct {
	State int    `json:"state"`
	Text  string `json:"text"`
}

type stateMapping struct {
	unitName   string
	stateNames []string
}

func init() {
	err := plugin.RegisterMetrics(
		&impl, "Systemd",
		"systemd.unit.get", "Returns the bulked info, usage: systemd.unit.get[unit,<interface>].",
		"systemd.unit.discovery", "Returns JSON array of discovered units, usage: systemd.unit.discovery[<type>].",
		"systemd.unit.info", "Returns the unit info, usage: systemd.unit.info[unit,<parameter>,<interface>].",
	)
	if err != nil {
		panic(errs.Wrap(err, "failed to register metrics"))
	}
}

func (p *Plugin) getConnection() (*dbus.Conn, error) {
	var err error
	var conn *dbus.Conn

	p.mutex.Lock()
	defer p.mutex.Unlock()

	if len(p.connections) == 0 {
		conn, err = dbus.SystemBusPrivate()
		if err != nil {
			return nil, err
		}
		err = conn.Auth(nil)
		if err != nil {
			conn.Close()
			return nil, err
		}

		err = conn.Hello()
		if err != nil {
			conn.Close()
			return nil, err
		}

	} else {
		conn = p.connections[len(p.connections)-1]
		p.connections = p.connections[:len(p.connections)-1]
	}

	return conn, nil
}

func (p *Plugin) releaseConnection(conn *dbus.Conn) {
	p.mutex.Lock()
	defer p.mutex.Unlock()
	p.connections = append(p.connections, conn)
}

func zbxNum2hex(c byte) byte {
	if c >= 10 {
		return c + 0x57 /* a-f */
	}
	return c + 0x30 /* 0-9 */
}

// Export -
func (p *Plugin) Export(key string, params []string, ctx plugin.ContextProvider) (interface{}, error) {
	conn, err := p.getConnection()

	if nil != err {
		return nil, fmt.Errorf("Cannot establish connection to any available bus: %s", err)
	}

	defer p.releaseConnection(conn)

	switch key {
	case "systemd.unit.get":
		return p.get(params, conn)
	case "systemd.unit.discovery":
		return p.discovery(params, conn)
	case "systemd.unit.info":
		return p.info(params, conn)
	default:
		return nil, plugin.UnsupportedMetricError
	}
}

func (p *Plugin) get(params []string, conn *dbus.Conn) (interface{}, error) {
	var unitType string
	var values map[string]interface{}

	if len(params) > 2 {
		return nil, fmt.Errorf("Too many parameters.")
	}

	if len(params) == 0 || len(params[0]) == 0 {
		return nil, fmt.Errorf("Invalid first parameter.")
	}

	if len(params) < 2 || len(params[1]) == 0 {
		unitType = "Unit"
	} else {
		unitType = params[1]
	}

	obj := conn.Object(
		"org.freedesktop.systemd1",
		dbus.ObjectPath("/org/freedesktop/systemd1/unit/"+getName(params[0])),
	)
	err := obj.Call("org.freedesktop.DBus.Properties.GetAll", 0, "org.freedesktop.systemd1."+unitType).Store(&values)
	if err != nil {
		return nil, fmt.Errorf("Cannot get unit property: %s", err)
	}

	if unitType == "Unit" {
		p.setUnitStates(values)
	}

	val, err := json.Marshal(values)
	if err != nil {
		return nil, fmt.Errorf("Cannot create JSON array: %s", err)
	}

	return string(val), nil
}

func (p *Plugin) getServiceType(name string, conn *dbus.Conn) string {
	serviceType, err := p.info([]string{name, "Type", "Service"}, conn)
	if err != nil {
		p.Debugf("failed to retrieve service type for %s, err:%s", name, err.Error())
		return ""
	}

	typeString, ok := serviceType.(string)
	if !ok {
		p.Debugf("unit service type is not string for %s", name)
		return ""
	}

	return typeString
}

func (p *Plugin) discovery(params []string, conn *dbus.Conn) (interface{}, error) {
	var ext string

	if len(params) > 1 {
		return nil, fmt.Errorf("Too many parameters.")
	}

	if len(params) == 0 || len(params[0]) == 0 {
		ext = ".service"
	} else {
		switch params[0] {
		case "service", "target", "automount", "device", "mount", "path", "scope", "slice", "snapshot", "socket", "swap", "timer":
			ext = "." + params[0]
		case "all":
			ext = ""
		default:
			return nil, fmt.Errorf("Invalid first parameter.")
		}
	}

	var units []unit
	obj := conn.Object("org.freedesktop.systemd1", dbus.ObjectPath("/org/freedesktop/systemd1"))
	err := obj.Call("org.freedesktop.systemd1.Manager.ListUnits", 0).Store(&units)
	if nil != err {
		return nil, fmt.Errorf("Cannot retrieve list of units: %s", err)
	}

	var unitFiles []unitFile
	if err = obj.Call("org.freedesktop.systemd1.Manager.ListUnitFiles", 0).Store(&unitFiles); err != nil {
		return nil, fmt.Errorf("Cannot retrieve list of unit files: %s", err)
	}

	var array []unitJson
	for _, u := range units {
		if len(ext) != 0 && ext != filepath.Ext(u.Name) {
			continue
		}

		UnitFileState, err := p.info([]string{u.Name, "UnitFileState"}, conn)
		if err != nil {
			p.Debugf("Failed to retrieve unit file state for %s, err:%s", u.Name, err.Error())
			continue
		}

		var state string
		switch reflect.TypeOf(UnitFileState).Kind() {
		case reflect.String:
			state = UnitFileState.(string)
		default:
			p.Debugf("Unit file state is not string for %s", u.Name)
			continue
		}

		serviceType := p.getServiceType(u.Name, conn)
		array = append(array, unitJson{
			u.Name, u.Description, u.LoadState, u.ActiveState,
			u.SubState, u.Followed, u.Path, u.JobID, u.JobType, u.JobPath, state, serviceType,
		})
	}

	for _, f := range unitFiles {
		unitFileExt := filepath.Ext(f.Name)
		basePath := filepath.Base(f.Name)
		if f.EnablementState != "disabled" || (len(ext) != 0 && ext != unitFileExt) ||
			strings.HasSuffix(strings.TrimSuffix(f.Name, unitFileExt), "@") || /* skip unit templates */
			isEnabledUnit(array, basePath) {
			continue
		}

		unitPath := "/org/freedesktop/systemd1/unit/" + getName(basePath)

		serviceType := p.getServiceType(f.Name, conn)

		array = append(array, unitJson{basePath, "", "", "inactive", "", "", unitPath, 0, "", "",
			f.EnablementState, serviceType})
	}

	jsonArray, err := json.Marshal(array)
	if nil != err {
		return nil, fmt.Errorf("Cannot create JSON array: %s", err)
	}

	return string(jsonArray), nil
}

func (p *Plugin) info(params []string, conn *dbus.Conn) (interface{}, error) {
	var property, unitType string
	var value interface{}

	if len(params) > 3 {
		return nil, fmt.Errorf("Too many parameters.")
	}

	if len(params) < 1 || len(params[0]) == 0 {
		return nil, fmt.Errorf("Invalid first parameter.")
	}

	if len(params) < 2 || len(params[1]) == 0 {
		property = "ActiveState"
	} else {
		property = params[1]
	}

	if len(params) < 3 || len(params[2]) == 0 || params[2] == "Unit" {
		unitType = "Unit"
	} else {
		unitType = params[2]
	}

	obj := conn.Object(
		"org.freedesktop.systemd1",
		dbus.ObjectPath("/org/freedesktop/systemd1/unit/"+getName(params[0])),
	)
	err := obj.Call("org.freedesktop.DBus.Properties.Get", 0, "org.freedesktop.systemd1."+unitType, property).
		Store(&value)
	if nil != err {
		return nil, fmt.Errorf("Cannot get unit property: %s", err)
	}

	switch reflect.TypeOf(value).Kind() {
	case reflect.Slice:
		fallthrough
	case reflect.Array:
		ret, err := json.Marshal(value)
		if nil != err {
			return nil, fmt.Errorf("Cannot create JSON array: %s", err)
		}

		return string(ret), nil
	}

	return value, nil
}

func getName(name string) string {
	nameEsc := make([]byte, len(name)*3)
	j := 0
	for i := 0; i < len(name); i++ {
		if (name[i] >= 'A' && name[i] <= 'Z') ||
			(name[i] >= 'a' && name[i] <= 'z') ||
			((name[i] >= '0' && name[i] <= '9') && i != 0) {
			nameEsc[j] = name[i]
			j++
			continue
		}
		nameEsc[j] = '_'
		j++
		nameEsc[j] = zbxNum2hex((name[i] >> 4) & 0xf)
		j++
		nameEsc[j] = zbxNum2hex(name[i] & 0xf)
		j++
	}
	return string(nameEsc[:j])
}

func (p *Plugin) setUnitStates(v map[string]interface{}) {
	mappings := []stateMapping{
		{"LoadState", []string{"loaded", "error", "masked"}},
		{"ActiveState", []string{"active", "reloading", "inactive", "failed", "activating", "deactivating"}},
		{
			"UnitFileState",
			[]string{
				"enabled",
				"enabled-runtime",
				"linked",
				"linked-runtime",
				"masked",
				"masked-runtime",
				"static",
				"disabled",
				"invalid",
			},
		},
	}

	for _, mapping := range mappings {
		p.createStateMapping(v, mapping.unitName, mapping.stateNames)
	}
}

func (p *Plugin) createStateMapping(v map[string]interface{}, key string, names []string) {
	if value, ok := v[key].(string); ok {
		for i, name := range names {
			if value == name {
				v[key] = &state{i + 1, value}
				return
			}
		}
		v[key] = &state{0, value}
		p.Debugf("cannot create mapping for '%s' unit state: unknown state '%s'", key, value)
	} else {
		p.Debugf("cannot create mapping for '%s' unit state: unit state with information type string not found", key)
	}
}

func isEnabledUnit(units []unitJson, p string) bool {
	for _, u := range units {
		if u.Name == p {
			return true
		}
	}
	return false
}