// Copyright 2015 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 (
	"bytes"
	"context"
	"database/sql"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"strings"
	"sync"
	"time"

	"github.com/docopt/docopt-go"
	"github.com/juju/errors"
	zaplog "github.com/pingcap/log"
	"github.com/pingcap/parser/ast"
	"github.com/pingcap/schrodinger-test/util"
	client "github.com/pingcap/schrodinger/pkg/test_client"
	"github.com/pingcap/tidb/session"
	"github.com/pingcap/tidb/sessionctx"
	"github.com/pingcap/tidb/util/logutil"
	"github.com/pingcap/tidb/util/mock"
	log "github.com/sirupsen/logrus"
)

var usage = `
	Usage: mysql_test [options] [<tests>...]

options:
	--record                       Record the test output in the result file
	--log-level=<loglevel>         set log level: info, warn, error, debug [default: error]
`

var (
	wg     sync.WaitGroup
	record bool
)

type query struct {
	firstWord string
	Query     string
	Line      int
	tp        int
}

type tester struct {
	mdb  *sql.DB
	name string

	tx *sql.Tx

	buf bytes.Buffer

	// enable query log will output origin statement into result file too
	// use --disable_query_log or --enable_query_log to control it
	enableQueryLog bool

	singleQuery bool

	queryVertical bool

	// Disable or enable warnings. This setting is enabled by default.
	// With this setting enabled, mysqltest uses SHOW WARNINGS to display
	// any warnings produced by SQL statements.
	enableWarning bool

	// check expected error, use --error before the statement
	// see http://dev.mysql.com/doc/mysqltest/2.0/en/writing-tests-expecting-errors.html
	expectedErrs []string

	// only for test, not record, every time we execute a statement, we should read the result
	// data to check correction.
	resultFD *os.File
	// ctx is used for Compile sql statement
	ctx sessionctx.Context
}

func newTester(name string) *tester {
	t := new(tester)

	t.name = name
	t.enableQueryLog = true
	// disable warning by default since our a lot of test cases
	// are ported wihtout explictly "disablewarning"
	t.enableWarning = false
	t.ctx = mock.NewContext()

	return t
}
func (t *tester) preProcess() {
	dbName := "test"
	// need to adjust
	mdb, err := OpenDBWithRetry("mysql", "root@tcp("+dbAddr+")/"+dbName+"?time_zone=%27Asia%2FShanghai%27")
	if err != nil {
		log.Fatalf("Open db err %v", err)
	}
	// log.Info("Create new db", mdb)

	if _, err = mdb.Exec(fmt.Sprintf("create database `%s`", t.name)); err != nil {
		log.Fatalf("Executing create db %s err[%v]", t.name, err)
	}

	if _, err = mdb.Exec(fmt.Sprintf("use `%s`", t.name)); err != nil {
		log.Fatalf("Executing Use test err[%v]", err)
	}
	if _, err = mdb.Exec("SET @@tidb_max_chunk_size=2"); err != nil {
		log.Fatalf("Executing \"SET @@tidb_max_chunk_size=2\" err[%v]", err)
	}
	if _, err = mdb.Exec("SET @@tidb_hash_join_concurrency=1"); err != nil {
		log.Fatalf("Executing \"SET @@tidb_hash_join_concurrency=2\" err[%v]", err)
	}
	t.mdb = mdb
}

func (t *tester) postProcess() {
	t.mdb.Exec(fmt.Sprintf("drop database `%s`", t.name))
	t.mdb.Close()
}

