// Copyright 2017 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"
	"database/sql"
	"errors"
	"flag"
	"fmt"
	"log"
	"math"

	_ "github.com/go-sql-driver/mysql"
	"github.com/pingcap/tidb-test/util"
)

var (
	host      = flag.String("h", "127.0.0.1", "host address")
	port      = flag.Int("P", 4001, "host port")
	user      = flag.String("u", "root", "set user of the database")
	password  = flag.String("p", "", "set password of the database ")
	dbname    = flag.String("D", "test", "set the default database name")
	tablename = flag.String("T", "analyze_test", "the default table name")
	rownum    = flag.Int64("R", 10000, "the row num of table")
)

func openDB() (*sql.DB, error) {
	dbaddr := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=utf8", *user, *password, *host, *port, *dbname)
	db, err := util.OpenDBWithRetry("mysql", dbaddr)
	if err != nil {
		return nil, err
	}
	return db, nil
}

func createTable(db *sql.DB) error {
	sql := fmt.Sprintf("DROP TABLE IF EXISTS %s", *tablename)
	if _, err := db.Exec(sql); err != nil {
		return err
	}
	sql = fmt.Sprintf(`CREATE TABLE %s (a BIGINT(64), b BIGINT(64),
			c BIGINT(64), PRIMARY KEY (a), KEY idx (b))`, *tablename)
	if _, err := db.Exec(sql); err != nil {
		return err
	}
	return nil
}

func insertTable(db *sql.DB) error {
	insertSQL := fmt.Sprintf("INSERT INTO %s (a, b, c) values", *tablename)
	var sqlStream bytes.Buffer
	var count int
	for id := int64(1); id <= *rownum; id++ {
		if sqlStream.Len() == 0 {
			sqlStream.WriteString(insertSQL)
		}

		sql := fmt.Sprintf("(%d, %d, %d)", id, id, id)
		sqlStream.WriteString(sql)

		count++
		if count >= 100 {
			count = 0
			if _, err := db.Exec(string(sqlStream.Bytes())); err != nil {
				return err
			}
			sqlStream.Reset()
		} else if id != *rownum {
			sqlStream.WriteString(",")
		}
	}

	// last batch
	if count > 0 {
		if _, err := db.Exec(string(sqlStream.Bytes())); err != nil {
			return err
		}
	}
	return nil
}

func analyzeTable(db *sql.DB) error {
	sql := fmt.Sprintf("ANALYZE TABLE %s", *tablename)
	if _, err := db.Exec(sql); err != nil {
		return err
	}
	return nil
}

func checkMeta(db *sql.DB) error {
	sql := fmt.Sprintf("SHOW STATS_META where Table_name = '%s'", *tablename)
	rows, err := db.Query(sql)
	if err != nil {
		return err
	}
	cols, err := rows.Columns()
	if err != nil {
		return err
	}
	var rowCount int64
	var non interface{}
	for rows.Next() {
		if len(cols) == 5 {
			// do not care the first 4 values.
			err = rows.Scan(&non, &non, &non, &non, &rowCount)
		} else {
			// extra one for partition name
			err = rows.Scan(&non, &non, &non, &non, &non, &rowCount)
		}
		if err != nil {
			return err
		}
		if rowCount != *rownum {
			return errors.New(fmt.Sprintf("expect rows count %d, got %d", *rownum, rowCount))
		}
	}
	return nil
}

func checkBuckets(db *sql.DB, colName string, is_index uint8) error {
	sql := fmt.Sprintf("SHOW STATS_BUCKETS where Table_name = '%s' and Column_name = '%s' and Is_index = %d", *tablename, colName, is_index)
	rows, err := db.Query(sql)
	if err != nil {
		return err
	}
	cols, err := rows.Columns()
	if err != nil {
		return err
	}
	var (
		non      interface{}
		count    int64
		preCount int64
		repeats  int64
		lower    int64
		upper    int64
		preUpper int64
	)
	for rows.Next() {
		if len(cols) == 9 {
			err = rows.Scan(&non, &non, &non, &non, &non, &count, &repeats, &lower, &upper)
		} else {
			// extra one for partition name
			err = rows.Scan(&non, &non, &non, &non, &non, &non, &count, &repeats, &lower, &upper, &non)
		}
		if err != nil {
			return err
		}
		if count < preCount {
			return errors.New(fmt.Sprintf("count %d should be larger than previous count %d", count, preCount))
		}
		if repeats != 1 {
			return errors.New(fmt.Sprintf("expect repeats %d, got %d", 1, repeats))
		}
		if lower < preUpper {
			return errors.New(fmt.Sprintf("lower bound %s should be larger than previous upper bound %s", lower, preUpper))
		}
		if lower > upper {
			return errors.New(fmt.Sprintf("lower bound %s should be not larger than upper bound %s", lower, upper))
		}
		preCount = count
		preUpper = upper
	}
	return nil
}

