package privileges_update

import (
	"context"
	"database/sql"
	"fmt"
	"strings"
	"time"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	_ "github.com/pingcap/endless/pkg/client/tidb"
	"github.com/pingcap/endless/pkg/matcher"
	"github.com/pingcap/endless/testcase/oltp/oltppkg"
	. "github.com/pingcap/endless/testcase/upgrade/utils"
	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/resource"
	"go.uber.org/zap"
)

// RESTRICTED_VARIABLES_ADMIN
var HiddenVariables = []string{
	"tidb_allow_remove_auto_inc",
	"ddl_slow_threshold",
	"tidb_check_mb4_value_in_utf8",
	"tidb_config",
	"tidb_enable_slow_log",
	"tidb_enable_telemetry",
	"tidb_expensive_query_time_threshold",
	"tidb_force_priority",
	"tidb_general_log",
	"tidb_metric_query_range_duration",
	"tidb_metric_query_step",
	"tidb_opt_write_row_id",
	"tidb_pprof_sql_cpu",
	"tidb_record_plan_in_slow_log",
	"tidb_row_format_version",
	"tidb_slow_query_file",
	"tidb_slow_query_file",
	"tidb_slow_log_threshold",
	"tidb_enable_collect_execution_info",
	"tidb_memory_usage_alarm_ratio",
	"tidb_redact_log",
}

// always "localhost" if enable sem
var ResetVariables = []string{
	"hostname",
}

// RESTRICTED_STATUS_ADMIN
var StatusHiddenVariables = []string{
	"tidb_gc_leader_desc",
}

//RESTRICTED_TABLES_ADMIN
var HiddenTables = []string{
	"information_schema.cluster_config",
	"information_schema.cluster_hardware",
	"information_schema.cluster_load",
	"information_schema.cluster_log",
	"information_schema.cluster_systeminfo",
	"information_schema.inspection_result",
	"information_schema.inspection_rules",
	"information_schema.inspection_summary",
	"information_schema.metrics_summary",
	"information_schema.metrics_summary_by_label",
	"information_schema.metrics_tables",
	"information_schema.tidb_hot_regions",
	"performance_schema.pd_profile_allocs",
	"performance_schema.pd_profile_block",
	"performance_schema.pd_profile_cpu",
	"performance_schema.pd_profile_goroutines",
	"performance_schema.pd_profile_memory",
	"performance_schema.pd_profile_mutex",
	"performance_schema.tidb_profile_allocs",
	"performance_schema.tidb_profile_block",
	"performance_schema.tidb_profile_cpu",
	"performance_schema.tidb_profile_goroutines",
	"performance_schema.tidb_profile_memory",
	"performance_schema.tidb_profile_mutex",
	"performance_schema.tikv_profile_cpu",
	"mysql.expr_pushdown_blacklist",
	"mysql.gc_delete_range",
	"mysql.gc_delete_range_done",
	"mysql.opt_rule_blacklist",
	"mysql.tidb",
	"mysql.global_variables",
	// all metrics_schema tables will be hidden
	"metrics_schema.node_memory_active",
}

var TableInstanceToID = []string{
	// "information_schema.cluster_deadlocks", // >=5.1.0
	"information_schema.cluster_info",
	"information_schema.cluster_processlist",
	"information_schema.cluster_slow_query",
	"information_schema.cluster_statements_summary",
	// "information_schema.cluster_statements_summary_evicted", // >=5.2.0
	"information_schema.cluster_statements_summary_history",
	"information_schema.cluster_tidb_trx",
}

var _ = Describe("[test_suite_privilege] upgrade cluster and enable sem", func() {
	It("should be success", func() {
		targetVersion := suiteTestCtx.Param("target_version", "v6.2.0-pre")
		tc := suiteTestCtx.Resource("tc").(resource.TiDBCluster)
		var config string = "[security]\nenable-sem=true"
		log.Info("start to upgrade the cluster and set config", zap.String("target_version", targetVersion), zap.String("config", config))
		UpdateClusterAndConfig(suiteTestCtx, tc, targetVersion, config)
	})
})

