package events

import (
	"crypto/sha512"
	"encoding/hex"
	"fmt"
	"github.com/olivere/elastic/v7"
	log "github.com/sirupsen/logrus"
	"github.com/spf13/viper"
	"iwarma.ru/console/correlator/config"
	"iwarma.ru/console/correlator/es"
	"iwarma.ru/console/correlator/util"
	"testing"
	"time"
)

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

	normalizedIndexName := "test-normalized"
	aggregatedIndexName := "test-aggregated"
	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",
	}

	viper.Set(config.ElasticAggregatedIndexName, aggregatedIndexName)
	viper.Set(config.ElasticNormalizedIndexName, normalizedIndexName)
	viper.Set(config.AggregatedFields, eventsFields)

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

	err = ClearIndex(client, normalizedIndexName)
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	err = ClearIndex(client, aggregatedIndexName)
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	// Add documents
	goodCount := 10
	err = FillNormalizedEventsForAggregation(normalizedIndexName, goodCount, client)
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	// Prepare query
	query := elastic.NewMatchAllQuery()

	raw, errs1 := client.Query(normalizedIndexName, query)
	events, errs2 := ParseEvents(raw, errs1, 1)

	// Prepare store
	window := TimeWindow{Begin: time.Now().UTC(), End: time.Now().UTC().Add(time.Second * 10)}
	store := NewAggregatedEventStore(client, window)

	// Add events
	store.AddEvents(events)

	// Check for errors
	for err = range errs2 {
		if err != nil {
			t.Errorf("%v", err)
		}
	}

	// Now, change window to update all events
	store.UpdateWindow(TimeWindow{Begin: time.Now().UTC(), End: time.Now().UTC().Add(time.Second * 10)})

	// Need to wait for elastic finish
	time.Sleep(time.Second)

	// Now, we can check what we have
	if store.stat.EventsProcessed != uint64(goodCount) {
		t.Errorf("Bad event processed count. Expect %v, got %v", goodCount, store.stat.EventsProcessed)
	}

	// We create only one aggregated event from all our normalized events
	if store.stat.EventsAggregated != uint64(1) {
		t.Errorf("Bad event aggregated count. Expect %v, got %v", 1, store.stat.EventsAggregated)
	}

	// Need to check what we have in index
	raw, errs1 = client.Query(normalizedIndexName, query)
	normalizedCount := 0
	for range raw {
		normalizedCount++
	}

	// Check for errors
	for err = range errs1 {
		if err != nil {
			t.Errorf("%v", err)
		}
	}

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

	// Now check aggregated events
	raw, errs1 = client.Query(aggregatedIndexName, query)
	aggregatedCount := 0
	for range raw {
		aggregatedCount++
	}

	// Check for errors
	for err = range errs1 {
		if err != nil {
			t.Errorf("%v", err)
		}
	}

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

	// Clear index
	err = client.DeleteIndex(normalizedIndexName)

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

	err = client.DeleteIndex(aggregatedIndexName)
	if err != nil {
		t.Errorf("%v", err)
	}
}

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

	normalizedIndexName := "test-normalized"
	aggregatedIndexName := "test-aggregated"

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

	err = ClearIndex(client, normalizedIndexName)
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	err = ClearIndex(client, aggregatedIndexName)
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	// Add documents
	goodCount := 1000
	err = FillNormalizedEventsForAggregation(normalizedIndexName, goodCount, client)
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	log.Info("Index is ready")

	// Prepare query
	query := elastic.NewMatchAllQuery()

	raw, errs1 := client.Query(normalizedIndexName, query)
	events, errs2 := ParseEvents(raw, errs1, 10)

	// Prepare store
	window := TimeWindow{Begin: time.Now().UTC(), End: time.Now().UTC().Add(time.Second * 10)}
	store := NewAggregatedEventStore(client, window)

	log.Info("Add event start")

	start := time.Now()
	// Add events
	store.AddEvents(events)

	log.Infof("Add event finish. Elapsed %v", time.Since(start))

	// Check for errors
	for err = range errs2 {
		if err != nil {
			t.Errorf("%v", err)
		}
	}

	log.Info("Bulk flush start")
	start = time.Now()
	// Now, change window to update all events
	store.UpdateWindow(TimeWindow{Begin: time.Now().UTC(), End: time.Now().UTC().Add(time.Second * 10)})
	log.Infof("Bulk flush finish. Elapsed %v", time.Since(start))

	// Need to wait for elastic finish
	time.Sleep(time.Second * 30)

	// Now, we can check what we have
	if store.stat.EventsProcessed != uint64(goodCount) {
		t.Errorf("Bad event processed count. Expect %v, got %v", goodCount, store.stat.EventsProcessed)
	}

	// We create only one aggregated event from all our normalized events
	if store.stat.EventsAggregated != uint64(1) {
		t.Errorf("Bad event aggregated count. Expect %v, got %v", 1, store.stat.EventsAggregated)
	}

	// Need to check what we have in index
	raw, errs1 = client.Query(normalizedIndexName, query)
	normalizedCount := 0
	for range raw {
		normalizedCount++
	}

	// Check for errors
	for err = range errs1 {
		if err != nil {
			t.Errorf("%v", err)
		}
	}

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

	// Now check aggregated events
	raw, errs1 = client.Query(aggregatedIndexName, query)
	aggregatedCount := 0
	for range raw {
		aggregatedCount++
	}

	// Check for errors
	for err = range errs1 {
		if err != nil {
			t.Errorf("%v", err)
		}
	}

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

}

