package pii

import (
	"encoding/json"
	"strings"
	"testing"
)

// ---------- 掩码函数单元测试 ----------

func TestMaskName(t *testing.T) {
	cases := map[string]string{
		"":     "",
		"张三":   "张***",
		"John": "J***",
		"李":    "李***",
	}
	for in, want := range cases {
		if got := maskName(in); got != want {
			t.Errorf("maskName(%q)=%q, want %q", in, got, want)
		}
	}
}

func TestMaskEmail(t *testing.T) {
	if got := maskEmail("john@corp.com"); got != "j***@corp.com" {
		t.Errorf("maskEmail got %q", got)
	}
	// 非法邮箱退化为完全脱敏
	if got := maskEmail("notanemail"); got != maskFull() {
		t.Errorf("maskEmail(invalid) got %q", got)
	}
	if got := maskEmail(""); got != "" {
		t.Errorf("maskEmail(empty) got %q", got)
	}
}

func TestMaskAccount(t *testing.T) {
	if got := maskAccount("abcdefgh"); got != "ab***gh" {
		t.Errorf("maskAccount got %q", got)
	}
	// 过短完全脱敏
	if got := maskAccount("abc"); got != maskFull() {
		t.Errorf("maskAccount(short) got %q", got)
	}
}

func TestMaskFullFixedLength(t *testing.T) {
	if got := maskFull(); got != strings.Repeat(MaskChar, fullMaskLen) {
		t.Errorf("maskFull got %q", got)
	}
}

// ---------- 优化 B：流式 token 脱敏测试 ----------

func TestMaskJSONBytes_FlatFields(t *testing.T) {
	in := []byte(`{
		"sn": "SN123",
		"userName": "Alice",
		"contactNumber": "13812345678",
		"email": "a@b.com",
		"longitude": 121.5,
		"latitude": 31.2,
		"address": "123 Main St",
		"city": "Shanghai",
		"province": "SH",
		"country": "CN",
		"iccid": "8986012345678901234",
		"imei": "356938035643809",
		"account": "installer01"
	}`)
	out, err := MaskJSONBytes(in)
	if err != nil {
		t.Fatalf("MaskJSONBytes err: %v", err)
	}
	var m map[string]any
	if err := json.Unmarshal(out, &m); err != nil {
		t.Fatalf("output not valid JSON: %v\n%s", err, out)
	}

	// 保留字段
	if m["sn"] != "SN123" {
		t.Errorf("sn should be preserved, got %v", m["sn"])
	}
	if m["iccid"] != "8986012345678901234" {
		t.Errorf("iccid should be preserved, got %v", m["iccid"])
	}
	if m["city"] != "Shanghai" || m["province"] != "SH" || m["country"] != "CN" {
		t.Errorf("admin region should be preserved: city=%v province=%v country=%v", m["city"], m["province"], m["country"])
	}

	// 完全脱敏字段
	full := maskFull()
	for _, k := range []string{"contactNumber", "longitude", "latitude", "address", "imei"} {
		if m[k] != full {
			t.Errorf("%s should be fully masked, got %v", k, m[k])
		}
	}
	// 姓名
	if m["userName"] != "A***" {
		t.Errorf("userName masked wrong: %v", m["userName"])
	}
	// 邮箱
	if m["email"] != "a***@b.com" {
		t.Errorf("email masked wrong: %v", m["email"])
	}
	// 账号
	if m["account"] != "in***01" {
		t.Errorf("account masked wrong: %v", m["account"])
	}
}

func TestMaskJSONBytes_Nested(t *testing.T) {
	in := []byte(`{
		"data": [
			{"userName": "Bob", "sn": "S1", "nested": {"phone": "111", "keep": "ok"}},
			{"userName": "Carol", "sn": "S2"}
		],
		"total": 2
	}`)
	out, err := MaskJSONBytes(in)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	var m map[string]any
	if err := json.Unmarshal(out, &m); err != nil {
		t.Fatalf("invalid JSON: %v\n%s", err, out)
	}
	arr := m["data"].([]any)
	first := arr[0].(map[string]any)
	if first["userName"] != "B***" {
		t.Errorf("nested list userName wrong: %v", first["userName"])
	}
	if first["sn"] != "S1" {
		t.Errorf("sn should be preserved in list: %v", first["sn"])
	}
	nested := first["nested"].(map[string]any)
	if nested["phone"] != maskFull() {
		t.Errorf("deeply nested phone should be masked: %v", nested["phone"])
	}
	if nested["keep"] != "ok" {
		t.Errorf("non-sensitive nested field should be preserved: %v", nested["keep"])
	}
	if m["total"].(float64) != 2 {
		t.Errorf("total wrong: %v", m["total"])
	}
}

func TestMaskJSONBytes_SensitiveSubtreeRedacted(t *testing.T) {
	// 敏感 key 对应对象/数组时，整个子树被脱敏
	in := []byte(`{"address": {"street": "secret", "no": 42}, "list_phone": ["a","b"]}`)
	// 注意 list_phone 不在规则表，仅测 address 子树
	out, err := MaskJSONBytes(in)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	var m map[string]any
	if err := json.Unmarshal(out, &m); err != nil {
		t.Fatalf("invalid JSON: %v\n%s", err, out)
	}
	addr := m["address"].(map[string]any)
	if addr["street"] != maskFull() {
		t.Errorf("address.street should be masked: %v", addr["street"])
	}
	// number 子字段被替换为掩码字符串
	if addr["no"] != maskFull() {
		t.Errorf("address.no should be masked: %v", addr["no"])
	}
}

func TestMaskJSONBytes_InvalidJSONFailsClosed(t *testing.T) {
	_, err := MaskJSONBytes([]byte(`{not valid json`))
	if err == nil {
		t.Fatal("expected error for invalid JSON (fail-closed)")
	}
}

func TestMaskJSONBytes_NullPreserved(t *testing.T) {
	in := []byte(`{"userName": null, "phone": null}`)
	out, err := MaskJSONBytes(in)
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	if !strings.Contains(string(out), `"userName":null`) {
		t.Errorf("null should be preserved: %s", out)
	}
}