func (t *tester) Run() error {
	t.preProcess()
	queries, err := t.loadQueries()
	if err != nil {
		return errors.Trace(err)
	}

	if err = t.openResult(); err != nil {
		return errors.Trace(err)
	}

	var s string
	defer func() {
		if t.tx != nil {
			log.Errorf("transaction is not committed correctly, rollback")
			t.rollback()
		}

		if t.resultFD != nil {
			t.resultFD.Close()
		}
	}()

	testCnt := 0
	startTime := time.Now()
	for _, q := range queries {
		s = q.Query
		switch q.tp {
		case Q_ENABLE_QUERY_LOG:
			t.enableQueryLog = true
		case Q_DISABLE_QUERY_LOG:
			t.enableQueryLog = false
		case Q_DISABLE_WARNINGS:
			t.enableWarning = false
		case Q_ENABLE_WARNINGS:
			t.enableWarning = true
		case Q_SINGLE_QUERY:
			t.singleQuery = true
		case Q_ERROR:
			t.expectedErrs = strings.Split(strings.TrimSpace(s), ",")
		case Q_ECHO:
			t.buf.WriteString(s)
			t.buf.WriteString("\n")
		case Q_QUERY_VERTICAL:
			t.queryVertical = true
			fallthrough
		case Q_QUERY:
			if err = t.execute(q); err != nil {
				return errors.Annotate(err, fmt.Sprintf("sql:%v", q.Query))
			}

			testCnt++

			// disable query vertical afterward.
			t.queryVertical = false
		}
	}

	fmt.Printf("%s: ok! %d test cases passed, take time %v s\n", t.testFileName(), testCnt, time.Since(startTime).Seconds())

	t.postProcess()

	return t.flushResult()
}

func (t *tester) loadQueries() ([]query, error) {
	data, err := ioutil.ReadFile(t.testFileName())
	if err != nil {
		return nil, err
	}

	seps := bytes.Split(data, []byte("\n"))
	queries := make([]string, 0, len(seps))
	newStmt := true
	for _, v := range seps {
		v := bytes.TrimSpace(v)
		s := string(v)
		// we will skip # comment here
		if strings.HasPrefix(s, "#") {
			newStmt = true
			continue
		} else if strings.HasPrefix(s, "--") {
			queries = append(queries, s)
			newStmt = true
			continue
		} else if len(s) == 0 {
			continue
		}

		if newStmt {
			queries = append(queries, s)
		} else {
			lastQuery := queries[len(queries)-1]
			lastQuery = fmt.Sprintf("%s\n%s", lastQuery, s)
			queries[len(queries)-1] = lastQuery
		}

		// if the line has a ; in the end, we will treat new line as the new statement.
		newStmt = strings.HasSuffix(s, ";")
	}

	return ParseQueries(queries...)
}

// parserErrorHandle handle mysql_test syntax `--error ER_PARSE_ERROR`, to allow following query
// return parser error.
func (t *tester) parserErrorHandle(query query, err error) error {
	offset := t.buf.Len()
	for _, expectedErr := range t.expectedErrs {
		if expectedErr == "ER_PARSE_ERROR" {
			if t.enableQueryLog {
				t.buf.WriteString(query.Query)
				t.buf.WriteString("\n")
			}

			t.buf.WriteString(fmt.Sprintf("%s\n", err))
			err = nil
			break
		}
	}

	if err != nil {
		return errors.Trace(err)
	}

	// clear expected errors after we execute the first query
	t.expectedErrs = nil
	t.singleQuery = false

	if !record {
		// check test result now
		gotBuf := t.buf.Bytes()[offset:]
		buf := make([]byte, t.buf.Len()-offset)
		if _, err = t.resultFD.ReadAt(buf, int64(offset)); err != nil {
			return errors.Trace(errors.Errorf("run \"%v\" at line %d err, we got \n%s\nbut read result err %s", query.Query, query.Line, gotBuf, err))
		}

		if !bytes.Equal(gotBuf, buf) {
			return errors.Trace(errors.Errorf("run \"%v\" around line %d err, we need(%v):\n%s\nbut got(%v):\n%s\n", query.Query, query.Line, len(buf), buf, len(gotBuf), gotBuf))
		}
	}

	return errors.Trace(err)
}