var _ = Describe("TestPrivilegesUpdate", func() {
	var (
		db     *sql.DB
		tc     resource.TiDBCluster
		ctx    context.Context
		cancel context.CancelFunc
	)

	BeforeEach(func() {
		tc = suiteTestCtx.Resource("tc").(resource.TiDBCluster)
		ctx, cancel = context.WithTimeout(suiteTestCtx, time.Minute*5)
		db, _ = tc.GetDB("test")
	})

	AfterEach(func() {
		db.Close()
		cancel()
	})

	Describe("Precheck super privileges with sem disable", func() {
		It("should the variables value correct", func() {
			var value int
			matcher.Try(db.QueryRow("select @@tidb_enable_enhanced_security;").Scan(&value))
			Expect(value).To(Equal(0))
		})

		It("should query all hidden variables and tables", func() {
			CheckSuperPrivilegesWithoutSem(ctx, tc, db)
		})

		It("can grants restricted privileges", func() {
			matcher.Try(db.Exec("drop user if exists u1, u2, u3"))
			matcher.Try(db.Exec("create user u1, u2, u3"))
			matcher.Try(db.Exec("grant all on *.* to u1,u2,u3"))
			matcher.Try(db.Exec("grant RESTRICTED_VARIABLES_ADMIN on *.* to u1;"))
			matcher.Try(db.Exec("grant RESTRICTED_STATUS_ADMIN on *.* to u2;"))
			matcher.Try(db.Exec("grant RESTRICTED_TABLES_ADMIN on *.* to u3;"))
			time.Sleep(3 * time.Second)

			db1, _ := tc.GetDB("test")
			Expect(CheckPrivilege(db1, "u1", "RESTRICTED_VARIABLES_ADMIN")).To(BeTrue())
			Expect(CheckPrivilege(db1, "u2", "RESTRICTED_STATUS_ADMIN")).To(BeTrue())
			Expect(CheckPrivilege(db1, "u3", "RESTRICTED_TABLES_ADMIN")).To(BeTrue())
		})
	})

	Describe("Postcheck super privileges with sem disable", func() {
		It("should the variables value correct", func() {
			var value int
			matcher.Try(db.QueryRow("select @@tidb_enable_enhanced_security;").Scan(&value))
			Expect(value).To(Equal(1))
		})
		It("should hide variables for super privilege", func() {
			for _, v := range HiddenVariables {
				CheckVariable(ctx, db, v, true, false)
			}
			_, err := tc.GetDB("metrics_schema")
			Expect(err).To(HaveOccurred())
		})
		It("should reset variables for super privilege", func() {
			var hostname string
			err := db.QueryRow("select @@hostname").Scan(&hostname)
			Expect(err).NotTo(HaveOccurred())
			Expect(hostname).To(Equal("localhost"))
		})
		It("should hide status variables for super privilege", func() {
			CheckStatusAdmin(ctx, db, true, false)
		})
		It("should hide tables for super privilege", func() {
			for _, v := range HiddenTables {
				CheckTable(ctx, db, v, true, false)
			}
		})
		It("should hide table instances for super privilege", func() {
			for _, v := range TableInstanceToID {
				CheckTableInstance(ctx, db, v, true, false)
			}
		})
		It("should set some info to null for privilege", func() {
			CheckNullResult(ctx, db, true, false)
		})
	})

	Describe("Postcheck restricted privileges with sem enable", func() {

		It("user with RESTRICTED_TABLES_ADMIN privileges can access all tables", func() {
			db1, err := GetDBWithUser(tc, "test", "u3")
			for _, table := range HiddenTables {
				CheckTable(ctx, db1, table, true, true)
			}
			_, err = GetDBWithUser(tc, "metrics_schema", "u3")
			Expect(err).NotTo(HaveOccurred())

			for _, table := range TableInstanceToID {
				CheckTableInstance(ctx, db1, table, true, true)
			}

			CheckNullResult(ctx, db1, true, true)
		})

		It("user with RESTRICTED_VARIABLES_ADMIN privileges can access all variables", func() {
			db1, _ := GetDBWithUser(tc, "test", "u1")
			for _, v := range HiddenVariables {
				CheckVariable(ctx, db1, v, true, true)
			}
		})
		It("user with RESTRICTED_STATUS_ADMIN privileges can access all status variables", func() {
			db1, _ := GetDBWithUser(tc, "test", "u2")
			CheckStatusAdmin(ctx, db1, true, true)
		})

	})
})