//func TestAggregatedEventStore_UpdateWindow(t *testing.T) {
//	util.SetupTest(t)
//	defer util.TearDownTest(t)
//
//	normalizedIndexName := "test-normalized"
//	aggregatedIndexName := "test-aggregated"
//
//	client, err := es.NewElastic()
//	if err != nil {
//		t.Errorf("%v", err)
//		return
//	}
//
//	err = ClearIndex(client, normalizedIndexName)
//	if err != nil {
//		t.Errorf("%v", err)
//		return
//	}
//
//	err = ClearIndex(client, aggregatedIndexName)
//	if err != nil {
//		t.Errorf("%v", err)
//		return
//	}
//
//	addEvents := func(count int) int {
//		err = FillNormalizedEventsForAggregation(normalizedIndexName, count, client)
//		if err != nil {
//			t.Errorf("%v", err)
//			return 0
//		}
//		return count
//	}
//
//	// Add documents
//	goodCount := 100
//	totalEvents := 0
//	aggregatedEventCount := 10
//	log.Info("Index is ready")
//
//	// Prepare store
//	window := TimeWindow{Begin: time.Now().UTC(), End: time.Now().UTC().Add(time.Second * 10)}
//	store := NewAggregatedEventStore(client, window)
//
//	// Prepare query
//	query := elastic.NewMatchAllQuery()
//	for i := 0; i < aggregatedEventCount; i++ {
//		// Add documents
//		totalEvents += addEvents(goodCount)
//
//		// Move aggregation window
//		log.Info("Bulk flush start")
//		start := time.Now()
//		// Now, change window to update all events
//		store.UpdateWindow(TimeWindow{Begin: time.Now().UTC(), End: time.Now().UTC().Add(time.Second * 10)})
//		log.Infof("Bulk flush finish. Elapsed %v", time.Since(start))
//
//		// Query for new events
//		raw, errs1 := client.Query(normalizedIndexName, query)
//		events, errs2 := ParseEvents(raw, errs1, 10)
//		log.Info("Add event start")
//		start = time.Now()
//		// Add events
//		store.AddEvents(events)
//		log.Infof("Add event finish. Elapsed %v", time.Since(start))
//
//		for err = range errs2 {
//			t.Errorf("%v", err)
//		}
//
//		// Need to wait for elastic finish
//		time.Sleep(time.Second * 20)
//	}
//
//	// Move aggregation window
//	log.Info("Bulk flush start")
//	start := time.Now()
//	// Now, change window to update all events
//	store.UpdateWindow(TimeWindow{Begin: time.Now().UTC(), End: time.Now().UTC().Add(time.Second * 10)})
//	log.Infof("Bulk flush finish. Elapsed %v", time.Since(start))
//
//	// Need to wait for elastic finish
//	time.Sleep(time.Second * 20)
//
//	// Now, we can check what we have
//	if store.stat.EventsProcessed != uint64(totalEvents) {
//		t.Errorf("Bad event processed count. Expect %v, got %v", goodCount, store.stat.EventsProcessed)
//	}
//
//	// We create only one aggregated event from all our normalized events
//	if store.stat.EventsAggregated != uint64(aggregatedEventCount) {
//		t.Errorf("Bad event aggregated count. Expect %v, got %v", 1, store.stat.EventsAggregated)
//	}
//
//	// Need to check what we have in index
//	raw, errs1 := client.Query(normalizedIndexName, query)
//	normalizedCount := 0
//	for range raw {
//		normalizedCount++
//	}
//
//	// Check for errors
//	for err = range errs1 {
//		if err != nil {
//			t.Errorf("%v", err)
//		}
//	}
//
//	if normalizedCount != 0 {
//		t.Errorf("Not all normalized events were deleted. We still have %v of them", normalizedCount)
//	}
//
//	// Now check aggregated events
//	raw, errs1 = client.Query(aggregatedIndexName, query)
//	aggregatedCount := 0
//	for range raw {
//		aggregatedCount++
//	}
//
//	// Check for errors
//	for err = range errs1 {
//		if err != nil {
//			t.Errorf("%v", err)
//		}
//	}
//
//	if aggregatedCount != aggregatedEventCount {
//		t.Errorf("Got wrong aggregated events count. Expect %v, got %v", aggregatedEventCount, aggregatedCount)
//	}
//
//}

