package mysqlutil

import (
	"bytes"
	"context"
	"database/sql"
	"database/sql/driver"
	"encoding/json"
	"errors"
	"fmt"
	"regexp"
	"strings"
	"time"
	"unicode/utf8"

	mysql "github.com/go-sql-driver/mysql"

	"github.com/mem9-ai/drive9/pkg/logger"
	"github.com/mem9-ai/drive9/pkg/metrics"
	"go.uber.org/zap"
)

const (
	RoleMeta       = "meta"
	RoleUser       = "user"
	RoleUserSchema = "user_schema"
	// RoleShared labels *sql.DB handles to shared-schema databases that serve
	// many tenants at once (fs_id row keys). Pool limits for this role must be
	// sized for the whole cluster, not for one tenant.
	RoleShared = "shared"
	// RoleSharedSchema labels the short-lived handle that applies shared-schema
	// bootstrap DDL to a physical shared DB, mirroring the RoleUser /
	// RoleUserSchema split. Schema work — CREATE TABLE, ADD INDEX, contract
	// repair, and the advisory lock that serializes them across pods — routinely
	// blocks for seconds on TiDB, so it runs on its own handle instead of the
	// RoleShared handle that serves tenant traffic. That keeps one cold pool
	// open out of the shared data-plane latency series while leaving schema
	// latency and failures observable under this role.
	RoleSharedSchema = "shared_schema"
)

const (
	maxSQLTraceLen      = 4096
	maxSQLTraceInputLen = maxSQLTraceLen * 2
)

func OpenInstrumented(ctx context.Context, dsn, role string) (*sql.DB, error) {
	return OpenInstrumentedForTenant(ctx, dsn, role, "")
}

func OpenInstrumentedForTenant(ctx context.Context, dsn, role, tenantID string) (*sql.DB, error) {
	return OpenInstrumentedForTenantWithOrg(ctx, dsn, role, tenantID, "")
}

func OpenInstrumentedForTenantWithOrg(ctx context.Context, dsn, role, tenantID, tidbCloudOrgID string) (*sql.DB, error) {
	return openInstrumented(ctx, dsn, role, func(db *sql.DB) {
		metrics.RegisterTenantDBWithOrg(role, tenantID, tidbCloudOrgID, db)
	})
}

func OpenInstrumentedForSharedDB(ctx context.Context, dsn, dbPoolUUID, tidbCloudOrgID string) (*sql.DB, error) {
	return openInstrumented(ctx, dsn, RoleShared, func(db *sql.DB) {
		metrics.RegisterSharedDBWithOrg(dbPoolUUID, tidbCloudOrgID, db)
	})
}

// OpenInstrumentedForSharedSchemaDB opens the dedicated handle used to apply
// shared-schema DDL to one physical shared DB. Callers close it as soon as the
// schema ensure finishes; see RoleSharedSchema for why this work does not share
// the data-plane handle.
func OpenInstrumentedForSharedSchemaDB(ctx context.Context, dsn, dbPoolUUID, tidbCloudOrgID string) (*sql.DB, error) {
	return openInstrumented(ctx, dsn, RoleSharedSchema, func(db *sql.DB) {
		metrics.RegisterSharedSchemaDBWithOrg(dbPoolUUID, tidbCloudOrgID, db)
	})
}

// NormalizeMetaDSNTimeContract applies the metadb time contract for API-key and
// other meta DATETIME columns that compare against UTC_TIMESTAMP(3):
//   - go-sql-driver loc=UTC so bound time.Time encodes as UTC wall clock
//   - session time_zone='+00:00' so DB-side CURRENT_TIMESTAMP(3) also writes UTC
//
// Only RoleMeta opens should use this. Tenant/shared/schema handles keep the
// caller DSN unchanged so DEFAULT/ON UPDATE CURRENT_TIMESTAMP deployments are
// not silently reinterpreted.
func NormalizeMetaDSNTimeContract(dsn string) (string, error) {
	cfg, err := mysql.ParseDSN(dsn)
	if err != nil {
		return "", fmt.Errorf("parse mysql dsn: %w", err)
	}
	cfg.Loc = time.UTC
	if cfg.Params == nil {
		cfg.Params = make(map[string]string)
	}
	// Quoted so the driver issues: SET time_zone='+00:00'
	cfg.Params["time_zone"] = "'+00:00'"
	return cfg.FormatDSN(), nil
}

