package mysqlutil

import (
	"context"
	"database/sql/driver"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"testing"
	"time"

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

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

func TestObserveDBOperationLogsSQLTraceWithoutArgs(t *testing.T) {
	t.Setenv("DRIVE9_DB_TRACE_LOG_ENABLED", "true")
	logger.ResetBenchTimingLogEnabledForTest()
	logger.ResetDBTraceLogEnabledForTest()
	t.Cleanup(logger.ResetBenchTimingLogEnabledForTest)
	t.Cleanup(logger.ResetDBTraceLogEnabledForTest)

	core, recorded := observer.New(zap.InfoLevel)
	prevLogger := logger.L()
	logger.Set(zap.New(core))
	t.Cleanup(func() { logger.Set(prevLogger) })

	ctx := context.Background()
	start := time.Now().Add(-350 * time.Millisecond)
	observeDBOperation(ctx, RoleUser, "exec", "  INSERT INTO file_nodes\n(path, name) VALUES (?, ?)  ", start, errors.New("duplicate key"))

	entries := recorded.FilterMessage("db_operation_timing").All()
	if len(entries) != 1 {
		t.Fatalf("db_operation_timing entries = %d, want 1", len(entries))
	}
	fields := entries[0].ContextMap()
	if fields["role"] != RoleUser {
		t.Errorf("role field = %v, want %s", fields["role"], RoleUser)
	}
	if fields["operation"] != "exec" {
		t.Errorf("operation field = %v, want exec", fields["operation"])
	}
	if fields["result"] != "error" {
		t.Errorf("result field = %v, want error", fields["result"])
	}
	if fields["sql"] != "INSERT INTO file_nodes (path, name) VALUES (?, ?)" {
		t.Errorf("sql field = %v", fields["sql"])
	}
	if _, ok := fields["args"]; ok {
		t.Errorf("db trace log must not include args: %#v", fields)
	}
	if fields["duration_ms"] == nil {
		t.Errorf("duration_ms field missing: %#v", fields)
	}
}

func TestObserveDBOperationLogsSQLTraceWhenDBTraceEnabled(t *testing.T) {
	t.Setenv("DRIVE9_BENCH_TIMING_LOG_ENABLED", "false")
	t.Setenv("DRIVE9_DB_TRACE_LOG_ENABLED", "true")
	logger.ResetBenchTimingLogEnabledForTest()
	logger.ResetDBTraceLogEnabledForTest()
	t.Cleanup(logger.ResetBenchTimingLogEnabledForTest)
	t.Cleanup(logger.ResetDBTraceLogEnabledForTest)

	core, recorded := observer.New(zap.InfoLevel)
	prevLogger := logger.L()
	logger.Set(zap.New(core))
	t.Cleanup(func() { logger.Set(prevLogger) })

	observeDBOperation(context.Background(), RoleUser, "query", "SELECT * FROM file_nodes WHERE path = ?", time.Now().Add(-350*time.Millisecond), nil)

	entries := recorded.FilterMessage("db_operation_timing").All()
	if len(entries) != 1 {
		t.Fatalf("db_operation_timing entries = %d, want 1", len(entries))
	}
	fields := entries[0].ContextMap()
	if fields["sql"] != "SELECT * FROM file_nodes WHERE path = ?" {
		t.Errorf("sql field = %v", fields["sql"])
	}
}

func TestNormalizeSQLForTraceRedactsLiterals(t *testing.T) {
	query := `SELECT * FROM files
		WHERE name = 'agent secret'
		AND path = "project literal"
		AND size > 123
		AND note = 'can''t leak'
		AND raw = 'can\'t leak either'`
	want := `SELECT * FROM files WHERE name = ? AND path = ? AND size > ? AND note = ? AND raw = ?`
	if got := normalizeSQLForTrace(query); got != want {
		t.Fatalf("normalizeSQLForTrace() = %q, want %q", got, want)
	}
}

func TestNormalizeSQLForTraceStripsComments(t *testing.T) {
	query := `SELECT '/* quoted@example.com */', "-- quoted-dash@example.com", '# quoted-hash@example.com'
		/* block@example.com */ FROM files -- dash@example.com
		WHERE id = 42 # hash@example.com`
	want := `SELECT ?, ?, ? FROM files WHERE id = ?`
	got := normalizeSQLForTrace(query)
	if got != want {
		t.Fatalf("normalizeSQLForTrace() = %q, want %q", got, want)
	}
	for _, secret := range []string{"block@example.com", "dash@example.com", "hash@example.com"} {
		if strings.Contains(got, secret) {
			t.Fatalf("normalizeSQLForTrace leaked comment secret %q in %q", secret, got)
		}
	}
}

func TestObserveDBOperationLogsSafeErrorDetails(t *testing.T) {
	t.Setenv("DRIVE9_DB_TRACE_LOG_ENABLED", "true")
	logger.ResetBenchTimingLogEnabledForTest()
	logger.ResetDBTraceLogEnabledForTest()
	t.Cleanup(logger.ResetBenchTimingLogEnabledForTest)
	t.Cleanup(logger.ResetDBTraceLogEnabledForTest)

	core, recorded := observer.New(zap.InfoLevel)
	prevLogger := logger.L()
	logger.Set(zap.New(core))
	t.Cleanup(func() { logger.Set(prevLogger) })

	const secret = "CHILD_SECRET"
	dbErr := &mysql.MySQLError{
		Number:  1062,
		Message: "Duplicate entry '" + secret + "' for key 'files.name'",
	}
	observeDBOperation(context.Background(), RoleUser, "exec", "INSERT INTO files(name) VALUES (?)", time.Now().Add(-350*time.Millisecond), dbErr)

	entries := recorded.FilterMessage("db_operation_timing").All()
	if len(entries) != 1 {
		t.Fatalf("db_operation_timing entries = %d, want 1", len(entries))
	}
	fields := entries[0].ContextMap()
	if _, ok := fields["error"]; ok {
		t.Errorf("db trace log must not include raw error field: %#v", fields)
	}
	if fields["db_error_type"] != "mysql" {
		t.Errorf("db_error_type = %v, want mysql", fields["db_error_type"])
	}
	if fmt.Sprint(fields["db_error_number"]) != "1062" {
		t.Errorf("db_error_number = %v, want 1062", fields["db_error_number"])
	}
	assertFieldsDoNotContain(t, fields, secret)
}

func TestObserveDBOperationBoundsSQLTraceField(t *testing.T) {
	t.Setenv("DRIVE9_DB_TRACE_LOG_ENABLED", "true")
	logger.ResetBenchTimingLogEnabledForTest()
	logger.ResetDBTraceLogEnabledForTest()
	t.Cleanup(logger.ResetBenchTimingLogEnabledForTest)
	t.Cleanup(logger.ResetDBTraceLogEnabledForTest)

	core, recorded := observer.New(zap.InfoLevel)
	prevLogger := logger.L()
	logger.Set(zap.New(core))
	t.Cleanup(func() { logger.Set(prevLogger) })

	query := "SELECT " + strings.Repeat("x", maxSQLTraceLen*3) + " FROM files"
	observeDBOperation(context.Background(), RoleUser, "query", query, time.Now().Add(-350*time.Millisecond), nil)

	entries := recorded.FilterMessage("db_operation_timing").All()
	if len(entries) != 1 {
		t.Fatalf("db_operation_timing entries = %d, want 1", len(entries))
	}
	fields := entries[0].ContextMap()
	sqlText, ok := fields["sql"].(string)
	if !ok {
		t.Fatalf("sql field = %T, want string", fields["sql"])
	}
	if len(sqlText) != maxSQLTraceLen {
		t.Fatalf("sql field len = %d, want %d", len(sqlText), maxSQLTraceLen)
	}
	if fields["sql_truncated"] != true {
		t.Errorf("sql_truncated = %v, want true", fields["sql_truncated"])
	}
	if fmt.Sprint(fields["sql_len"]) != fmt.Sprint(len(query)) {
		t.Errorf("sql_len = %v, want %d", fields["sql_len"], len(query))
	}
}

func TestObserveDBOperationSkipsTraceFieldsWhenBenchTimingDisabled(t *testing.T) {
	t.Setenv("DRIVE9_BENCH_TIMING_LOG_ENABLED", "false")
	t.Setenv("DRIVE9_DB_TRACE_LOG_ENABLED", "false")
	logger.ResetBenchTimingLogEnabledForTest()
	logger.ResetDBTraceLogEnabledForTest()
	t.Cleanup(logger.ResetBenchTimingLogEnabledForTest)
	t.Cleanup(logger.ResetDBTraceLogEnabledForTest)

	core, recorded := observer.New(zap.InfoLevel)
	prevLogger := logger.L()
	logger.Set(zap.New(core))
	t.Cleanup(func() { logger.Set(prevLogger) })

	res := &recordingResult{}
	observeDBOperation(context.Background(), RoleUser, "exec", "SELECT 'secret'", time.Now(), nil, res)

	if entries := recorded.FilterMessage("db_operation_timing").All(); len(entries) != 0 {
		t.Fatalf("db_operation_timing entries = %d, want 0", len(entries))
	}
	if res.rowsAffectedCalled {
		t.Fatal("RowsAffected called while DB trace logs are disabled")
	}
}

func TestObserveDBOperationSkipsTraceWhenOnlyBenchTimingEnabled(t *testing.T) {
	t.Setenv("DRIVE9_BENCH_TIMING_LOG_ENABLED", "true")
	t.Setenv("DRIVE9_DB_TRACE_LOG_ENABLED", "false")
	logger.ResetBenchTimingLogEnabledForTest()
	logger.ResetDBTraceLogEnabledForTest()
	t.Cleanup(logger.ResetBenchTimingLogEnabledForTest)
	t.Cleanup(logger.ResetDBTraceLogEnabledForTest)

	core, recorded := observer.New(zap.InfoLevel)
	prevLogger := logger.L()
	logger.Set(zap.New(core))
	t.Cleanup(func() { logger.Set(prevLogger) })

	observeDBOperation(context.Background(), RoleUser, "query", "SELECT * FROM file_nodes", time.Now(), nil)

	if entries := recorded.FilterMessage("db_operation_timing").All(); len(entries) != 0 {
		t.Fatalf("db_operation_timing entries = %d, want 0", len(entries))
	}
}

func TestObserveDBOperationSkipsTraceBelowDefaultSlowThreshold(t *testing.T) {
	t.Setenv("DRIVE9_DB_TRACE_LOG_ENABLED", "true")
	logger.ResetDBTraceLogEnabledForTest()
	logger.ResetDBSlowTraceThresholdForTest()
	t.Cleanup(logger.ResetDBTraceLogEnabledForTest)
	t.Cleanup(logger.ResetDBSlowTraceThresholdForTest)

	core, recorded := observer.New(zap.InfoLevel)
	prevLogger := logger.L()
	logger.Set(zap.New(core))
	t.Cleanup(func() { logger.Set(prevLogger) })

	observeDBOperation(context.Background(), RoleUser, "query", "SELECT * FROM file_nodes", time.Now().Add(-50*time.Millisecond), nil)

	if entries := recorded.FilterMessage("db_operation_timing").All(); len(entries) != 0 {
		t.Fatalf("db_operation_timing entries = %d, want 0", len(entries))
	}
}

func TestObserveDBOperationLogsTraceAboveDefaultSlowThreshold(t *testing.T) {
	logger.ResetDBTraceLogEnabledForTest()
	logger.ResetDBSlowTraceThresholdForTest()
	t.Cleanup(logger.ResetDBTraceLogEnabledForTest)
	t.Cleanup(logger.ResetDBSlowTraceThresholdForTest)

	core, recorded := observer.New(zap.InfoLevel)
	prevLogger := logger.L()
	logger.Set(zap.New(core))
	t.Cleanup(func() { logger.Set(prevLogger) })

	observeDBOperation(context.Background(), RoleUser, "query", "SELECT * FROM file_nodes", time.Now().Add(-350*time.Millisecond), nil)

	if entries := recorded.FilterMessage("db_operation_timing").All(); len(entries) != 1 {
		t.Fatalf("db_operation_timing entries = %d, want 1", len(entries))
	}
}

func TestObserveDBOperationLogsTraceWhenSlowThresholdIsZero(t *testing.T) {
	t.Setenv("DRIVE9_DB_TRACE_LOG_ENABLED", "true")
	t.Setenv("DRIVE9_DB_SLOW_TRACE_MS", "0")
	logger.ResetDBTraceLogEnabledForTest()
	logger.ResetDBSlowTraceThresholdForTest()
	t.Cleanup(logger.ResetDBTraceLogEnabledForTest)
	t.Cleanup(logger.ResetDBSlowTraceThresholdForTest)

	core, recorded := observer.New(zap.InfoLevel)
	prevLogger := logger.L()
	logger.Set(zap.New(core))
	t.Cleanup(func() { logger.Set(prevLogger) })

	observeDBOperation(context.Background(), RoleUser, "query", "SELECT * FROM file_nodes", time.Now(), nil)

	if entries := recorded.FilterMessage("db_operation_timing").All(); len(entries) != 1 {
		t.Fatalf("db_operation_timing entries = %d, want 1", len(entries))
	}
}

type recordingResult struct {
	rowsAffectedCalled bool
}

func (r *recordingResult) LastInsertId() (int64, error) {
	return 0, driver.ErrSkip
}

func (r *recordingResult) RowsAffected() (int64, error) {
	r.rowsAffectedCalled = true
	return 1, nil
}

func assertFieldsDoNotContain(t *testing.T, fields map[string]interface{}, secret string) {
	t.Helper()
	for key, value := range fields {
		if strings.Contains(fmt.Sprint(value), secret) {
			t.Fatalf("field %q leaked %q in %#v", key, secret, fields)
		}
	}
}

func TestNormalizeDSNLocUTCForcesUTC(t *testing.T) {
	in := "user:pass@tcp(127.0.0.1:3306)/db?parseTime=true&loc=Asia%2FShanghai"
	out, err := NormalizeDSNLocUTC(in)
	if err != nil {
		t.Fatal(err)
	}
	cfg, err := mysql.ParseDSN(out)
	if err != nil {
		t.Fatal(err)
	}
	if cfg.Loc == nil || cfg.Loc.String() != "UTC" {
		t.Fatalf("Loc=%v, want UTC", cfg.Loc)
	}
	// ParseTime and other params preserved.
	if !cfg.ParseTime {
		t.Fatal("ParseTime should remain true")
	}
}

func TestNormalizeDSNLocUTCRejectsInvalid(t *testing.T) {
	if _, err := NormalizeDSNLocUTC(":::not-a-dsn"); err == nil {
		t.Fatal("want error for invalid dsn")
	}
}

func TestNormalizeMetaDSNTimeContractForcesLocAndSessionUTC(t *testing.T) {
	in := "user:pass@tcp(127.0.0.1:3306)/db?parseTime=true&loc=Asia%2FShanghai&time_zone=%27%2B08%3A00%27"
	out, err := NormalizeMetaDSNTimeContract(in)
	if err != nil {
		t.Fatal(err)
	}
	cfg, err := mysql.ParseDSN(out)
	if err != nil {
		t.Fatal(err)
	}
	if cfg.Loc == nil || cfg.Loc.String() != "UTC" {
		t.Fatalf("Loc=%v, want UTC", cfg.Loc)
	}
	if cfg.Params["time_zone"] != "'+00:00'" {
		t.Fatalf("time_zone param=%q, want '+00:00'", cfg.Params["time_zone"])
	}
	if !cfg.ParseTime {
		t.Fatal("ParseTime should remain true")
	}
}

func TestNormalizeMetaDSNTimeContractRejectsInvalid(t *testing.T) {
	if _, err := NormalizeMetaDSNTimeContract(":::not-a-dsn"); err == nil {
		t.Fatal("want error for invalid dsn")
	}
}

func TestApplyClientSQLDefaultsEnablesInterpolateParams(t *testing.T) {
	in := "user:pass@tcp(127.0.0.1:3306)/db?parseTime=true&tls=skip-verify"
	out, err := applyClientSQLDefaults(in)
	if err != nil {
		t.Fatal(err)
	}
	cfg, err := mysql.ParseDSN(out)
	if err != nil {
		t.Fatal(err)
	}
	if !cfg.InterpolateParams {
		t.Fatal("InterpolateParams should be true")
	}
	if !cfg.ParseTime {
		t.Fatal("ParseTime should remain true")
	}
	if cfg.TLSConfig != "skip-verify" {
		t.Fatalf("TLSConfig=%q, want skip-verify", cfg.TLSConfig)
	}
}

func TestApplyClientSQLDefaultsOverridesExplicitFalse(t *testing.T) {
	in := "user:pass@tcp(127.0.0.1:3306)/db?parseTime=true&interpolateParams=false"
	out, err := applyClientSQLDefaults(in)
	if err != nil {
		t.Fatal(err)
	}
	cfg, err := mysql.ParseDSN(out)
	if err != nil {
		t.Fatal(err)
	}
	if !cfg.InterpolateParams {
		t.Fatal("InterpolateParams should be forced true")
	}
}

func TestApplyClientSQLDefaultsRejectsInvalid(t *testing.T) {
	if _, err := applyClientSQLDefaults(":::not-a-dsn"); err == nil {
		t.Fatal("want error for invalid dsn")
	}
}

func TestNormalizeOpenDSNEnablesInterpolateParamsForAllRoles(t *testing.T) {
	in := "user:pass@tcp(127.0.0.1:3306)/db?parseTime=true"
	roles := []string{RoleMeta, RoleUser, RoleUserSchema, RoleShared, RoleSharedSchema, "unknown"}
	for _, role := range roles {
		out, err := normalizeOpenDSN(in, role)
		if err != nil {
			t.Fatalf("role %s: %v", role, err)
		}
		cfg, err := mysql.ParseDSN(out)
		if err != nil {
			t.Fatalf("role %s parse: %v", role, err)
		}
		if !cfg.InterpolateParams {
			t.Fatalf("role %s InterpolateParams=false, want true", role)
		}
		if !cfg.ParseTime {
			t.Fatalf("role %s ParseTime=false, want true", role)
		}
	}
}

func TestNormalizeOpenDSNAppliesMetaTimeContractOnlyForMeta(t *testing.T) {
	in := "user:pass@tcp(127.0.0.1:3306)/db?parseTime=true&loc=Asia%2FShanghai&time_zone=%27%2B08%3A00%27"

	metaOut, err := normalizeOpenDSN(in, RoleMeta)
	if err != nil {
		t.Fatal(err)
	}
	metaCfg, err := mysql.ParseDSN(metaOut)
	if err != nil {
		t.Fatal(err)
	}
	if !metaCfg.InterpolateParams {
		t.Fatal("meta InterpolateParams should be true")
	}
	if metaCfg.Loc == nil || metaCfg.Loc.String() != "UTC" {
		t.Fatalf("meta Loc=%v, want UTC", metaCfg.Loc)
	}
	if metaCfg.Params["time_zone"] != "'+00:00'" {
		t.Fatalf("meta time_zone=%q, want '+00:00'", metaCfg.Params["time_zone"])
	}

	userOut, err := normalizeOpenDSN(in, RoleUser)
	if err != nil {
		t.Fatal(err)
	}
	userCfg, err := mysql.ParseDSN(userOut)
	if err != nil {
		t.Fatal(err)
	}
	if !userCfg.InterpolateParams {
		t.Fatal("user InterpolateParams should be true")
	}
	if userCfg.Loc == nil || userCfg.Loc.String() != "Asia/Shanghai" {
		t.Fatalf("user Loc=%v, want Asia/Shanghai", userCfg.Loc)
	}
	if userCfg.Params["time_zone"] != "'+08:00'" {
		t.Fatalf("user time_zone=%q, want '+08:00'", userCfg.Params["time_zone"])
	}
}

func TestNormalizeOpenDSNRejectsInvalid(t *testing.T) {
	if _, err := normalizeOpenDSN(":::not-a-dsn", RoleUser); err == nil {
		t.Fatal("want error for invalid user dsn")
	}
	if _, err := normalizeOpenDSN(":::not-a-dsn", RoleMeta); err == nil {
		t.Fatal("want error for invalid meta dsn")
	}
}

func TestLooksLikeJSON(t *testing.T) {
	cases := []struct {
		name string
		in   []byte
		want bool
	}{
		{"object", []byte(`{}`), true},
		{"object_kv", []byte(`{"k":"v"}`), true},
		{"array_ws", []byte("  [0.1,0.2]"), true},
		{"null", []byte("null"), true},
		{"null_ws", []byte("  null  "), true},
		{"bool", []byte("true"), false},
		{"quoted", []byte(`"quoted"`), false},
		{"empty", []byte(""), false},
		{"nil", nil, false},
		{"png", []byte{0x89, 'P', 'N', 'G'}, false},
		{"ascii_blob", []byte("ciphertext-not-json"), false},
		{"brace_invalid_utf8", []byte{'{', 0xff, 0x00, 0xfe}, false},
		{"bracket_invalid_utf8", []byte{'[', 0xff, 0x00, 0xfe}, false},
		{"bracket_then_nul", []byte{'[', 0x00}, false},
		{"gcm_style_brace", []byte{'{', 0x9a, 0x10, 0x03, 0x7b, 0x02}, false},
		{"gcm_style_bracket", []byte{'[', 0x9a, 0x10, 0x03, 0x7b, 0x02}, false},
		{"truncated_object", []byte(`{"k":`), false},
	}
	for _, tc := range cases {
		if got := looksLikeJSON(tc.in); got != tc.want {
			t.Errorf("%s looksLikeJSON(%q)=%v, want %v", tc.name, tc.in, got, tc.want)
		}
	}
}

func TestRewriteJSONBytesArgsConvertsJSONButKeepsBlobs(t *testing.T) {
	label := []byte(`{"env":"test"}`)
	blob := []byte{0x00, 0x01, 0xff}
	braceBlob := []byte{'{', 0xff, 0x00, 0xfe}
	gcmBlob := []byte{'[', 0x9a, 0x10, 0x03, 0x7b, 0x02}
	raw := json.RawMessage(`{"ok":true}`)
	in := []driver.NamedValue{
		{Ordinal: 1, Value: "tenant-1"},
		{Ordinal: 2, Value: label},
		{Ordinal: 3, Value: blob},
		{Ordinal: 4, Value: raw},
		{Ordinal: 5, Value: []byte("null")},
		{Ordinal: 6, Value: int64(7)},
		{Ordinal: 7, Value: braceBlob},
		{Ordinal: 8, Value: gcmBlob},
	}
	out := rewriteInterpolatedJSONArgs(in)
	if out[0].Value != "tenant-1" {
		t.Fatalf("string arg changed: %#v", out[0].Value)
	}
	if got, ok := out[1].Value.(string); !ok || got != `{"env":"test"}` {
		t.Fatalf("label = %#v, want utf8 string", out[1].Value)
	}
	gotBlob, ok := out[2].Value.([]byte)
	if !ok || string(gotBlob) != string(blob) {
		t.Fatalf("blob = %#v, want original []byte", out[2].Value)
	}
	if got, ok := out[3].Value.(json.RawMessage); !ok || string(got) != `{"ok":true}` {
		t.Fatalf("RawMessage = %#v, want unchanged RawMessage (driver interpolates it)", out[3].Value)
	}
	if got, ok := out[4].Value.(string); !ok || got != "null" {
		t.Fatalf("json null = %#v, want string null", out[4].Value)
	}
	if out[5].Value != int64(7) {
		t.Fatalf("int arg changed: %#v", out[5].Value)
	}
	gotBrace, ok := out[6].Value.([]byte)
	if !ok || string(gotBrace) != string(braceBlob) {
		t.Fatalf("brace-prefixed invalid UTF-8 = %#v, want original []byte", out[6].Value)
	}
	gotGCM, ok := out[7].Value.([]byte)
	if !ok || string(gotGCM) != string(gcmBlob) {
		t.Fatalf("GCM-style ciphertext = %#v, want original []byte", out[7].Value)
	}
}

func TestRewriteJSONBytesArgsNoCopyWhenUnchanged(t *testing.T) {
	in := []driver.NamedValue{{Ordinal: 1, Value: "x"}, {Ordinal: 2, Value: []byte{0x01}}}
	out := rewriteInterpolatedJSONArgs(in)
	if &out[0] != &in[0] {
		t.Fatal("expected original slice when no JSON args")
	}
}