<commit_before>package stan_output

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/google/uuid"
	"github.com/karimra/gnmic/collector"
	"github.com/karimra/gnmic/outputs"
	"github.com/mitchellh/mapstructure"
	"github.com/nats-io/nats.go"
	"github.com/nats-io/stan.go"
	"github.com/openconfig/gnmi/proto/gnmi"
	"github.com/prometheus/client_golang/prometheus"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/reflect/protoreflect"
)

const (
	stanDefaultTimeout      = 10
	stanDefaultPingInterval = 5
	stanDefaultPingRetry    = 2

	defaultSubjectName = "gnmic-telemetry"
)

func init() {
	outputs.Register("stan", func() outputs.Output {
		return &StanOutput{
			Cfg: &Config{},
		}
	})
}

// StanOutput //
type StanOutput struct {
	Cfg      *Config
	conn     stan.Conn
	metrics  []prometheus.Collector
	logger   *log.Logger
	stopChan chan struct{}
}

// Config //
type Config struct {
	Name          string `mapstructure:"name,omitempty"`
	Address       string `mapstructure:"address,omitempty"`
	SubjectPrefix string `mapstructure:"subject-prefix,omitempty"`
	Subject       string `mapstructure:"subject,omitempty"`
	Username      string `mapstructure:"username,omitempty"`
	Password      string `mapstructure:"password,omitempty"`
	ClusterName   string `mapstructure:"cluster-name,omitempty"`
	Timeout       int    `mapstructure:"timeout,omitempty"`
	PingInterval  int    `mapstructure:"ping-interval,omitempty"`
	PingRetry     int    `mapstructure:"ping-retry,omitempty"`
	Format        string `mapstructure:"format,omitempty"`
}

func (s *StanOutput) String() string {
	b, err := json.Marshal(s)
	if err != nil {
		return ""
	}
	return string(b)
}

// Init //
func (s *StanOutput) Init(cfg map[string]interface{}, logger *log.Logger) error {
	err := mapstructure.Decode(cfg, s.Cfg)
	if err != nil {
		return err
	}
	if s.Cfg.Name == "" {
		s.Cfg.Name = "gnmic-" + uuid.New().String()
	}
	if s.Cfg.ClusterName == "" {
		return fmt.Errorf("clusterName is mandatory")
	}
	if s.Cfg.Subject == "" && s.Cfg.SubjectPrefix == "" {
		s.Cfg.Subject = defaultSubjectName
	}
	s.conn, err = s.createSTANConn(s.Cfg)
	if err != nil {
		return err
	}
	s.logger = log.New(os.Stderr, "stan_output ", log.LstdFlags|log.Lmicroseconds)
	if logger != nil {
		s.logger.SetOutput(logger.Writer())
		s.logger.SetFlags(logger.Flags())
	}
	if s.Cfg.Format == "" {
		s.Cfg.Format = "event"
	}
	if !(s.Cfg.Format == "event" || s.Cfg.Format == "json" || s.Cfg.Format == "proto") {
		return fmt.Errorf("unsupported output format: %s", s.Cfg.Format)
	}
	s.stopChan = make(chan struct{})
	s.logger.Printf("initialized stan producer: %s", s.String())
	return nil
}