// NormalizeDSNLocUTC is kept as a thin alias for tests that only need loc=UTC.
// Prefer NormalizeMetaDSNTimeContract for production metadb opens.
func NormalizeDSNLocUTC(dsn string) (string, error) {
	cfg, err := mysql.ParseDSN(dsn)
	if err != nil {
		return "", fmt.Errorf("parse mysql dsn: %w", err)
	}
	cfg.Loc = time.UTC
	return cfg.FormatDSN(), nil
}

// applyClientSQLDefaults turns on client-side placeholder interpolation so
// database/sql Query/Exec with args stay on COM_QUERY instead of falling
// back to COM_STMT_PREPARE + COM_STMT_EXECUTE + COM_STMT_CLOSE. That extra
// protocol round-trip is what makes TiDB SQL QPS look ~3x a logical statement.
//
// InterpolateParams is forced on even if the caller DSN set it false: every
// instrumented open shares this contract. Rollback is a revert, not a flag.
// The driver still returns ErrSkip (and database/sql falls back to prepare)
// for unsupported value types or payloads larger than maxAllowedPacket, so
// prepare metrics may not go fully to zero on large blob writes.
func applyClientSQLDefaults(dsn string) (string, error) {
	cfg, err := mysql.ParseDSN(dsn)
	if err != nil {
		return "", fmt.Errorf("parse mysql dsn: %w", err)
	}
	cfg.InterpolateParams = true
	return cfg.FormatDSN(), nil
}

// normalizeOpenDSN applies the role-specific DSN contract, then the client
// defaults that every instrumented open shares.
func normalizeOpenDSN(dsn, role string) (string, error) {
	openDSN := dsn
	if role == RoleMeta {
		normalized, err := NormalizeMetaDSNTimeContract(dsn)
		if err != nil {
			return "", err
		}
		openDSN = normalized
	}
	return applyClientSQLDefaults(openDSN)
}

func openInstrumented(ctx context.Context, dsn, role string, register func(*sql.DB)) (*sql.DB, error) {
	openDSN, err := normalizeOpenDSN(dsn, role)
	if err != nil {
		return nil, err
	}
	connector, err := (&mysql.MySQLDriver{}).OpenConnector(openDSN)
	if err != nil {
		return nil, fmt.Errorf("open mysql connector: %w", err)
	}
	db := sql.OpenDB(instrumentedConnector{base: connector, role: role})
	ApplyPoolDefaults(db, role)
	register(db)
	if err := db.PingContext(ctx); err != nil {
		metrics.UnregisterDB(db)
		_ = db.Close()
		return nil, fmt.Errorf("ping mysql: %w", err)
	}
	return db, nil
}

func CloseInstrumented(db *sql.DB) error {
	if db == nil {
		return nil
	}
	metrics.UnregisterDB(db)
	return db.Close()
}

type instrumentedConnector struct {
	base driver.Connector
	role string
}

func (c instrumentedConnector) Connect(ctx context.Context) (driver.Conn, error) {
	start := time.Now()
	conn, err := c.base.Connect(ctx)
	if !errors.Is(err, driver.ErrSkip) {
		observeDBOperation(ctx, c.role, "connect", "", start, err)
	}
	if err != nil {
		return nil, err
	}
	return instrumentedConn{base: conn, role: c.role}, nil
}

func (c instrumentedConnector) Driver() driver.Driver {
	return c.base.Driver()
}

type instrumentedConn struct {
	base driver.Conn
	role string
}

func (c instrumentedConn) Prepare(query string) (driver.Stmt, error) {
	start := time.Now()
	stmt, err := c.base.Prepare(query)
	if !errors.Is(err, driver.ErrSkip) {
		observeDBOperation(context.Background(), c.role, "prepare", query, start, err)
	}
	if err != nil {
		return nil, err
	}
	return instrumentedStmt{base: stmt, role: c.role, query: query}, nil
}

func (c instrumentedConn) Close() error {
	return c.base.Close()
}

func (c instrumentedConn) Begin() (driver.Tx, error) {
	return c.BeginTx(context.Background(), driver.TxOptions{})
}

func (c instrumentedConn) Ping(ctx context.Context) error {
	pinger, ok := c.base.(driver.Pinger)
	if !ok {
		return nil
	}
	start := time.Now()
	err := pinger.Ping(ctx)
	if !errors.Is(err, driver.ErrSkip) {
		observeDBOperation(ctx, c.role, "ping", "", start, err)
	}
	return err
}