func (t *tester) execute(query query) error {
	if len(query.Query) == 0 {
		return nil
	}

	list, err := session.Parse(t.ctx, query.Query)
	if err != nil {
		return t.parserErrorHandle(query, err)
	}

	for _, st := range list {
		var qText string
		if t.singleQuery {
			qText = query.Query
		} else {
			qText = st.Text()
		}
		offset := t.buf.Len()
		if t.enableQueryLog {
			t.buf.WriteString(qText)
			t.buf.WriteString("\n")
		}

		switch st.(type) {
		case *ast.BeginStmt:
			t.tx, err = t.mdb.Begin()
			if err != nil {
				t.rollback()
				break
			}
		case *ast.CommitStmt:
			err = t.commit()
			if err != nil {
				t.rollback()
				break
			}
		case *ast.RollbackStmt:
			err = t.rollback()
			if err != nil {
				break
			}
		default:
			if t.tx != nil {
				_, err = filterWarning(t.executeStmt(qText), t.enableWarning)
				if err != nil {
					break
				}
			} else {
				// if begin or following commit fails, we don't think
				// this error is the expected one.
				if t.tx, err = t.mdb.Begin(); err != nil {
					t.rollback()
					break
				}

				var warn bool
				warn, err = filterWarning(t.executeStmt(qText), t.enableWarning)
				if err != nil && !warn {
					t.rollback()
					break
				} else {
					warn, commitErr := filterWarning(t.commit(), t.enableWarning)
					if err == nil && commitErr != nil {
						err = commitErr
					}
					if commitErr != nil && !warn {
						t.rollback()
						break
					}
				}
			}
		}

		if err != nil && len(t.expectedErrs) > 0 {
			// TODO: check whether this err is expected.
			// but now we think it is.

			// output expected err
			t.buf.WriteString(fmt.Sprintf("%s\n", err))
			err = nil
		}
		// clear expected errors after we execute the first query
		t.expectedErrs = nil
		t.singleQuery = false

		if err != nil {
			return errors.Trace(errors.Errorf("run \"%v\" at line %d err %v", st.Text(), query.Line, err))
		}

		if !record {
			// check test result now
			gotBuf := t.buf.Bytes()[offset:]

			buf := make([]byte, t.buf.Len()-offset)
			if _, err = t.resultFD.ReadAt(buf, int64(offset)); err != nil {
				return errors.Trace(errors.Errorf("run \"%v\" at line %d err, we got \n%s\nbut read result err %s", st.Text(), query.Line, gotBuf, err))
			}

			if !bytes.Equal(gotBuf, buf) {
				return errors.Trace(errors.Errorf("failed to run query \n\"%v\" \n around line %d, \nwe need(%v):\n%s\nbut got(%v):\n%s\n", query.Query, query.Line, len(buf), buf, len(gotBuf), gotBuf))
			}
		}
	}

	return errors.Trace(err)
}

// MySQLWarning is an error type which represents a single MySQL warning.
// Warnings are returned in groups only. See MySQLWarnings
type MySQLWarning struct {
	Level   string
	Code    string
	Message string
}

// MySQLWarnings is an error type which represents a group of one or more MySQL
// warnings
type MySQLWarnings []MySQLWarning

// LineBreak is linebreak
const LineBreak = "\n"

func (mws MySQLWarnings) Error() string {
	var msg string
	for i, warning := range mws {
		if i > 0 {
			msg += LineBreak
		}
		msg += fmt.Sprintf(
			"%s %s: %s",
			warning.Level,
			warning.Code,
			warning.Message,
		)
	}
	return msg
}

func filterWarning(err error, enableWarning bool) (isWarn bool, outErr error) {
	if err == nil {
		return false, err
	}
	causeErr := errors.Cause(err)
	_, isWarn = causeErr.(MySQLWarnings)
	// enableWarning is enabled by default
	if !enableWarning {
		if isWarn {
			return isWarn, nil
		}
	}
	return isWarn, err
}

func (t *tester) commit() error {
	err := t.tx.Commit()
	if err != nil {
		return err
	}
	t.tx = nil
	return nil
}

func (t *tester) rollback() error {
	if t.tx == nil {
		return nil
	}
	err := t.tx.Rollback()
	t.tx = nil
	return err
}

func (t *tester) writeQueryResult(rows *sql.Rows) error {
	cols, err := rows.Columns()
	if err != nil {
		return errors.Trace(err)
	}

	for i, c := range cols {
		t.buf.WriteString(c)
		if i != len(cols)-1 {
			t.buf.WriteString("\t")
		}
	}
	t.buf.WriteString("\n")

	values := make([][]byte, len(cols))
	scanArgs := make([]interface{}, len(values))
	for i := range values {
		scanArgs[i] = &values[i]
	}

	for rows.Next() {
		err := rows.Scan(scanArgs...)
		if err != nil {
			return errors.Trace(err)
		}

		var value string
		for i, col := range values {
			// Here we can check if the value is nil (NULL value)
			if col == nil {
				value = "NULL"
			} else {
				value = string(col)
			}
			t.buf.WriteString(value)
			if i < len(values)-1 {
				t.buf.WriteString("\t")
			}
		}
		t.buf.WriteString("\n")
	}
	err = rows.Err()
	if err != nil {
		return errors.Trace(err)
	}

	return nil
}

