package main

import (
	"context"
	"database/sql"
	"flag"
	"fmt"
	"os"
	"os/signal"
	"strings"
	"syscall"

	_ "github.com/go-sql-driver/mysql"
	"github.com/ngaut/log"
	"github.com/pingcap/errors"
	"github.com/pingcap/schrodinger/pkg/test_client"
)

type ScopeModifier string

const (
	GLOBAL_KEYWORD          ScopeModifier = "GLOBAL "
	GLOBAL_ANNOTATION       ScopeModifier = "@@GLOBAL."
	SESSION_DEFAULT         ScopeModifier = ""
	SESSION_KEYWORD         ScopeModifier = "SESSION "
	SESSION_ANNOTATION      ScopeModifier = "@@SESSION."
	SESSION_ANNOTATION_ABBR ScopeModifier = "@@"
)

func StrQuote(v string) string {
	return "'" + v + "'"
}

type SysVar struct {
	Name    string
	Default string
	Escape  func(string) string
}

func (v SysVar) Set(db *sql.DB, value string) error {
	return v.SetInScope(db, value, SESSION_DEFAULT)
}

func (v SysVar) SetDefault(db *sql.DB) error {
	err := v.Set(db, v.Default)
	if err != nil {
		return err
	}
	return v.SetInScope(db, v.Default, GLOBAL_ANNOTATION)
}

func (v SysVar) SetInScope(db *sql.DB, value string, scope ScopeModifier) error {
	if v.Escape != nil {
		value = v.Escape(value)
	}
	stmt := "SET " + string(scope) + v.Name + " = " + value
	//client.Info(">", stmt)
	_, err := db.Exec(stmt)
	return errors.Annotate(err, "set var by "+stmt)
}

func (v SysVar) WithValue(db *sql.DB, value string, cb func(db *sql.DB) error) (err error) {
	return v.WithValueInScope(db, value, SESSION_DEFAULT, cb)
}

func (v SysVar) WithValueInScope(db *sql.DB, value string, scope ScopeModifier, cb func(db *sql.DB) error) error {
	var orig string
	var err error
	var ann ScopeModifier
	switch scope {
	case GLOBAL_KEYWORD, GLOBAL_ANNOTATION:
		orig, err = v.GetGlobal(db)
		ann = GLOBAL_ANNOTATION
	default:
		orig, err = v.Get(db)
		ann = SESSION_ANNOTATION
	}
	if err != nil {
		return err
	}
	err = v.SetInScope(db, value, scope)
	if err != nil {
		return err
	}

	err = cb(db)
	if err != nil {
		return err
	}
	err = v.SetInScope(db, orig, ann)
	if err != nil {
		return errors.Annotate(err, "recover var")
	}

	return nil
}

func (v SysVar) Get(db *sql.DB) (string, error) {
	var val string
	stmt := "SELECT " + string(SESSION_ANNOTATION) + v.Name
	err := db.QueryRow(stmt).Scan(&val)
	return val, errors.Annotate(err, "get var by "+stmt)
}

func (v SysVar) GetGlobal(db *sql.DB) (string, error) {
	var val string
	stmt := "SELECT " + string(GLOBAL_ANNOTATION) + v.Name
	err := db.QueryRow(stmt).Scan(&val)
	return val, errors.Annotate(err, "get var by "+stmt)
}

type TestCase interface {
	Init(dsn string) error
	Run() error
}

type BaseTestCase struct {
	DSN string
	Var SysVar
}

func (t *BaseTestCase) Init(dsn string) error {
	t.DSN = dsn
	return nil
}

func (t *BaseTestCase) Run() error {
	return t.TestBasic(t.Var.Default)
}

func (t *BaseTestCase) TestBasic(val string) error {
	db1, err := sql.Open("mysql", t.DSN)
	if err != nil {
		return errors.Annotate(err, "open database")
	}
	defer db1.Close()

	// test default value
	orig, err := t.Var.Get(db1)
	if err != nil {
		return err
	}
	if orig != t.Var.Default {
		client.Warnf("[sysvars] default value mismatch: exp=%s act=%s", t.Var.Default, orig)
	}

	// test set by different ways
	for _, c := range []struct {
		scope ScopeModifier
		get   func(db *sql.DB) (string, error)
	}{
		{GLOBAL_KEYWORD, t.Var.GetGlobal},
		{GLOBAL_ANNOTATION, t.Var.GetGlobal},
		{SESSION_DEFAULT, t.Var.Get},
		{SESSION_KEYWORD, t.Var.Get},
		{SESSION_ANNOTATION, t.Var.Get},
		{SESSION_ANNOTATION_ABBR, t.Var.Get},
	} {
		err := t.Var.WithValueInScope(db1, val, c.scope, func(db *sql.DB) error {
			act, err := c.get(db)
			if err != nil {
				return err
			}
			if val != act {
				client.Errorf("failed to set variable: exp=%s act=%s modifier=%s",
					val, act, c.scope)
				return errors.New("failed to set var")
			}
			return nil
		})
		if err != nil {
			return errors.Annotate(err, "set "+string(c.scope)+t.Var.Name+" to "+val)
		}
	}

	// test session scope
	if orig != val {
		db2, err := sql.Open("mysql", t.DSN)
		if err != nil {
			return errors.Annotate(err, "open database")
		}
		defer db2.Close()

		err = t.Var.WithValue(db1, val, func(db *sql.DB) error {
			v, err := t.Var.Get(db2)
			if err != nil {
				return err
			}
			if v == val {
				return errors.New("session variable got updated in another connection")
			}
			return nil
		})
		if err != nil {
			return err
		}
	}

	// test global scope
	if orig != val {
		db3, err := sql.Open("mysql", t.DSN)
		if err != nil {
			return errors.Annotate(err, "open database")
		}
		defer db3.Close()

		err = t.Var.WithValueInScope(db1, val, GLOBAL_ANNOTATION, func(db *sql.DB) error {
			vs, err := t.Var.Get(db3)
			if err != nil {
				return err
			}
			if vs == val {
				return errors.New("session variable got updated in another connection")
			}
			vg, err := t.Var.GetGlobal(db3)
			if err != nil {
				return err
			}
			if vg != val {
				return errors.New("global variable unchanged")
			}

			return nil
		})

		if err != nil {
			return err
		}
	}

	return nil
}