func (c instrumentedConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
	prep, ok := c.base.(driver.ConnPrepareContext)
	if !ok {
		return c.Prepare(query)
	}
	start := time.Now()
	stmt, err := prep.PrepareContext(ctx, query)
	if !errors.Is(err, driver.ErrSkip) {
		observeDBOperation(ctx, c.role, "prepare", query, start, err)
	}
	if err != nil {
		return nil, err
	}
	return instrumentedStmt{base: stmt, role: c.role, query: query}, nil
}

func (c instrumentedConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
	if beginner, ok := c.base.(driver.ConnBeginTx); ok {
		start := time.Now()
		tx, err := beginner.BeginTx(ctx, opts)
		if !errors.Is(err, driver.ErrSkip) {
			observeDBOperation(ctx, c.role, "begin", "", start, err)
		}
		if err != nil {
			return nil, err
		}
		return instrumentedTx{base: tx, role: c.role}, nil
	}
	if opts.Isolation != driver.IsolationLevel(0) || opts.ReadOnly {
		return nil, fmt.Errorf("driver does not support non-default transaction options")
	}
	return nil, driver.ErrSkip
}

func (c instrumentedConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
	execer, ok := c.base.(driver.ExecerContext)
	if !ok {
		return nil, driver.ErrSkip
	}
	start := time.Now()
	res, err := execer.ExecContext(ctx, query, rewriteInterpolatedJSONArgs(args))
	if !errors.Is(err, driver.ErrSkip) {
		observeDBOperation(ctx, c.role, "exec", query, start, err, res)
	}
	return res, err
}

func (c instrumentedConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
	queryer, ok := c.base.(driver.QueryerContext)
	if !ok {
		return nil, driver.ErrSkip
	}
	start := time.Now()
	rows, err := queryer.QueryContext(ctx, query, rewriteInterpolatedJSONArgs(args))
	if !errors.Is(err, driver.ErrSkip) {
		observeDBOperation(ctx, c.role, "query", query, start, err)
	}
	return rows, err
}

func (c instrumentedConn) ResetSession(ctx context.Context) error {
	resetter, ok := c.base.(driver.SessionResetter)
	if !ok {
		return nil
	}
	return resetter.ResetSession(ctx)
}

func (c instrumentedConn) IsValid() bool {
	validator, ok := c.base.(driver.Validator)
	if !ok {
		return true
	}
	return validator.IsValid()
}

func (c instrumentedConn) CheckNamedValue(nv *driver.NamedValue) error {
	checker, ok := c.base.(driver.NamedValueChecker)
	if !ok {
		return driver.ErrSkip
	}
	return checker.CheckNamedValue(nv)
}

type instrumentedStmt struct {
	base  driver.Stmt
	role  string
	query string
}

func (s instrumentedStmt) Close() error {
	return s.base.Close()
}

func (s instrumentedStmt) NumInput() int {
	return s.base.NumInput()
}

func (s instrumentedStmt) Exec(args []driver.Value) (driver.Result, error) {
	return s.ExecContext(context.Background(), toNamedValues(args))
}

func (s instrumentedStmt) Query(args []driver.Value) (driver.Rows, error) {
	return s.QueryContext(context.Background(), toNamedValues(args))
}

func (s instrumentedStmt) ExecContext(ctx context.Context, args []driver.NamedValue) (driver.Result, error) {
	execer, ok := s.base.(driver.StmtExecContext)
	if !ok {
		return nil, driver.ErrSkip
	}
	start := time.Now()
	res, err := execer.ExecContext(ctx, args)
	if !errors.Is(err, driver.ErrSkip) {
		observeDBOperation(ctx, s.role, "exec", s.query, start, err, res)
	}
	return res, err
}

func (s instrumentedStmt) QueryContext(ctx context.Context, args []driver.NamedValue) (driver.Rows, error) {
	queryer, ok := s.base.(driver.StmtQueryContext)
	if !ok {
		return nil, driver.ErrSkip
	}
	start := time.Now()
	rows, err := queryer.QueryContext(ctx, args)
	if !errors.Is(err, driver.ErrSkip) {
		observeDBOperation(ctx, s.role, "query", s.query, start, err)
	}
	return rows, err
}

