package main

import (
	"context"
	"math"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"

	"github.com/pingcap/errors"
	"github.com/pingcap/log"
	"github.com/pingcap/tidb/pkg/domain"
	"github.com/pingcap/tidb/pkg/plugin"
	"github.com/pingcap/tidb/pkg/session"
	"github.com/pingcap/tidb/pkg/sessionctx"
	"github.com/pingcap/tidb/pkg/sessionctx/vardef"
	"github.com/pingcap/tidb/pkg/sessionctx/variable"
	"github.com/pingcap/tidb/pkg/util/chunk"
	"github.com/pingcap/tidb/pkg/util/sqlexec"
	"github.com/tikv/client-go/v2/util"
	"go.uber.org/zap"
)

const (
	createTable = `CREATE TABLE IF NOT EXISTS mysql.tidb_audit_table_access (
  id int not null auto_increment, 
  user varchar(32) NOT NULL comment "username need be audit(regex)",
  db varchar(32) NOT NULL comment "db name need be audit(regex)",
  tbl varchar(32) NOT NULL comment "table name need be audit(regex)",
  access_type varchar(1024) NOT NULL comment "table access types split by comma",
  create_time timestamp NULL DEFAULT CURRENT_TIMESTAMP comment "record create time",
  update_time timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP comment "record last update time",
  UNIQUE KEY(user, db, tbl),
  PRIMARY KEY (id)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin`
	selectTable               = `select user, db, tbl, access_type from mysql.tidb_audit_table_access`
	TiDBAuditRedactLog        = "tidb_audit_redact_log" // based on tidb_redact_log which applies to regular logs
	TiDBAuditLogMaxSize       = "tidb_audit_log_max_size"
	TiDBAuditLogGeneralEvents = "tidb_audit_log_general_events"
	TiDBAuditLogIncludeExtra  = "tidb_audit_log_include_extras"
	DefAuditLogGeneralEvents  = "STARTING,COMPLETED,ERROR"
)

type AuditLogEnabledEvents struct {
	generalEvents atomic.Value
}

func (a *AuditLogEnabledEvents) SetSysVarTidbAuditLogGeneralEvents(val string) error {
	val, err := CheckAndNormalizeAuditLogGeneralEvents(val)
	if err != nil {
		return err
	}

	strs := strings.Split(val, ",")
	events := make([]bool, int(plugin.GeneralEventCount))
	for _, s := range strs {
		event, err := plugin.GeneralEventFromString(s)
		if err != nil {
			return err
		}
		events[int(event)] = true
	}
	a.generalEvents.Store(events)
	return nil
}

func (a *AuditLogEnabledEvents) IsGeneralEventEnabled(event plugin.GeneralEvent) bool {
	events, ok := a.generalEvents.Load().([]bool)
	if !ok {
		return false
	}

	idx := int(event)
	return idx < len(events) && events[idx]
}

func CheckAndNormalizeAuditLogGeneralEvents(val string) (string, error) {
	strs := strings.Split(val, ",")
	events := make([]string, 0, len(strs))
	for _, s := range strs {
		s = strings.TrimSpace(s)
		if s == "" {
			continue
		}

		event, err := plugin.GeneralEventFromString(s)
		if err != nil {
			return "", err
		}
		events = append(events, event.String())
	}
	return strings.Join(events, ","), nil
}

var (
	redactLog             int32
	auditLogMaxSize       uint64
	auditLogEnabledEvents AuditLogEnabledEvents
	auditLogExtras        atomic.Pointer[map[string]struct{}]
)

const (
	auditLogExtrasClose = "CLOSE"
)

var redactLogSv = &variable.SysVar{
	Name:  TiDBAuditRedactLog,
	Value: "ON", // ON by default, for backward compatibility with previous audit log behavior.
	Scope: vardef.ScopeGlobal,
	Type:  vardef.TypeBool,
	SetGlobal: func(_ context.Context, s *variable.SessionVars, v string) error {
		var newVal int32
		if variable.TiDBOptOn(v) {
			newVal = 1
		}
		atomic.StoreInt32(&redactLog, newVal)
		return nil
	},
}

