package coordinator_test

import (
	"strings"
	"testing"

	"chat-agent-system/internal/coordinator"
	"chat-agent-system/internal/ctxmgr"
)

func TestRewriteQuery_Rules(t *testing.T) {
	proj := coordinator.ProjectInfo{
		Name:       "数字化高频超声手术吸引系统",
		DeviceType: "超声手术设备",
	}

	// -------------------------------------------------------------
	// 1. 测试 R1: 省略补全 (标准法规编号替换与谓语继承)
	// -------------------------------------------------------------
	stR1 := ctxmgr.NewConversationTaskState()
	stR1.TopicAnchor = "GB 9706.1 对超声设备适用吗"

	resR1 := coordinator.RewriteQuery("YY 0505 呢？", stR1, nil, proj)
	if !containsRule(resR1.AppliedRules, "ellipsis_completion") {
		t.Errorf("expected ellipsis_completion rule, got: %v", resR1.AppliedRules)
	}
	if !strings.Contains(resR1.Rewritten, "YY 0505") || !strings.Contains(resR1.Rewritten, "适用吗") {
		t.Errorf("expected YY 0505 predicate completion, got: %s", resR1.Rewritten)
	}

	// -------------------------------------------------------------
	// 2. 测试 R2: 代词消解 (Pronoun Resolution)
	// -------------------------------------------------------------
	stR2 := ctxmgr.NewConversationTaskState()
	stR2.TopicAnchor = "数字化高频超声手术吸引系统方案研讨"

	resR2 := coordinator.RewriteQuery("它的预算是多少？", stR2, nil, proj)
	if !containsRule(resR2.AppliedRules, "pronoun_resolution") {
		t.Errorf("expected pronoun_resolution rule, got: %v", resR2.AppliedRules)
	}
	if !strings.Contains(resR2.Rewritten, "数字化高频超声手术吸引系统") {
		t.Errorf("expected project entity in rewritten query, got: %s", resR2.Rewritten)
	}

	// -------------------------------------------------------------
	// 3. 测试 R3: 主题延续合并 (Topic Continuation)
	// -------------------------------------------------------------
	stR3 := ctxmgr.NewConversationTaskState()
	stR3.TopicAnchor = "对比A厂家和B厂家的换能器报价"

	resR3 := coordinator.RewriteQuery("再加上C厂家呢", stR3, nil, proj)
	if !containsRule(resR3.AppliedRules, "topic_continuation") {
		t.Errorf("expected topic_continuation rule, got: %v", resR3.AppliedRules)
	}
	if !strings.Contains(resR3.Rewritten, "对比A厂家和B厂家的换能器报价") || !strings.Contains(resR3.Rewritten, "再加上C厂家呢") {
		t.Errorf("expected merged continuation query, got: %s", resR3.Rewritten)
	}

	// -------------------------------------------------------------
	// 4. 测试 R4: 槽位上下文补全 (Slot Context Completion for SearchQuery)
	// -------------------------------------------------------------
	stR4 := ctxmgr.NewConversationTaskState()
	stR4.CollectedSlots["device_type"] = "超声手术设备"
	stR4.CollectedSlots["target_market"] = "国内+CE"

	resR4 := coordinator.RewriteQuery("型检周期一般多久", stR4, nil, proj)
	if !containsRule(resR4.AppliedRules, "slot_context_completion") {
		t.Errorf("expected slot_context_completion in search query, got: %v", resR4.AppliedRules)
	}
	if !strings.Contains(resR4.SearchQuery, "器械=超声手术设备") || !strings.Contains(resR4.SearchQuery, "市场=国内+CE") {
		t.Errorf("expected slot constraints in SearchQuery, got: %s", resR4.SearchQuery)
	}

	// -------------------------------------------------------------
	// 5. 独立完整长句：不应被改写
	// -------------------------------------------------------------
	stStand := ctxmgr.NewConversationTaskState()
	stStand.TopicAnchor = "前期沟通记录"
	fullSentence := "请帮我详细阐述一下压电换能器逆压电效应的物理机理以及在超声刀上的实现方式"
	resStand := coordinator.RewriteQuery(fullSentence, stStand, nil, proj)
	if resStand.Rewritten != fullSentence {
		t.Errorf("standalone full sentence should not be rewritten, got: %s", resStand.Rewritten)
	}
}

func containsRule(rules []string, target string) bool {
	for _, r := range rules {
		if r == target {
			return true
		}
	}
	return false
}