func (s instrumentedStmt) CheckNamedValue(nv *driver.NamedValue) error {
	checker, ok := s.base.(driver.NamedValueChecker)
	if !ok {
		return driver.ErrSkip
	}
	return checker.CheckNamedValue(nv)
}

type instrumentedTx struct {
	base driver.Tx
	role string
}

func (tx instrumentedTx) Commit() error {
	start := time.Now()
	err := tx.base.Commit()
	if !errors.Is(err, driver.ErrSkip) {
		observeDBOperation(context.Background(), tx.role, "commit", "", start, err)
	}
	return err
}

func (tx instrumentedTx) Rollback() error {
	start := time.Now()
	err := tx.base.Rollback()
	if !errors.Is(err, driver.ErrSkip) {
		observeDBOperation(context.Background(), tx.role, "rollback", "", start, err)
	}
	return err
}

func observeDBOperation(ctx context.Context, role, operation, query string, start time.Time, err error, results ...driver.Result) {
	elapsed := time.Since(start)
	result := dbResult(err)
	metrics.RecordDBOperation(role, operation, result, elapsed)

	if !logger.DBTraceLogEnabled() {
		return
	}
	if threshold := logger.DBSlowTraceThreshold(); threshold > 0 && elapsed < threshold {
		return
	}

	fields := []zap.Field{
		zap.String("role", role),
		zap.String("operation", operation),
		zap.String("result", result),
		zap.Float64("duration_ms", float64(elapsed.Microseconds())/1000.0),
	}
	if sqlText, truncated, sourceLen := sqlTraceText(query); sqlText != "" {
		fields = append(fields, zap.String("sql", sqlText))
		fields = append(fields, zap.Int("sql_len", sourceLen))
		if truncated {
			fields = append(fields, zap.Bool("sql_truncated", true))
		}
	}
	if len(results) > 0 && results[0] != nil {
		if affected, rowsErr := results[0].RowsAffected(); rowsErr == nil {
			fields = append(fields, zap.Int64("rows_affected", affected))
		}
	}
	if err != nil {
		fields = append(fields, safeDBErrorFields(err)...)
	}
	logger.Info(ctx, "db_operation_timing", fields...)
}

func dbResult(err error) string {
	// Delegate to metrics.ResultForError so DB operation labels match worker
	// labels, especially for transient bad connections.
	return metrics.ResultForError(err)
}

func toNamedValues(args []driver.Value) []driver.NamedValue {
	if len(args) == 0 {
		return nil
	}
	named := make([]driver.NamedValue, 0, len(args))
	for idx, arg := range args {
		named = append(named, driver.NamedValue{
			Ordinal: idx + 1,
			Value:   arg,
		})
	}
	return named
}

// rewriteInterpolatedJSONArgs converts JSON object/array/null []byte args to
// string so interpolateParams quotes them as utf8mb4 instead of _binary'...'.
// MySQL 8 rejects JSON values created from a binary-charset string (error
// 3144). Real BLOBs stay []byte and keep the _binary prefix. json.RawMessage
// is left alone: the driver already interpolates it without _binary.
// A first-byte `{`/`[` check is not enough — GCM ciphertext and arbitrary
// file bytes hit those prefixes ~1/128 of the time — so []byte is rewritten
// only when it is valid UTF-8 JSON object/array/null.
func rewriteInterpolatedJSONArgs(args []driver.NamedValue) []driver.NamedValue {
	var out []driver.NamedValue
	for i, arg := range args {
		b, ok := arg.Value.([]byte)
		if !ok || !looksLikeJSON(b) {
			if out != nil {
				out[i] = arg
			}
			continue
		}
		if out == nil {
			out = make([]driver.NamedValue, len(args))
			copy(out, args[:i])
		}
		arg.Value = string(b)
		out[i] = arg
	}
	if out == nil {
		return args
	}
	return out
}

func looksLikeJSON(b []byte) bool {
	s := bytes.TrimSpace(b)
	if len(s) == 0 {
		return false
	}
	switch s[0] {
	case '{', '[':
	default:
		if len(s) != 4 || string(s) != "null" {
			return false
		}
	}
	return utf8.Valid(s) && json.Valid(s)
}

var sqlWhitespace = regexp.MustCompile(`\s+`)

