package aggregator

import (
	"encoding/json"
	"fmt"
	"github.com/olivere/elastic/v7"
	"github.com/spf13/viper"
	"iwarma.ru/console/correlator/config"
	"iwarma.ru/console/correlator/es"
	"iwarma.ru/console/correlator/events"
	"iwarma.ru/console/correlator/util"
	"testing"
	"time"
)

func TestAggregator(t *testing.T) {
	util.SetupTest(t)
	defer util.TearDownTest(t)

	client, err := es.NewElastic()
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	err = events.ClearIndex(client, viper.GetString(config.ElasticNormalizedIndexName), viper.GetString(config.ElasticAggregatedIndexName))
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	stop := make(chan interface{})
	duration := time.Second * 10

	eventsStore := events.NewAggregatedEventStore(client, events.TimeWindow{
		Begin: time.Now().UTC(),
		End:   time.Now().UTC().Add(duration),
	})

	go func() {
		Aggregator(eventsStore, duration, stop)
	}()

	N := 3
	for i := 0; i < N; i++ {
		err = events.FillNormalizedEventsForAggregation(viper.GetString(config.ElasticNormalizedIndexName), N/(i+1), client)
		if err != nil {
			t.Errorf("%v", err)
			return
		}
		time.Sleep(duration + time.Second)
	}

	stop <- struct{}{}

	time.Sleep(time.Second * 5)

	// Let's check what we have
	eventsResult, errs := client.Query(viper.GetString(config.ElasticNormalizedIndexName), elastic.NewMatchAllQuery())
	normalizedCount := 0
	for range eventsResult {
		normalizedCount++
	}

	if normalizedCount != 0 {
		t.Errorf("Not all normalized events were deleted. Still have %v", normalizedCount)
	}

	for err = range errs {
		if err != nil {
			t.Errorf("%v", err)
		}
	}

	// And Aggregated
	eventsResult, errs = client.Query(viper.GetString(config.ElasticAggregatedIndexName), elastic.NewMatchAllQuery())
	aggregatedCount := 0
	for range eventsResult {
		aggregatedCount++
	}

	if aggregatedCount != N {
		t.Errorf("Bad aggregated events count. Expect %v, got %v", N, aggregatedCount)
	}

	for err = range errs {
		if err != nil {
			t.Errorf("%v", err)
		}
	}

}

func TestAggregatorSeveralNormalizedIndex(t *testing.T) {
	util.SetupTest(t)
	defer util.TearDownTest(t)

	index1 := "arma-1"
	index2 := "arma-2"

	viper.Set(config.ElasticNormalizedIndexName, "arma-*")

	client, err := es.NewElastic()
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	err = events.ClearIndex(client, viper.GetString(config.ElasticNormalizedIndexName), viper.GetString(config.ElasticAggregatedIndexName))
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	stop := make(chan interface{})
	duration := time.Second * 10

	eventsStore := events.NewAggregatedEventStore(client, events.TimeWindow{
		Begin: time.Now().UTC(),
		End:   time.Now().UTC().Add(duration),
	})

	go func() {
		Aggregator(eventsStore, duration, stop)
	}()

	N := 3
	for i := 0; i < N; i++ {
		err = events.FillNormalizedEventsForAggregation(index1, N, client)
		if err != nil {
			t.Errorf("%v", err)
			return
		}

		err = events.FillNormalizedEventsForAggregation(index2, N, client)
		if err != nil {
			t.Errorf("%v", err)
			return
		}
	}

	time.Sleep(duration)

	stop <- struct{}{}
	eventsStore.SendUpdateBulk()

	// Let's check correlator stat
	stat := eventsStore.GetStat()

	if stat.EventsProcessed != uint64(N*N*2) {
		t.Errorf("Bad event processed count: %v", stat.EventsProcessed)
	}

	// All normalized events must be aggregated in one agregated event
	if stat.EventsAggregated != 1 {
		t.Errorf("Bad event aggregated count. Expect %v, got %v", 1, stat.EventsAggregated)
	}

	// Now, check indices
	count, err := client.CountDocuments(index1, index1)
	if err != nil {
		t.Errorf("%v", err)
	}

	if count != 0 {
		t.Errorf("Not all normalized events were deleted")
	}

	count, err = client.CountDocuments(viper.GetString(config.ElasticAggregatedIndexName))
	if err != nil {
		t.Errorf("%v", err)
	}

	if count != 1 {
		t.Errorf("Bad aggregated events count. Expect %v, got %v", 1, count)
	}

}

func TestSendUpdateBulk(t *testing.T) {
	util.SetupTest(t)
	defer util.TearDownTest(t)
	eventsFields := []string{
		"event_severity",
		"event_protocol",
		"message",
		"device_vendor",
		"device_product",
		"device_action",
		"device_version",
		"device_timezone",
		"sign_id",
		"sign_category",
		"sign_subcategory",
		"application",
		"source_ip",
		"source_host",
		"source_port",
		"source_mac",
		"source_timezone",
		"source_software",
		"source_action",
		"destination_ip",
		"destination_mac",
		"destination_timezone",
		"destination_software",
		"destination_action",
		"destination_host",
		"destination_port",
		"destination_user",
		"cs1",
		"cs1Label",
		"cs2",
		"cs2Label",
		"object_type",
		"event_last",
		"source_user",
		"@created",
		"event_hash",
		"index",
		"timestamp",
		"aggregated_id",
		"event_src_msg",
		"type",
		"event_count",
		"event_hash",
		"event_timestamp",
		"timestamp",
		"event_id",
		"event_first",
	}

	client, err := es.NewElastic()
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	err = events.ClearIndex(client, viper.GetString(config.ElasticNormalizedIndexName), viper.GetString(config.ElasticAggregatedIndexName))
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	stop := make(chan interface{})
	duration := time.Second * 10

	eventsStore := events.NewAggregatedEventStore(client, events.TimeWindow{
		Begin: time.Now().UTC(),
		End:   time.Now().UTC().Add(duration),
	})

	go func() {
		Aggregator(eventsStore, duration, stop)
	}()

	N := 3
	for i := 0; i < N; i++ {
		err = events.FillNormalizedEventsForAggregation(viper.GetString(config.ElasticNormalizedIndexName), N/(i+1), client)
		if err != nil {
			t.Errorf("%v", err)
			return
		}
		time.Sleep(duration + time.Second)
	}

	stop <- struct{}{}
	eventsResult, errs := client.Query(viper.GetString(config.ElasticAggregatedIndexName), elastic.NewMatchAllQuery())
	errorsCount := 0
	for item := range eventsResult {
		var event events.Event
		err := json.Unmarshal(item.Source, &event)
		if err != nil {
			t.Errorf("%v", err)
		}
		for key, _ := range event {
			if !contains(eventsFields, key) {
				errorsCount++
				t.Errorf("Bad field -> %v", key)
			}
		}

	}
	if errorsCount != 0 {
		t.Errorf("Bad fields in events")
		return
	}
	for err = range errs {
		if err != nil {
			t.Errorf("%v", err)
		}
	}

	// Clear after test
	err = events.ClearIndex(client, viper.GetString(config.ElasticNormalizedIndexName), viper.GetString(config.ElasticAggregatedIndexName))
	if err != nil {
		t.Errorf("%v", err)
	}
}

func contains(fields []string, item string) bool {
	for _, field := range fields {
		if field == item {
			fmt.Printf("%v - %v \n", field, item)
			return true
		}
	}
	return false
}