func (t *tester) writeQueryResultVertically(rows *sql.Rows) error {
	cols, err := rows.Columns()
	if err != nil {
		return errors.Trace(err)
	}
	// each buffer buffers each column's result
	// say we have a table with c1 and c2 columns and data is 1, 2; 3, 4
	buffers := make([][]byte, len(cols))
	for i := range buffers {
		buffers[i] = make([]byte, 0, 1024)
	}

	for i, c := range cols {
		buffers[i] = append(buffers[i], c...)
		buffers[i] = append(buffers[i], "\t"...)
	}

	values := make([][]byte, len(cols))
	scanArgs := make([]interface{}, len(values))
	for i := range values {
		scanArgs[i] = &values[i]
	}

	rowCnt := 0
	for rows.Next() {
		err := rows.Scan(scanArgs...)
		if err != nil {
			return errors.Trace(err)
		}

		// fetch first row data: 1, 3 and append them to their corresponding
		// buffer. After this for loop, buffer should look like the following:
		//       buf1: c1 1 2
		//       buf2: c2 3 4
		for i, col := range values {
			// Here we can check if the value is nil (NULL value)
			if col == nil {
				buffers[i] = append(buffers[i], "NULL"...)
			} else {
				buffers[i] = append(buffers[i], col...)
			}
			buffers[i] = append(buffers[i], "\t"...)
		}
		rowCnt++
	}

	// the last step is to write temporary buffers into tester's buffer.
	for _, buf := range buffers {
		buf[len(buf)-1] = '\n'
		t.buf.Write(buf)
	}

	err = rows.Err()
	if err != nil {
		return errors.Trace(err)
	}

	return nil
}

func (t *tester) executeStmt(query string) error {
	if util.IsQuery(query) {
		rows, err := t.tx.Query(query)
		if err != nil {
			return errors.Trace(err)
		}

		if t.queryVertical {
			return t.writeQueryResultVertically(rows)
		} else {
			return t.writeQueryResult(rows)
		}
	} else {
		// TODO: rows affected and last insert id
		_, err := t.tx.Exec(query)
		if err != nil {
			return errors.Trace(err)
		}
	}
	return nil
}

func (t *tester) openResult() error {
	if record {
		return nil
	}

	var err error
	t.resultFD, err = os.Open(t.resultFileName())
	return err
}

func (t *tester) flushResult() error {
	if !record {
		return nil
	}
	return ioutil.WriteFile(t.resultFileName(), t.buf.Bytes(), 0644)
}

func (t *tester) testFileName() string {
	// test and result must be in current ./t the same as MySQL
	return fmt.Sprintf("./t/%s.test", t.name)
}

func (t *tester) resultFileName() string {
	// test and result must be in current ./r, the same as MySQL
	return fmt.Sprintf("./r/%s.result", t.name)
}

func loadAllTests() ([]string, error) {
	// tests must be in t folder
	files, err := ioutil.ReadDir("./t")
	if err != nil {
		return nil, err
	}

	tests := make([]string, 0, len(files))
	for _, f := range files {
		if f.IsDir() {
			continue
		}

		// the test file must have a suffix .test
		name := f.Name()
		if strings.HasSuffix(name, ".test") {
			name = strings.TrimSuffix(name, ".test")

			// if we use record and the result file exists, skip generating
			if record && resultExists(name) {
				continue
			}

			tests = append(tests, name)
		}
	}

	return tests, nil
}

func resultExists(name string) bool {
	resultFile := fmt.Sprintf("./r/%s.result", name)

	if _, err := os.Stat(resultFile); err != nil {
		if os.IsNotExist(err) {
			return false
		}
	}
	return true
}