func sqlTraceText(query string) (string, bool, int) {
	sourceLen := len(query)
	if sourceLen == 0 {
		return "", false, 0
	}
	truncated := false
	if len(query) > maxSQLTraceInputLen {
		query = query[:maxSQLTraceInputLen]
		truncated = true
	}
	out := normalizeSQLForTrace(query)
	if len(out) > maxSQLTraceLen {
		out = out[:maxSQLTraceLen]
		truncated = true
	}
	return out, truncated, sourceLen
}

func normalizeSQLForTrace(query string) string {
	return sqlWhitespace.ReplaceAllString(strings.TrimSpace(redactSQLLiterals(query)), " ")
}

func redactSQLLiterals(query string) string {
	var b strings.Builder
	b.Grow(len(query))
	for i := 0; i < len(query); {
		switch query[i] {
		case '\'', '"':
			b.WriteByte('?')
			i = skipQuotedSQLLiteral(query, i, query[i])
		case '/':
			if i+1 < len(query) && query[i+1] == '*' {
				b.WriteByte(' ')
				i = skipBlockSQLComment(query, i)
				continue
			}
			b.WriteByte(query[i])
			i++
		case '-':
			if i+1 < len(query) && query[i+1] == '-' {
				b.WriteByte(' ')
				i = skipLineSQLComment(query, i)
				continue
			}
			b.WriteByte(query[i])
			i++
		case '#':
			b.WriteByte(' ')
			i = skipLineSQLComment(query, i)
		default:
			if isNumericSQLLiteralStart(query, i) {
				b.WriteByte('?')
				i = skipNumericSQLLiteral(query, i)
				continue
			}
			b.WriteByte(query[i])
			i++
		}
	}
	return b.String()
}

func safeDBErrorFields(err error) []zap.Field {
	var mysqlErr *mysql.MySQLError
	if errors.As(err, &mysqlErr) {
		fields := []zap.Field{
			zap.String("db_error_type", "mysql"),
			zap.Uint16("db_error_number", mysqlErr.Number),
		}
		if mysqlErr.SQLState != [5]byte{} {
			fields = append(fields, zap.String("db_error_sql_state", string(mysqlErr.SQLState[:])))
		}
		return fields
	}
	switch {
	case errors.Is(err, context.Canceled):
		return []zap.Field{zap.String("db_error_type", "context_canceled")}
	case errors.Is(err, context.DeadlineExceeded):
		return []zap.Field{zap.String("db_error_type", "context_deadline_exceeded")}
	case errors.Is(err, driver.ErrBadConn):
		return []zap.Field{zap.String("db_error_type", "bad_conn")}
	default:
		return []zap.Field{zap.String("db_error_type", fmt.Sprintf("%T", err))}
	}
}

func skipQuotedSQLLiteral(query string, start int, quote byte) int {
	for i := start + 1; i < len(query); i++ {
		switch query[i] {
		case '\\':
			if i+1 < len(query) {
				i++
			}
		case quote:
			if i+1 < len(query) && query[i+1] == quote {
				i++
				continue
			}
			return i + 1
		}
	}
	return len(query)
}

func skipBlockSQLComment(query string, start int) int {
	for i := start + 2; i < len(query)-1; i++ {
		if query[i] == '*' && query[i+1] == '/' {
			return i + 2
		}
	}
	return len(query)
}

func skipLineSQLComment(query string, start int) int {
	for i := start + 1; i < len(query); i++ {
		if query[i] == '\n' || query[i] == '\r' {
			return i
		}
	}
	return len(query)
}

func isNumericSQLLiteralStart(query string, i int) bool {
	if query[i] < '0' || query[i] > '9' {
		return false
	}
	return i == 0 || !isSQLIdentByte(query[i-1])
}

func skipNumericSQLLiteral(query string, start int) int {
	i := start
	for i < len(query) {
		c := query[i]
		if (c >= '0' && c <= '9') ||
			(c >= 'a' && c <= 'f') ||
			(c >= 'A' && c <= 'F') ||
			c == 'x' || c == 'X' || c == '.' || c == '+' || c == '-' {
			i++
			continue
		}
		break
	}
	return i
}

func isSQLIdentByte(c byte) bool {
	return (c >= 'a' && c <= 'z') ||
		(c >= 'A' && c <= 'Z') ||
		(c >= '0' && c <= '9') ||
		c == '_' ||
		c == '$'
}