/*
 * Copyright 2021 PingCAP, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package main

import (
	"context"
	"crypto/md5"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"strings"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	k8s "k8s.io/client-go/kubernetes"
	"k8s.io/client-go/rest"

	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/sqs"
	_ "go.uber.org/automaxprocs"
	_ "gorm.io/driver/mysql"
)

func main() {
	var cfgPath string
	flag.StringVar(&cfgPath, "c", "./", "snapshoter -c=/path/to/config.yaml")
	flag.Parse()
	if err := InitConfig(cfgPath); err != nil {
		log.Fatalf("parse config error %+v", err)
	}
	// retrieve central secret from `central-secret-xxx`
	// under the central's namespace before connecting db if necessary.
	if cfg.Mysql.NeedMergePWD {
		k8sCfg, err := rest.InClusterConfig()
		if err != nil {
			log.Fatalf("init cluster config error %+v", err)
		}
		clientSet, err := k8s.NewForConfig(k8sCfg)
		if err != nil {
			log.Fatalf("new client set error %+v", err)
		}
		timeout := int64(5)
		sl, err := clientSet.CoreV1().Secrets(cfg.CentralNamespace).List(context.Background(), metav1.ListOptions{
			TimeoutSeconds: &timeout,
			FieldSelector:  fields.OneTermEqualSelector("type", "Opaque").String(),
		})
		if err != nil {
			log.Fatalf("list secert error %+v", err)
		}
		for _, s := range sl.Items {
			if !strings.HasPrefix(s.Name, "tidb-secret") {
				continue
			}
			cfg.Mysql.MergePWD(string(s.Data["central"]))
			break
		}
	}
	db, err := Dial(cfg.Mysql)
	if err != nil {
		log.Fatalf("dial database with dsn %+v error %+v", cfg.Mysql.DSN, err)
	}
	snapshoter, err := NewSnapshoter(db, cfg.Tables, 500)
	if err != nil {
		log.Fatalf("new snapshoter error %+v", err)
	}
	sess := session.Must(session.NewSessionWithOptions(session.Options{
		SharedConfigState: session.SharedConfigEnable,
	}))
	sqsCli := sqs.New(sess)
	sink := snapshoter.getSink()
	go func() {
		if err := snapshoter.collect(context.Background()); err != nil {
			log.Printf("snapshoter collect error %+v\n", err)
		}
	}()
	success := 0
	fail := 0
	for {
		select {
		case record, ok := <-sink:
			if !ok {
				log.Printf("success: %+v, fail: %+v, total: %+v\n", success, fail, success+fail)
				return
			}
			// send message info to sqs
			bs, err := json.Marshal(record.RowData)
			if err != nil {
				fail++
				continue
			}
			msg := new(sqs.SendMessageInput)
			msg.SetMessageBody(string(bs))
			msg.SetQueueUrl(cfg.SqsUrl)
			msg.SetMessageGroupId(record.MessageGroupID)
			dupID := fmt.Sprintf("%x", md5.Sum(bs))
			msg.SetMessageDeduplicationId(dupID)
			if _, err = sqsCli.SendMessage(msg); err != nil {
				log.Printf("send message error %+v\n", err)
				fail++
			} else {
				success++
			}
		}
	}
}