// convertTestsToTestTasks convert all test cases into several testBatches.
// If we have 11 cases and batchSize is 5, then we will have 4 testBatches.
func convertTestsToTestTasks(tests []string) []testBatch {
	batchSize := 30
	total := (len(tests) / batchSize) + 1
	// the extra 1 is for sub_query_more test
	tTasks := make([]testBatch, total+1)
	testIdx := 0
	for i := 0; i < total; i++ {
		tTasks[i] = make(testBatch, 0, batchSize)
		for j := 0; j <= batchSize && testIdx < len(tests); j++ {
			// skip sub_query_more test, since it consumes the most time
			// we better use a separate goroutine to run it
			if tests[testIdx] != "sub_query_more" && tests[testIdx] != "show" && tests[testIdx] != "infoschema" {
				tTasks[i] = append(tTasks[i], tests[testIdx])
			}
			testIdx++
		}
	}

	tTasks[total] = testBatch{"sub_query_more"}

	return tTasks
}

var done = make(chan bool)
var msgs = make(chan testTask)

type testTask struct {
	err  error
	test string
}

type testBatch []string

func (t testBatch) Run() {
	defer wg.Done()
	for _, test := range t {
		tr := newTester(test)
		msgs <- testTask{
			test: test,
			err:  tr.Run(),
		}
	}
}

func (t testBatch) String() string {
	return strings.Join([]string(t), ", ")
}

func executeTests(tasks []testBatch) {
	// show and infoschema have to be executed first, since the following
	// tests will create database using their own name.
	show := newTester("show")
	msgs <- testTask{
		test: "show",
		err:  show.Run(),
	}

	infoschema := newTester("infoschema")
	msgs <- testTask{
		test: "infoschema",
		err:  infoschema.Run(),
	}

	wg.Add(len(tasks))
	for _, t := range tasks {
		go t.Run()
	}

	wg.Wait()
}

func consumeError() {
	for {
		if t, more := <-msgs; more {
			if t.err != nil {
				log.Fatalf("run test [%s] err: %v", t.test, t.err)
			} else {
				log.Infof("run test [%s] ok", t.test)
			}

		} else {
			done <- true
			return
		}
	}

}

var (
	dbAddr = "localhost:4001"
)

func downloadTestFile() error {
	testFileURL := "http://pingcap-dev.hk.ufileos.com/schrodinger/mysql_test_files.tar.gz"
	testFileName := path.Base(testFileURL)
	if err := exec.Command("curl", testFileURL, "-o", testFileName).Run(); err != nil {
		log.Error("download test file error")
		return err
	}
	log.Info("successfully download test file")
	if err := exec.Command("tar", "-xzf", testFileName).Run(); err != nil {
		log.Error("extract test file error")
		return err
	}
	log.Info("successfully extract test file")
	return nil
}

func init() {
	if err := downloadTestFile(); err != nil {
		log.Fatal("downloadTestFile Error: ", err)
	}
	log.Info("start get schrodinger config")
	cli := client.New(context.Background(), client.ModeOnline)
	if tidbAddr := cli.TiDBService(); tidbAddr != "" {
		dbAddr = tidbAddr
	}
	log.Infof("schrodinger dbAddrs are: %v", dbAddr)
}

func main() {
	args, err := docopt.Parse(usage, nil, true, "mysql_test v0.1", true)
	if err != nil {
		log.Fatal(err)
	}
	var logLevel string
	if args["--log-level"] != nil {
		logLevel = args["--log-level"].(string)
	}
	logutil.InitLogger(&logutil.LogConfig{
		Config: zaplog.Config{
			Level: logLevel,
		},
	})

	var tests []string
	if v := args["<tests>"]; v != nil {
		tests = v.([]string)
	}

	if v := args["--record"]; v != nil {
		record = v.(bool)
	}

	// we will run all tests if no tests assigned
	if len(tests) == 0 {
		if tests, err = loadAllTests(); err != nil {
			log.Fatalf("load all tests err %v", err)
		}
	}

	if !record {
		log.Infof("running tests: %v", tests)
	} else {
		log.Infof("recording tests: %v", tests)
	}

	go consumeError()
	executeTests(convertTestsToTestTasks(tests))
	close(msgs)
	<-done
	println("\nGreat, All tests passed")
}