package rules

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

// Check that NewRuleStore create and init new store
func TestNewRuleStore(t *testing.T) {
	util.SetupTest(t)
	defer util.TearDownTest(t)

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

	store := NewRuleStore(client)

	if len(store.rules) != 0 {
		t.Errorf("Got bad rules count: %v", len(store.rules))
	}

	if store.client == nil {
		t.Error("Client was not set")
	}

	if store.stat.AverageRuleTime == nil {
		t.Error("Store stat rules map was not created")
	}

	if store.cl == nil {
		t.Errorf("Logger wasn't created")
	}
}

// Check that clear store function is working
func TestClearStore(t *testing.T) {
	util.SetupTest(t)
	defer util.TearDownTest(t)

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

	store := NewRuleStore(client)

	var rule Rule

	store.AddRules(rule)

	if len(store.rules) != 1 {
		t.Errorf("Bad store sizr after add: %v", len(store.rules))
	}

	store.ClearStore()

	if len(store.rules) != 0 {
		t.Error("Can't clear store")
	}
}

// Check that we can update rule
func TestUpdateRule(t *testing.T) {
	util.SetupTest(t)
	defer util.TearDownTest(t)
	rule := Rule{Id: "10", Name: "Test"}

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

	store := NewRuleStore(client)

	store.AddRules(rule)

	rule2 := Rule{Id: "10", Name: "Test2222"}
	store.AddRules(rule2)

	if len(store.rules) != 1 {
		t.Errorf("Bad rules count, need 1, got %v", len(store.rules))
	}

	if store.rules[0].Name != "Test2222" {
		t.Error("Rule was not updated")
	}
}

// Check thar we can add several rules at a time
func TestAddRules(t *testing.T) {
	util.SetupTest(t)
	defer util.TearDownTest(t)

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

	store := NewRuleStore(client)

	rules := make([]Rule, 10)
	for i := 0; i < 10; i++ {
		rules[i] = Rule{Id: fmt.Sprintf("%v", i), Name: fmt.Sprintf("Test_%v", i)}
	}

	store.AddRules(rules...)

	if len(store.rules) != 10 {
		t.Errorf("Bad rules count, neet 10, got %v", len(store.rules))
	}
}

// Check that we can update several rules
func TestUpdateRules(t *testing.T) {
	util.SetupTest(t)
	defer util.TearDownTest(t)

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

	store := NewRuleStore(client)

	rules := make([]Rule, 10)
	for i := 0; i < 10; i++ {
		rules[i] = Rule{Id: fmt.Sprintf("%v", i), Name: fmt.Sprintf("Test_%v", i)}
	}

	store.AddRules(rules...)

	for i := 0; i < len(store.rules); i++ {
		rules[i].Name = "DDDD"
	}

	store.AddRules(rules...)

	for _, cur := range store.rules {
		if cur.Name != "DDDD" {
			t.Errorf("Rule %v wasn't updated", cur)
		}
	}
}

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

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

	store := NewRuleStore(client)

	stat := store.GetStat()

	if stat.AverageRuleTime == nil {
		t.Error("Got nil AverageRuleTime")
	}
}

// Check that Do function works with single action mode
func TestRule_Do(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, events.GetAggregatedIndexName())
	if err != nil {
		t.Errorf("%v", err)
	}

	err = events.FillAggregatedEvents(events.GetAggregatedIndexName(), 10, client)
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	store := NewRuleStore(client)

	actions := make([]Action, 1)
	action := TestAction{
		PerformError:        false,
		ParseInterfaceError: false,
		MarshalError:        false,
		UnmarshalError:      false,
	}
	actions[0] = &action

	store.AddRules(Rule{
		Id:          "1",
		Name:        "TestRule",
		Description: "Some long description",
		Multi:       false,
		Depth:       time.Second * 30,
		Predicate:   NewPredicate("", "device_vendor:TestDevice"),
		Actions:     actions,
	})

	PerformCountClear()
	EventsProcessedClear()

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

	time.Sleep(time.Second * 2)

	goodCount := 1
	if GetPerformCount() != goodCount {
		t.Errorf("Bad perform count. Expect %v, got %v", goodCount, GetPerformCount())
	}

	if GetEventsProcessed() != goodCount {
		t.Errorf("Bad events processed count. Expect %v, got %v", goodCount, GetEventsProcessed())
	}
}

// Check that Do function works with single action mode and multiple rules
func TestMultipleRule_Do(t *testing.T) {
	util.PrepareElastic()
	index := "test-aggregated"
	client, err := es.NewElastic()
	if err != nil {
		t.Errorf("%v", err)
		return
	}

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

	err = events.FillAggregatedEvents(events.GetAggregatedIndexName(), 2, client)
	if err != nil {
		t.Errorf("%v", err)
		return
	}

	store := NewRuleStore(client)

	actions := make([]Action, 1)
	action := TestAction{
		PerformError:        false,
		ParseInterfaceError: false,
		MarshalError:        false,
		UnmarshalError:      false,
	}
	actions[0] = &action

	store.AddRules(Rule{
		Id:          "1",
		Name:        "TestRule #1",
		Description: "Some long description",
		Multi:       false,
		Depth:       time.Second * 30,
		Predicate:   NewPredicate("", "device_vendor:TestDevice"),
		Actions:     actions,
	},
		Rule{
			Id:          "2",
			Name:        "TestRule #2",
			Description: "Some long description",
			Multi:       false,
			Depth:       time.Second * 30,
			Predicate:   NewPredicate("", "source_ip:127.0.0.1"),
			Actions:     actions,
		},
	)

	PerformCountClear()
	EventsProcessedClear()

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

	time.Sleep(time.Second * 2)

	goodCount := 2
	if GetPerformCount() != goodCount {
		t.Errorf("Bad perform count. Expect %v, got %v", goodCount, GetPerformCount())
	}
	query := elastic.NewMatchAllQuery()

	raw, errs1 := client.Query(index, query)
	events, _ := events.ParseEvents(raw, errs1, 2)

	errorsCount := 0
	for event := range events {
		tags, ok := event.GetValue("rule_tag").([]string)
		if ok {
			if len(tags) != 2 {
				errorsCount += 1
			}
		}

	}
	if errorsCount != 0 {
		t.Errorf("Got errors while count rule tags . Expect 0 errors, got %v", errorsCount)
	}

}

// Check that Do function works with multi action mode
func TestRule_Do_multi(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, events.GetAggregatedIndexName())
	if err != nil {
		t.Errorf("%v", err)
	}

	// Events count
	N := 10

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

	store := NewRuleStore(client)

	actions := make([]Action, 1)
	action := TestAction{
		PerformError:        false,
		ParseInterfaceError: false,
		MarshalError:        false,
		UnmarshalError:      false,
	}
	actions[0] = &action

	store.AddRules(Rule{
		Id:          "1",
		Name:        "TestRule",
		Description: "Some long description",
		Multi:       true,
		Depth:       time.Second * 30,
		Predicate:   NewPredicate("", "device_vendor:TestDevice"),
		Actions:     actions,
	})

	PerformCountClear()
	EventsProcessedClear()

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

	time.Sleep(time.Second * 2)

	goodCount := 1
	if GetPerformCount() != goodCount {
		t.Errorf("Bad perform count. Expect %v, got %v", goodCount, GetPerformCount())
	}

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