func checkHistograms(db *sql.DB) error {
	sql := fmt.Sprintf("SHOW STATS_HISTOGRAMS where Table_name = '%s'", *tablename)
	rows, err := db.Query(sql)
	if err != nil {
		return err
	}
	if err != nil {
		return err
	}
	var (
		non         interface{}
		column_name string
		is_index    uint8
		ndv         int64
		null        int64
		avg_size    float64
	)
	for rows.Next() {
		err = rows.Scan(&non, &non, &non, &column_name, &is_index, &non, &ndv, &null, &avg_size,
			&non, &non, &non, &non, &non, &non)
		if err != nil {
			return err
		}
		if is_index == 1 && ndv != *rownum {
			return errors.New(fmt.Sprintf("expect distinct value %d, got %d", *rownum, ndv))
		}
		if is_index == 0 && math.Abs(float64(ndv-*rownum))/float64(*rownum) > 0.01 {
			return errors.New(fmt.Sprintf("actual distinct value %d, got %d", *rownum, ndv))
		}
		if null != 0 {
			return errors.New(fmt.Sprintf("expect null value %d, got %d", 0, null))
		}
		if is_index == 1 && avg_size != 0 {
			return errors.New(fmt.Sprintf("expect avg col size 0, got %f", avg_size))
		}
		if is_index == 0 && column_name == "a" && avg_size != 8 {
			return errors.New(fmt.Sprintf("expect avg col size 8, got %f", avg_size))
		}
		if err := checkBuckets(db, column_name, is_index); err != nil {
			return err
		}
	}
	return nil
}

func checkAnalyze(db *sql.DB) error {
	if err := checkMeta(db); err != nil {
		return err
	}
	if err := checkHistograms(db); err != nil {
		return err
	}
	return nil
}

func checkVirtualColumnAnalyze(db *sql.DB) error {
	dropTable := fmt.Sprintf("DROP TABLE IF EXISTS %s", *tablename)
	if _, err := db.Exec(dropTable); err != nil {
		return err
	}
	createTable := fmt.Sprintf(`
CREATE TABLE %s (
	id INT AUTO_INCREMENT PRIMARY KEY,
	first_name VARCHAR(50) NOT NULL,
	last_name VARCHAR(50) NOT NULL,
	fullname varchar(101) GENERATED ALWAYS AS (CONCAT(first_name,' ',last_name)) VIRTUAL NOT NULL,
	fullname2 varchar(101) GENERATED ALWAYS AS (CONCAT(first_name,' ',last_name)) STORED NOT NULL,
	email VARCHAR(100) NOT NULL
)
`, *tablename)

	if _, err := db.Exec(createTable); err != nil {
		return err
	}
	insertData := fmt.Sprintf("INSERT INTO %s(first_name,last_name, email) VALUES('john','doe','test@pingcap.com')", *tablename)
	if _, err := db.Exec(insertData); err != nil {
		return err
	}
	return analyzeTable(db)
}

func main() {
	flag.Parse()
	db, err := openDB()
	if err != nil {
		log.Fatal(err)
	}
	if err := createTable(db); err != nil {
		log.Fatal(err)
	}
	if err := insertTable(db); err != nil {
		log.Fatal(err)
	}
	if err := analyzeTable(db); err != nil {
		log.Fatal(err)
	}
	if err := checkAnalyze(db); err != nil {
		log.Fatal(err)
	}
	if err := checkVirtualColumnAnalyze(db); err != nil {
		log.Fatal(err)
	}
}