var auditLogMaxSizeSv = &variable.SysVar{
	Name:     TiDBAuditLogMaxSize,
	Value:    strconv.Itoa(50 << 20), // 50MB
	Scope:    vardef.ScopeGlobal,
	Type:     vardef.TypeUnsigned,
	MinValue: 1 << 20,             // 1MB
	MaxValue: math.MaxInt32 << 20, // limit by lumberjack
	GetGlobal: func(_ context.Context, s *variable.SessionVars) (string, error) {
		return strconv.FormatUint(atomic.LoadUint64(&auditLogMaxSize), 10), nil
	},
	SetGlobal: func(_ context.Context, s *variable.SessionVars, v string) error {
		val, err := strconv.ParseUint(v, 10, 64)
		if err != nil {
			return err
		}
		val = val >> 20 << 20
		if !atomic.CompareAndSwapUint64(&auditLogMaxSize, val, val) {
			atomic.StoreUint64(&auditLogMaxSize, val)
			initLog()
		}
		return nil
	},
}

var auditLogGeneralEvents = &variable.SysVar{
	Name:  TiDBAuditLogGeneralEvents,
	Scope: vardef.ScopeGlobal,
	Value: DefAuditLogGeneralEvents,
	Type:  vardef.TypeStr,
	Validation: func(vars *variable.SessionVars, normalizedValue string, originalValue string, flag vardef.ScopeFlag) (string, error) {
		return CheckAndNormalizeAuditLogGeneralEvents(normalizedValue)
	},
	SetGlobal: func(_ context.Context, s *variable.SessionVars, v string) error {
		return auditLogEnabledEvents.SetSysVarTidbAuditLogGeneralEvents(v)
	},
}

var auditLogIncludeExtrasSv = &variable.SysVar{
	Name:  TiDBAuditLogIncludeExtra,
	Scope: vardef.ScopeGlobal,
	Value: "",
	Type:  vardef.TypeStr,
	Validation: func(vars *variable.SessionVars, normalizedValue string, originalValue string, flag vardef.ScopeFlag) (string, error) {
		// validate the extras
		if normalizedValue == "" {
			return "", nil
		}
		extras := strings.Split(normalizedValue, ",")
		validExtras := map[string]struct{}{
			auditLogExtrasClose: {},
		}
		for _, extra := range extras {
			extra = strings.TrimSpace(strings.ToUpper(extra))
			if _, ok := validExtras[extra]; !ok {
				return "", errors.Errorf("invalid extra %q", extra)
			}
		}
		return normalizedValue, nil
	},
	SetGlobal: func(_ context.Context, s *variable.SessionVars, v string) error {
		newAtomicExtras := make(map[string]struct{})
		if v != "" {
			extras := strings.Split(v, ",")
			for _, extra := range extras {
				extra = strings.TrimSpace(strings.ToUpper(extra))
				newAtomicExtras[extra] = struct{}{}
			}
		}
		auditLogExtras.Store(&newAtomicExtras)
		return nil
	},
}

var global = &handle{}

type handle struct {
	mu struct {
		sync.RWMutex
		dom *domain.Domain
	}
	value atomic.Value
}

func (h *handle) init() error {
	h.mu.RLock()
	dom := h.mu.dom
	h.mu.RUnlock()
	if dom != nil {
		return nil
	}

	h.mu.Lock()
	defer h.mu.Unlock()
	dom, err := session.GetDomain(nil)
	if err != nil {
		return err
	}
	h.mu.dom = dom

	pool := h.mu.dom.SysSessionPool()
	sctx, err := pool.Get()
	if err != nil {
		return err
	}
	defer pool.Put(sctx)

	ctx := util.WithInternalSourceType(context.Background(), "audit_plugin")
	rss, err := sctx.(sqlexec.SQLExecutor).Execute(ctx, createTable)
	if err != nil {
		return err
	}
	for _, rs := range rss {
		_ = rs.Close()
	}
	return h.update()
}

func (h *handle) update() error {
	pool := h.mu.dom.SysSessionPool()
	tmp, err := pool.Get()
	if err != nil {
		return err
	}

	defer pool.Put(tmp)
	sctx := tmp.(sessionctx.Context)

	var audits AuditObjects
	err = audits.loadFromTable(sctx)
	if err != nil {
		return err
	}
	h.updateRedactLog(sctx)
	h.updateAuditLogMaxSize(sctx)
	h.value.Store(&audits)
	return nil
}

// Update the cached value of tidb_audit_redact_log.
func (h *handle) updateRedactLog(sctx sessionctx.Context) {
	val, err := sctx.GetSessionVars().GlobalVarsAccessor.GetGlobalSysVar(TiDBAuditRedactLog)
	if err != nil {
		return
	}
	var redacted int32 = 1
	if !variable.TiDBOptOn(val) {
		redacted = 0
	}
	atomic.StoreInt32(&redactLog, redacted)
}