// TestHashFuncWithAllFields The check is performed on all standard event fields
func TestHashFuncWithAllFields(t *testing.T) {
	util.SetupTest(t)
	defer util.TearDownTest(t)
	normalizedIndexName := "test-normalized"
	aggregatedIndexName := "test-aggregated"
	eventsFields := []string{
		"event_severity",
		"event_protocol",
		"message",
		"device_vendor",
		"device_product",
		"device_action",
		"message",
		"device_version",
		"sign_id",
		"sign_category",
		"sign_subcategory",
		"application",
		"source_ip",
		"source_host",
		"source_port",
		"destination_ip",
		"destination_host",
		"destination_port",
		"destination_user",
	}

	viper.Set(config.ElasticAggregatedIndexName, aggregatedIndexName)
	viper.Set(config.ElasticNormalizedIndexName, normalizedIndexName)
	viper.Set(config.AggregatedFields, eventsFields)

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

	err = ClearIndex(client, normalizedIndexName)
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	err = ClearIndex(client, aggregatedIndexName)
	if err != nil {
		t.Errorf("%v", err)
		return
	}
	goodCount := 10
	err = FillNormalizedEventsForAggregation(normalizedIndexName, goodCount, client)
	if err != nil {
		t.Errorf("%v", err)
		return
	}
	query := elastic.NewMatchAllQuery()

	raw, errs1 := client.Query(normalizedIndexName, query)
	events, errs2 := ParseEvents(raw, errs1, 10)

	for item := range errs2 {
		if item != nil {
			t.Errorf("%v", item)
		}
	}
	var eventString string
	for item := range events {

		for _, field := range eventsFields {
			eventString += fmt.Sprintf("%v", item.GetValue(field))
		}
		fieldStr := sha512.Sum512_256([]byte(fmt.Sprintf("%v", eventString)))
		//
		hashOld := hex.EncodeToString(fieldStr[:])
		hashNew := item.Hash()
		if hashOld != hashNew {
			t.Errorf("%v[%v->%v]", "The hashes don't match", hashNew, hashOld)
		}
		eventString = ""

	}

}

//TestHashFuncWithPartFields The check is performed with an incomplete list of fields
func TestHashFuncWithPartFields(t *testing.T) {
	util.SetupTest(t)
	defer util.TearDownTest(t)
	normalizedIndexName := "test-normalized"
	aggregatedIndexName := "test-aggregated"
	eventsFields := []string{
		"event_severity",
		"event_protocol",
		"device_vendor",
		"device_product",
		"device_action",
		"device_version",
		"sign_id",
		"sign_category",
		"sign_subcategory",
		"application",
		"source_ip",
		"source_host",
	}

	viper.Set(config.ElasticAggregatedIndexName, aggregatedIndexName)
	viper.Set(config.ElasticNormalizedIndexName, normalizedIndexName)
	viper.Set(config.AggregatedFields, eventsFields)

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

	err = ClearIndex(client, normalizedIndexName)
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	err = ClearIndex(client, aggregatedIndexName)
	if err != nil {
		t.Errorf("%v", err)
		return
	}
	goodCount := 10
	err = FillNormalizedEventsForAggregation(normalizedIndexName, goodCount, client)
	if err != nil {
		t.Errorf("%v", err)

	}
	query := elastic.NewMatchAllQuery()

	raw, errs1 := client.Query(normalizedIndexName, query)
	events, errs2 := ParseEvents(raw, errs1, 10)

	for err = range errs2 {
		if err != nil {
			t.Errorf("%v", err)
		}
	}
	var eventString string
	for item := range events {

		for _, field := range eventsFields {
			eventString += fmt.Sprintf("%v", item.GetValue(field))
		}
		fieldStr := sha512.Sum512_256([]byte(fmt.Sprintf("%v", eventString)))
		hashOld := hex.EncodeToString(fieldStr[:])
		hashNew := item.Hash()
		if hashOld != hashNew {
			t.Errorf("%v[%v->%v]", "The hashes don't match", hashNew, hashOld)
		}
		eventString = ""
	}
}

// TestHashFuncWithEmptyFieldsList The test works completely without aggregation fields.
//In this case, the hash is calculated from an empty string
func TestHashFuncWithEmptyFieldsList(t *testing.T) {
	util.SetupTest(t)
	defer util.TearDownTest(t)

	normalizedIndexName := "test-normalized"
	aggregatedIndexName := "test-aggregated"
	viper.Set(config.ElasticAggregatedIndexName, aggregatedIndexName)
	viper.Set(config.ElasticNormalizedIndexName, normalizedIndexName)
	viper.Set(config.AggregatedFields, make([]string, 0))

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

	err = ClearIndex(client, normalizedIndexName)
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	err = ClearIndex(client, aggregatedIndexName)
	if err != nil {
		t.Errorf("%v", err)
		return
	}
	goodCount := 10
	err = FillNormalizedEventsForAggregation(normalizedIndexName, goodCount, client)
	if err != nil {
		t.Errorf("%v", err)

	}
	query := elastic.NewMatchAllQuery()

	raw, errs1 := client.Query(normalizedIndexName, query)
	events, errs2 := ParseEvents(raw, errs1, 10)

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

	for item := range events {
		fieldStr := sha512.Sum512_256([]byte(fmt.Sprint("")))
		hashOld := hex.EncodeToString(fieldStr[:])
		hashNew := item.Hash()
		if hashOld != hashNew {
			t.Errorf("%v[%v->%v]", "The hashes don't match", hashNew, hashOld)
		}
	}
}