func CheckSuperPrivilegesWithoutSem(ctx context.Context, tc resource.TiDBCluster, db *sql.DB) {
	for _, variable := range HiddenVariables {
		CheckVariable(ctx, db, variable, false, false)
	}
	for _, tables := range HiddenTables {
		CheckTable(ctx, db, tables, false, false)
	}

	for _, tables := range TableInstanceToID {
		CheckTableInstance(ctx, db, tables, false, false)
	}

	CheckStatusAdmin(ctx, db, false, false)
	CheckNullResult(ctx, db, false, false)

	// https://github.com/pingcap/tidb/issues/32918
	if r := oltppkg.CompareVersionWithCurrent(ctx, tc, "v6.0.0"); r < 0 {
		var hostname string
		err := db.QueryRow("select @@hostname").Scan(&hostname)
		Expect(err).NotTo(HaveOccurred())
		Expect(hostname).NotTo(Equal("localhost"))
	}

}

func CheckVariable(ctx context.Context, db *sql.DB, variable string, sem bool, variableAdmin bool) {
	var value string
	log.Info("check variable", zap.String("variable", variable))
	err := db.QueryRowContext(ctx, "select @@"+variable).Scan(&value)
	if sem && !variableAdmin {
		Expect(err).To(HaveOccurred())
		Expect(err.Error()).Should(ContainSubstring("Access denied"))
	} else {
		Expect(err).NotTo(HaveOccurred())
	}
}

func CheckTable(ctx context.Context, db *sql.DB, table string, sem bool, tableAdmin bool) {
	var (
		value int
		sql   string
	)
	if table == "information_schema.cluster_log" {
		sql = fmt.Sprintf("select count(*) from information_schema.cluster_log where time > '%s' and time < '%s' and message like '%%sem%%' limit 1;", time.Now().Add(-time.Hour).Format("2006-01-02 15:04:05"), time.Now().Format("2006-01-02 15:04:05"))
	} else {
		sql = "select count(*) from " + table + " limit 1;"
	}
	log.Info("check table", zap.String("sql", sql))
	err := db.QueryRowContext(ctx, sql).Scan(&value)
	if sem && !tableAdmin {
		Expect(err).To(HaveOccurred())
		Expect(err.Error()).To(ContainSubstring("SELECT command denied to user"))
	} else {
		Expect(err).NotTo(HaveOccurred())
	}
}

func CheckTableInstance(ctx context.Context, db *sql.DB, table string, sem bool, tableAdmin bool) {
	var value string
	log.Info("check instance", zap.String("table", table))
	rows, err := db.QueryContext(ctx, "select instance from "+table)
	Expect(err).NotTo(HaveOccurred())
	for rows.Next() {
		_ = rows.Scan(&value)
		if value == "" {
			continue
		}
		if sem && !tableAdmin {
			Expect(value).NotTo(ContainSubstring("."), "instance: %s", value)
		} else {
			Expect(value).To(ContainSubstring("."), "instance: %s", value)
		}
	}
}

func CheckNullResult(ctx context.Context, db *sql.DB, sem bool, tableAdmin bool) {
	var address, capacity, available, start_ts, uptime string
	var ip string

	rows1, err1 := db.QueryContext(ctx, "select address, capacity, available, start_ts, uptime from information_schema.tikv_store_status")
	rows2, err2 := db.QueryContext(ctx, "select ip from information_schema.tidb_servers_info")
	Expect(err1).NotTo(HaveOccurred())
	Expect(err2).NotTo(HaveOccurred())
	for rows1.Next() {
		_ = rows1.Scan(&address, &capacity, &available, &start_ts, &uptime)
		for _, col := range []string{address, capacity, available, start_ts, uptime} {
			if sem && !tableAdmin {
				Expect(col).To(Equal(""))
			} else {
				Expect(col).NotTo(Equal(""))
			}
		}
	}
	for rows2.Next() {
		_ = rows2.Scan(&ip)
		if sem && !tableAdmin {
			Expect(ip).To(Equal(""))
		} else {
			Expect(ip).NotTo(Equal(""))
		}
	}
}

func CheckStatusAdmin(ctx context.Context, db *sql.DB, sem bool, statusAdmin bool) {
	var name, value string
	err := db.QueryRowContext(ctx, "show status like '%%tidb_gc_leader_desc%%';").Scan(&name, &value)
	if sem && !statusAdmin {
		Expect(err).To(HaveOccurred())
		Expect(err.Error()).To(ContainSubstring("no rows in result set"))
	} else {
		Expect(err).NotTo(HaveOccurred())
	}
}

func CheckPrivilege(db *sql.DB, user, privilege string) bool {
	var pri string
	if rows, err := db.Query("show grants for " + user); err != nil {
		return false
	} else {
		for rows.Next() {
			if err := rows.Scan(&pri); err != nil {
				return false
			}
			if strings.Contains(pri, privilege) {
				return true
			}
		}
	}
	return false
}