package main

import (
	"flag"
	"fmt"
	"net/http"
	"os"
	"runtime"
	"runtime/pprof"
	"strings"
	"sync"
	"time"

	"iwarma.ru/console/correlator/aggregator"
	"iwarma.ru/console/correlator/api"
	"iwarma.ru/console/correlator/config"
	"iwarma.ru/console/correlator/correlator"
	"iwarma.ru/console/correlator/es"
	"iwarma.ru/console/correlator/events"
	"iwarma.ru/console/correlator/mapping"
	"iwarma.ru/console/correlator/rules"

	"github.com/spf13/viper"

	log "github.com/sirupsen/logrus"

	"github.com/gorilla/mux"
)

var (
	VERSION = "1.1.6"
)

var (
	configFile    = flag.String("config", "config.json", "Path to config file")
	cpuprofile    = flag.String("cpuprofile", "", "write cpu profile to `file`")
	memprofile    = flag.String("memprofile", "", "write memory profile to `file`")
	noMappingFlag = flag.Bool("no-mapping", false, "Don't apply mapping")
	versionFlag   = flag.Bool("version", false, "Show version and exit")
	noCorrelator  = flag.Bool("no-correlator", false, "Disable correlator")
	noAggregator  = flag.Bool("no-aggregator", false, "Disable aggregator")
)

func main() {

	cl := log.WithFields(log.Fields{"part": "main", "ver": VERSION})

	//##############################
	//    Read and parse config
	//##############################

	flag.Parse()

	if *versionFlag == true {
		cl.Infof("Correlator version %v", VERSION)
		return
	}

	// Viper config
	err := config.ReadCfg(*configFile)
	if err != nil {
		cl.Fatalf("%v", err)
	}

	//##############################
	//       Init logging
	//##############################
	err = initLogging()
	if err != nil {
		cl.Fatalf("%v", err)
	}

	//##############################
	//       Profiling
	//##############################
	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			cl.Fatal("could not create CPU profile: ", err)
		}
		defer f.Close() // error handling omitted for example
		if err := pprof.StartCPUProfile(f); err != nil {
			cl.Fatal("could not start CPU profile: ", err)
		}
		cl.Infof("Start CPU profile")
		defer pprof.StopCPUProfile()
	}

	//##############################
	//     Memory profiling
	//##############################
	if *memprofile != "" {
		f, err := os.Create(*memprofile)
		if err != nil {
			cl.Fatal("could not create memory profile: ", err)
		}
		defer f.Close() // error handling omitted for example
		runtime.GC()    // get up-to-date statistics
		if err := pprof.WriteHeapProfile(f); err != nil {
			cl.Fatal("could not write memory profile: ", err)
		}

		cl.Infof("Start Memory profile")
	}

	//##############################
	//   Connecting to elastic
	//##############################
	client, err := es.NewElastic()
	if err != nil {
		cl.Fatalf("Can't conect to elastic: %v", err)
	}

	//##############################
	//   Prepare store
	//##############################

	// Correlator's store
	ruleStore := rules.NewRuleStore(client)

	// Aggregator's store
	window := viper.GetDuration(config.AggregatorWindow)
	eventsStore := events.NewAggregatedEventStore(client, events.TimeWindow{
		Begin: time.Now().UTC(),
		End:   time.Now().UTC().Add(window),
	})

	//##############################
	//           Api
	//##############################
	cl.Infof("Starting API on port %v", viper.GetInt(config.ApiPort))
	router := mux.NewRouter()
	router.HandleFunc("/add_many/", api.AddRulesHandler(ruleStore)).Methods("POST")
	router.HandleFunc("/clear/", api.ClearStoreHandler(ruleStore)).Methods("GET")
	router.HandleFunc("/stats/", api.StatHandler(eventsStore, ruleStore)).Methods("GET")
	router.Use(api.LogUrlMiddleware)

	// TODO: Add error control
	go http.ListenAndServe(fmt.Sprintf(":%v", viper.GetInt(config.ApiPort)), router)
	cl.Infof("API server started on port %v", viper.GetInt(config.ApiPort))

	//##############################
	//     Prepare elastic
	//##############################
	cl.Info("Checking indices")

	var normalizedEventIndexName, aggregatedEventIndexName string
	if strings.Contains(events.GetNormalizedIndexName(), "*") || strings.Contains(events.GetAggregatedIndexName(), "*") {
		// Need to replace * with current date
		normalizedEventIndexName = strings.Replace(events.GetNormalizedIndexName(), "*", time.Now().UTC().Format("2006.01.02"), 1)
		aggregatedEventIndexName = strings.Replace(events.GetAggregatedIndexName(), "*", time.Now().UTC().Format("2006.01.02"), 1)
	} else {
		normalizedEventIndexName = events.GetNormalizedIndexName()
		aggregatedEventIndexName = events.GetAggregatedIndexName()
	}

	err = client.CheckAndCreateIndex(normalizedEventIndexName, aggregatedEventIndexName)
	if err != nil {
		cl.Fatalf("Can't check index for normalized events: %v", err)
	}

	//##############################
	//     Mapping
	//##############################
	if *noMappingFlag {
		cl.Info("No mapping will be applied")
	} else {

		var mappingArray []mapping.Mapping
		// See https://github.com/spf13/viper/issues/196
		err := viper.UnmarshalKey("mapping", &mappingArray)
		if err != nil {
			cl.Errorf("Can't read mapping from config")
		}

		err = mapping.UpdateMapping(&mappingArray)
		if err != nil {
			cl.Errorf("Can't apply mapping: %v", err)
		}
	}

	cl.Info("Indices are ready")

	//##############################
	//    Debug dumping prepare
	//##############################
	if viper.GetBool(config.DebugDumpRequest) {
		if _, err := os.Stat(viper.GetString(config.DebugDumpPath)); os.IsNotExist(err) {
			log.Fatalf("Dump path does not exist")
		}
	}

	var wg sync.WaitGroup

	wgCount := 2
	if *noAggregator == true {
		wgCount -= 1
	}

	if *noCorrelator == true {
		wgCount -= 1
	}

	wg.Add(wgCount)

	//##############################
	//         Aggregator
	//##############################

	aggregatorStop := make(chan interface{})
	if *noAggregator == false {
		go func() {
			aggregator.Aggregator(eventsStore, window, aggregatorStop)
			wg.Done()
		}()
	}

	//##############################
	//         Correlator
	//##############################

	correlatorStop := make(chan interface{})
	if *noCorrelator == false {
		go func() {
			correlator.Correlator(ruleStore, correlatorStop)
			wg.Done()
		}()
	}

	wg.Wait()
}