// Write //
func (s *StanOutput) Write(rsp protoreflect.ProtoMessage, meta outputs.Meta) {
	if rsp == nil {
		return
	}
	if format, ok := meta["format"]; ok {
		if format == "textproto" {
			return
		}
	}
	ssb := strings.Builder{}
	ssb.WriteString(s.Cfg.SubjectPrefix)
	if s.Cfg.SubjectPrefix != "" {
		if s, ok := meta["source"]; ok {
			source := strings.ReplaceAll(fmt.Sprintf("%s", s), ".", "-")
			source = strings.ReplaceAll(source, " ", "_")
			ssb.WriteString(".")
			ssb.WriteString(source)
		}
		if subname, ok := meta["subscription-name"]; ok {
			ssb.WriteString(".")
			ssb.WriteString(fmt.Sprintf("%s", subname))
		}
	} else if s.Cfg.Subject != "" {
		ssb.WriteString(s.Cfg.Subject)
	}
	subject := strings.ReplaceAll(ssb.String(), " ", "_")
	b := make([]byte, 0)
	var err error
	switch s.Cfg.Format {
	case "proto":
		b, err = proto.Marshal(rsp)
	case "json":
		b, err = protojson.MarshalOptions{Multiline: true, Indent: "  "}.Marshal(rsp)
	case "event":
		switch sub := rsp.ProtoReflect().Interface().(type) {
		case *gnmi.SubscribeResponse:
			var subscriptionName string
			var ok bool
			if subscriptionName, ok = meta["subscription-name"]; !ok {
				subscriptionName = "default"
			}
			switch sub.Response.(type) {
			case *gnmi.SubscribeResponse_Update:
				events, err := collector.ResponseToEventMsgs(subscriptionName, sub, meta)
				if err != nil {
					s.logger.Printf("failed converting response to events: %v", err)
					return
				}
				b, err = json.MarshalIndent(events, "", "  ")
				if err != nil {
					s.logger.Printf("failed marshaling events: %v", err)
					return
				}
			case *gnmi.SubscribeResponse_SyncResponse:
				s.logger.Printf("received subscribe syncResponse with %v", meta)
			case *gnmi.SubscribeResponse_Error:
				gnmiErr := sub.GetError()
				s.logger.Printf("received subscribe response error with %v, code=%d, message=%v, data=%v ",
					meta, gnmiErr.Code, gnmiErr.Message, gnmiErr.Data)
			}
		}
	}
	if err != nil {
		s.logger.Printf("failed marshaling event: %v", err)
		return<commit_msg>remove indentation from stan output<commit_after>package stan_output

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"strings"

	"github.com/google/uuid"
	"github.com/karimra/gnmic/collector"
	"github.com/karimra/gnmic/outputs"
	"github.com/mitchellh/mapstructure"
	"github.com/nats-io/nats.go"
	"github.com/nats-io/stan.go"
	"github.com/openconfig/gnmi/proto/gnmi"
	"github.com/prometheus/client_golang/prometheus"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/proto"
	"google.golang.org/protobuf/reflect/protoreflect"
)

const (
	stanDefaultTimeout      = 10
	stanDefaultPingInterval = 5
	stanDefaultPingRetry    = 2

	defaultSubjectName = "gnmic-telemetry"
)

func init() {
	outputs.Register("stan", func() outputs.Output {
		return &StanOutput{
			Cfg: &Config{},
		}
	})
}

// StanOutput //
type StanOutput struct {
	Cfg      *Config
	conn     stan.Conn
	metrics  []prometheus.Collector
	logger   *log.Logger
	stopChan chan struct{}
}

// Config //
type Config struct {
	Name          string `mapstructure:"name,omitempty"`
	Address       string `mapstructure:"address,omitempty"`
	SubjectPrefix string `mapstructure:"subject-prefix,omitempty"`
	Subject       string `mapstructure:"subject,omitempty"`
	Username      string `mapstructure:"username,omitempty"`
	Password      string `mapstructure:"password,omitempty"`
	ClusterName   string `mapstructure:"cluster-name,omitempty"`
	Timeout       int    `mapstructure:"timeout,omitempty"`
	PingInterval  int    `mapstructure:"ping-interval,omitempty"`
	PingRetry     int    `mapstructure:"ping-retry,omitempty"`
	Format        string `mapstructure:"format,omitempty"`
}

func (s *StanOutput) String() string {
	b, err := json.Marshal(s)
	if err != nil {
		return ""
	}
	return string(b)
}

// Init //
func (s *StanOutput) Init(cfg map[string]interface{}, logger *log.Logger) error {
	err := mapstructure.Decode(cfg, s.Cfg)
	if err != nil {
		return err
	}
	if s.Cfg.Name == "" {
		s.Cfg.Name = "gnmic-" + uuid.New().String()
	}
	if s.Cfg.ClusterName == "" {
		return fmt.Errorf("clusterName is mandatory")
	}
	if s.Cfg.Subject == "" && s.Cfg.SubjectPrefix == "" {
		s.Cfg.Subject = defaultSubjectName
	}
	s.conn, err = s.createSTANConn(s.Cfg)
	if err != nil {
		return err
	}
	s.logger = log.New(os.Stderr, "stan_output ", log.LstdFlags|log.Lmicroseconds)
	if logger != nil {
		s.logger.SetOutput(logger.Writer())
		s.logger.SetFlags(logger.Flags())
	}
	if s.Cfg.Format == "" {
		s.Cfg.Format = "event"
	}
	if !(s.Cfg.Format == "event" || s.Cfg.Format == "json" || s.Cfg.Format == "proto") {
		return fmt.Errorf("unsupported output format: %s", s.Cfg.Format)
	}
	s.stopChan = make(chan struct{})
	s.logger.Printf("initialized stan producer: %s", s.String())
	return nil
}

// Write //
func (s *StanOutput) Write(rsp protoreflect.ProtoMessage, meta outputs.Meta) {
	if rsp == nil {
		return
	}
	if format, ok := meta["format"]; ok {
		if format == "textproto" {
			return
		}
	}
	ssb := strings.Builder{}
	ssb.WriteString(s.Cfg.SubjectPrefix)
	if s.Cfg.SubjectPrefix != "" {
		if s, ok := meta["source"]; ok {
			source := strings.ReplaceAll(fmt.Sprintf("%s", s), ".", "-")
			source = strings.ReplaceAll(source, " ", "_")
			ssb.WriteString(".")
			ssb.WriteString(source)
		}
		if subname, ok := meta["subscription-name"]; ok {
			ssb.WriteString(".")
			ssb.WriteString(fmt.Sprintf("%s", subname))
		}
	} else if s.Cfg.Subject != "" {
		ssb.WriteString(s.Cfg.Subject)
	}
	subject := strings.ReplaceAll(ssb.String(), " ", "_")
	b := make([]byte, 0)
	var err error
	switch s.Cfg.Format {
	case "proto":
		b, err = proto.Marshal(rsp)
	case "json":
		b, err = protojson.Marshal(rsp)
	case "event":
		switch sub := rsp.ProtoReflect().Interface().(type) {
		case *gnmi.SubscribeResponse:
			var subscriptionName string
			var ok bool
			if subscriptionName, ok = meta["subscription-name"]; !ok {
				subscriptionName = "default"
			}
			switch sub.Response.(type) {
			case *gnmi.SubscribeResponse_Update:
				events, err := collector.ResponseToEventMsgs(subscriptionName, sub, meta)
				if err != nil {
					s.logger.Printf("failed converting response to events: %v", err)
					return
				}
				b, err = json.Marshal(events)
				if err != nil {
					s.logger.Printf("failed marshaling events: %v", err)
					return
				}
			case *gnmi.SubscribeResponse_SyncResponse:
				s.logger.Printf("received subscribe syncResponse with %v", meta)
			case *gnmi.SubscribeResponse_Error:
				gnmiErr := sub.GetError()
				s.logger.Printf("received subscribe response error with %v, code=%d, message=%v, data=%v ",
					meta, gnmiErr.Code, gnmiErr.Message, gnmiErr.Data)
			}
		}
	}
	if err != nil {
		s.logger.Printf("failed marshaling event: %v", err)
		return
{'fraction_non_alphanumeric': 0.15663157894736843, 'fraction_numerical': 0.0010526315789473684, 'mean_word_length': 3.715136476426799, 'pattern_counts': {'":': 6, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2472147', 'n_tokens_mistral': 3574, 'n_tokens_neox': 3308, 'n_words': 678}