var (
	cases = make(map[string]TestCase)
)

func RegisterTestCase(n string, c TestCase) {
	if _, ok := cases[n]; ok {
		client.Warnf("[sysvars] overwrite existing case: name=%s var")
	}
	cases[n] = c
}

type Step struct {
	SQL        string
	Annotation string
}

func ExecAll(db *sql.DB, steps []Step) error {
	for _, s := range steps {
		_, err := db.Exec(s.SQL)
		if err != nil {
			if len(s.Annotation) > 0 {
				return errors.Annotate(err, s.Annotation)
			} else {
				return errors.Trace(err)
			}
		}
	}
	return nil
}

func main() {
	var opts struct {
		dsn   string
		cases string
		sets  string
	}
	flag.StringVar(&opts.dsn, "d", "", "overwrite dsn")
	flag.StringVar(&opts.cases, "t", "none", "test cases to run (comma separated)")
	flag.StringVar(&opts.sets, "s", "", "set global vars (pipe separated)")
	flag.Parse()

	client.Infof("[sysvars] start with options: %+v", opts)

	cs := make(map[string]TestCase)
	switch opts.cases {
	case "none":
	case "all":
		cs = cases
	default:
		for _, n := range strings.Split(opts.cases, ",") {
			n = strings.TrimSpace(n)
			if c, ok := cases[n]; ok {
				cs[n] = c
			} else {
				log.Warn("[sysvars] unknown case: " + n)
			}
		}
	}

	ctx, cancel := context.WithCancel(context.Background())

	sc := make(chan os.Signal, 1)
	signal.Notify(sc, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
	go func() {
		sig := <-sc
		client.Infof("[sysvars] terminate due to signal: %s", sig)
		cancel()
		os.Exit(1)
	}()

	if len(opts.dsn) == 0 {
		cli := client.New(ctx, client.ModeOnline)
		opts.dsn = fmt.Sprintf("root:@tcp(%s)/", cli.TiDBService())
		client.Infof("[sysvars] use dsn: %s", opts.dsn)
	}

	for n, c := range cs {
		client.Infof("[sysvars][%s] begin", n)
		err := c.Init(opts.dsn)
		if err != nil {
			log.Fatalf("[sysvars][%s] init error: %s", n, err.Error())
		}
		err = c.Run()
		if err != nil {
			log.Fatalf("[sysvars][%s] test error: %s", n, err.Error())
		}
		client.Infof("[sysvars][%s] end", n)
	}

	if len(cs) > 0 {
		client.Info("[sysvars] all tests passed")
	}

	var db *sql.DB
	var err error
	for _, s := range strings.Split(opts.sets, "|") {
		s = strings.TrimSpace(s)
		if len(s) == 0 {
			continue
		}
		kv := strings.SplitN(s, "=", 2)
		if len(kv) != 2 {
			log.Warn("[sysvars] invalid var assignment: ", s)
			continue
		}
		name, value := strings.TrimSpace(kv[0]), strings.TrimSpace(kv[1])
		tag := "[sysvars][" + name + "]\t"
		if db == nil {
			db, err = sql.Open("mysql", opts.dsn)
			if err != nil {
				log.Fatalf(tag+"failed to open database: dsn=%s err=%s",
					opts.dsn, err.Error())
			}
		}
		stmt := "SET " + string(GLOBAL_KEYWORD) + name + " = " + value
		v := SysVar{Name: name}

		orig, err := v.GetGlobal(db)
		if err != nil {
			client.Errorf(tag+"failed to get var: %v", err)
		}
		err = v.SetInScope(db, value, GLOBAL_KEYWORD)
		if err != nil {
			client.Errorf(tag+"failed to set var: %v", err)
			client.Errorf(tag+"exec: %-50s\tFAIL", stmt)
			continue
		} else {
			client.Infof(tag+"exec: %-50s\tOK", stmt)
		}
		cur, err := v.GetGlobal(db)
		if err != nil {
			client.Errorf(tag+"failed to get var: ", err)
		} else {
			client.Infof(tag+"value: %s -> %s", orig, cur)
		}

	}
}