// Update the cached value of audit_log_max_size.
func (h *handle) updateAuditLogMaxSize(sctx sessionctx.Context) {
	v, err := sctx.GetSessionVars().GlobalVarsAccessor.GetGlobalSysVar(TiDBAuditLogMaxSize)
	if err != nil {
		return
	}
	val, err := strconv.ParseUint(v, 10, 64)
	val = val >> 20 << 20
	if err == nil && !atomic.CompareAndSwapUint64(&auditLogMaxSize, val, val) {
		atomic.StoreUint64(&auditLogMaxSize, val)
		initLog()
	}
}

func (h *handle) get() *AuditObjects {
	return h.value.Load().(*AuditObjects)
}

type AuditObjects []AuditItem

type AuditItem struct {
	AccessType map[string]bool
	Table      *regexp.Regexp
	User       *regexp.Regexp
	DB         *regexp.Regexp
}

// userIsAuditable is similar to check() except that
// I had intended to cache matches to avoid MatchString() calls.
// However, currently the performance overhead did not show up in
// a simple sysbench OLTP test. So I will keep with a simpler design
// until it is proven to be an issue.
func (a *AuditObjects) userIsAuditable(user string) bool {
	for _, item := range *a {
		if item.User == nil || item.User.MatchString(user) {
			return true
		}
	}
	return false
}

func (a *AuditObjects) check(user, db, tbl, subClass string) bool {
	for _, item := range *a {
		if (item.Table == nil || item.Table.MatchString(tbl)) &&
			(item.DB == nil || item.DB.MatchString(db)) &&
			(item.User == nil || item.User.MatchString(user)) &&
			(len(item.AccessType) == 0 || item.AccessType[strings.ToLower(subClass)]) {
			return true
		}
	}
	return false
}

func (a *AuditObjects) loadFromTable(sctx sessionctx.Context) error {
	ctx := util.WithInternalSourceType(context.Background(), "audit_plugin")
	tmp, err := sctx.(sqlexec.SQLExecutor).Execute(ctx, selectTable)
	if err != nil {
		return err
	}
	rs := tmp[0]
	defer rs.Close()

	chk := rs.NewChunk(nil)
	for {
		err = rs.Next(context.TODO(), chk)
		if err != nil {
			return err
		}
		if chk.NumRows() == 0 {
			return nil
		}
		it := chunk.NewIterator4Chunk(chk)
		for row := it.Begin(); row != it.End(); row = it.Next() {
			err = a.decodeTableRow(row)
			if err != nil {
				log.Error("", zap.Error(err))
				// Ignore this row and continue...
				continue
			}
		}
		// NOTE: decodeTableRow decodes data from a chunk Row, that is a shallow copy.
		// The result will reference memory in the chunk, so the chunk must not be reused
		// here, otherwise some werid bug will happen!
		chk = chunk.Renew(chk, sctx.GetSessionVars().MaxChunkSize)
	}
}

func (a *AuditObjects) decodeTableRow(row chunk.Row) error {
	user := row.GetString(0)
	db := row.GetString(1)
	tbl := row.GetString(2)
	accessType := row.GetString(3)
	var (
		userRegex, tblRegex, dbRegex *regexp.Regexp
		err                          error
	)
	if len(user) > 0 {
		userRegex, err = regexp.Compile(user)
		if err != nil {
			return err
		}
	}
	if len(tbl) > 0 {
		tblRegex, err = regexp.Compile(tbl)
		if err != nil {
			return err
		}
	}
	if len(db) > 0 {
		dbRegex, err = regexp.Compile(db)
		if err != nil {
			return err
		}
	}
	accessMap := make(map[string]bool)
	if len(accessType) > 0 {
		for _, item := range strings.Split(accessType, ",") {
			accessMap[strings.ToLower(item)] = true
		}
	}
	item := AuditItem{
		User:       userRegex,
		DB:         dbRegex,
		Table:      tblRegex,
		AccessType: accessMap,
	}
	*a = append(*a, item)
	return nil
}

func auditLogExtrasIncludes(extra string) bool {
	extras := auditLogExtras.Load()
	if extras == nil {
		return false
	}
	_, ok := (*extras)[extra]
	return ok
}