package workbench

import (
	"context"
	"strings"
	"testing"

	"github.com/Tencent/WeKnora/internal/agent/token"
	"github.com/Tencent/WeKnora/internal/types"
)

type readingStub struct {
	budgetSearchStub
	overviewCalls int
}

func (s *readingStub) KnowledgeReadingOverview(context.Context, []string, []string, []types.TagScope) (*types.KnowledgeReadingOverview, error) {
	s.overviewCalls++
	return &types.KnowledgeReadingOverview{DocumentCount: 3, SummaryCount: 2, Fingerprint: "same", Entries: []types.KnowledgeReadingEntry{
		{Title: "第一份", Summary: "已有摘要"},
		{Title: "第二份", Summary: strings.Repeat("无法完整容纳的摘要", 2000)},
		{Title: "第三份"},
	}}, nil
}

func TestLayeredReadingCountsActualIncludedSummaries(t *testing.T) {
	s := NewService(nil)
	stub := &readingStub{}
	s.sessions = stub
	ctx, report := withReading(context.Background(), "full", &ReadingReport{Fingerprint: "same"})
	text, err := s.readingOverview(ctx, []string{"kb"}, nil, nil, 900)
	if err != nil {
		t.Fatal(err)
	}
	estimator, _ := token.NewEstimator()
	if report.Documents != 3 || report.AvailableSummaries != 2 || report.IncludedSummaries != 1 || report.IncludedDocuments != 3 || !report.ScopeUnchanged || estimator.EstimateString(text) > 300 {
		t.Fatalf("incorrect coverage: %+v", report)
	}
	quick, _ := withReading(context.Background(), "quick", report)
	if text, err = s.readingOverview(quick, []string{"kb"}, nil, nil, 900); err != nil || text != "" || stub.overviewCalls != 1 {
		t.Fatal("quick mode unexpectedly reads summaries")
	}
}

func TestAnchorMustMatchExactAcceptedSourceAndFingerprint(t *testing.T) {
	seg := Segment{ID: "a", Title: "基础", Content: "前文定义是准确引用，不可改写成另一种结论。"}
	anchor := ContinuityAnchor{SourceHash: segmentHash(seg), Quotes: []string{"前文定义是准确引用"}}
	if err := validateAnchor(anchor, seg); err != nil {
		t.Fatal(err)
	}
	anchor.Quotes = []string{"前文定义不准确"}
	if validateAnchor(anchor, seg) == nil {
		t.Fatal("invented anchor accepted")
	}
	anchor.Quotes = []string{"前文定义是准确引用"}
	seg.Content += "正文已变更。"
	if validateAnchor(anchor, seg) == nil {
		t.Fatal("stale anchor accepted")
	}
}

func TestLongContinuityReusesAnchorsAndOmitsSkippedText(t *testing.T) {
	s, ctx, _, model := generationHarness(t)
	d := TaskData{ContinuityAnchors: map[string]ContinuityAnchor{}, Segments: []Segment{
		{ID: "a", Title: "开头", Status: "accepted", Content: strings.Repeat("已讲授定义。", 2500)},
		{ID: "skip", Title: "废弃", Status: "skipped", Content: "禁止把此内容带入后文"},
		{ID: "b", Title: "中间", Status: "accepted", Content: "最近完整正文甲"},
		{ID: "c", Title: "承接", Status: "accepted", Content: "最近完整正文乙"},
		{ID: "d", Title: "本段", Status: "pending"},
	}}
	d.ContinuityAnchors["a"] = ContinuityAnchor{SourceHash: segmentHash(d.Segments[0]), Quotes: []string{"已讲授定义。"}}
	text, err := s.continuityContext(ctx, &d, 4, "agent", false)
	if err != nil {
		t.Fatal(err)
	}
	if len(model.messages) != 0 || strings.Contains(text, "禁止把此内容") || !strings.Contains(text, "最近完整正文甲") || !strings.Contains(text, "最近完整正文乙") || !strings.Contains(text, "已讲授定义") || !strings.Contains(text, "最后一个未跳过") {
		t.Fatal("anchor reuse or accepted context violated")
	}
}

func TestSegmentPositionExcludesSkippedOutlinesAndClosesFinalSegment(t *testing.T) {
	d := TaskData{Segments: []Segment{{Title: "当前"}, {Title: "废弃主题", Status: "skipped"}, {Title: "预算落地", Status: "pending"}}}
	first := segmentPosition(d, 0)
	if strings.Contains(first, "废弃主题") || !strings.Contains(first, "预算落地") || strings.Contains(first, "最后一个未跳过") {
		t.Fatal("position included a skipped topic or closed too early")
	}
	d.Segments[2].Status = "skipped"
	if !strings.Contains(segmentPosition(d, 0), "最后一个未跳过") {
		t.Fatal("last retained segment lacks closure instructions")
	}
}

func TestManualCategorySelectionUsesAssignedAgentScope(t *testing.T) {
	options := []SegmentCategory{{ID: "a"}, {ID: "b"}}
	d := TaskData{CategoryOptions: options, AgentCategoryOptions: map[string][]SegmentCategory{"writer": {options[0]}, "editor": {options[1]}}, Segments: []Segment{{ID: "s1", AgentID: "writer", CategoryID: "a", Status: "accepted", Content: "正文"}}}
	r := &Resource{Kind: "content-tasks", Data: asJSON(d)}
	d.Segments[0].AgentID = "editor"
	if applyManualUpdate(r, resourceRequest{Name: "任务", Data: asJSON(d)}) == nil {
		t.Fatal("agent change retained a category outside its scope")
	}
	d.Segments[0].CategoryID = "b"
	if err := applyManualUpdate(r, resourceRequest{Name: "任务", Data: asJSON(d)}); err != nil {
		t.Fatal(err)
	}
	updated, _ := decode[TaskData](r.Data)
	if updated.Segments[0].Status != "needs_review" {
		t.Fatal("agent and category change did not invalidate approval")
	}
	updated.AgentCategoryOptions["editor"] = options
	if applyManualUpdate(r, resourceRequest{Name: "任务", Data: asJSON(updated)}) == nil {
		t.Fatal("client widened category scope")
	}
}

func TestCategoryRecommendationAndManualUpdateProtection(t *testing.T) {
	options := []SegmentCategory{{ID: "cat", KnowledgeBaseID: "kb", Name: "课程/基础"}}
	segments := []Segment{{ID: "a", AgentID: "agent", Status: "accepted", Content: "正文", RecommendedCategoryID: "cat", CategoryReason: "适合本段"}}
	if err := normalizeCategoryRecommendations(segments, options); err != nil || segments[0].CategoryID != "cat" {
		t.Fatal(err)
	}
	d := TaskData{CategoryOptions: options, Segments: segments}
	r := &Resource{Kind: "content-tasks", Data: asJSON(d), Status: "completed"}
	d.Segments[0].CategoryID = ""
	if err := applyManualUpdate(r, resourceRequest{Name: "任务", Data: asJSON(d)}); err != nil {
		t.Fatal(err)
	}
	got, _ := decode[TaskData](r.Data)
	if got.Segments[0].Status != "needs_review" {
		t.Fatal("category change kept stale approval")
	}
	got.CategoryOptions = append(got.CategoryOptions, SegmentCategory{ID: "foreign"})
	if applyManualUpdate(r, resourceRequest{Name: "任务", Data: asJSON(got)}) == nil {
		t.Fatal("client injected category")
	}
	segments[0].RecommendedCategoryID = "foreign"
	if normalizeCategoryRecommendations(segments, options) == nil {
		t.Fatal("unknown AI category accepted")
	}
}

func TestRecommendationsCannotBorrowAnotherAgentsCategory(t *testing.T) {
	options := map[string][]SegmentCategory{
		"writer": {{ID: "writing", KnowledgeBaseID: "kb-writing", Name: "写作"}},
		"editor": {{ID: "editing", KnowledgeBaseID: "kb-editing", Name: "编辑"}},
	}
	segments := []Segment{{AgentID: "writer", RecommendedCategoryID: "editing"}}
	if validateAgentCategoryRecommendations(segments, options) == nil {
		t.Fatal("recommendation expanded the assigned Agent's knowledge scope")
	}
	segments[0].RecommendedCategoryID = "writing"
	if err := validateAgentCategoryRecommendations(segments, options); err != nil {
		t.Fatal(err)
	}
	if normalizeCategoryRecommendations(segments, options["writer"]) == nil {
		t.Fatal("category recommendation accepted without a reason")
	}
	segments[0].CategoryReason = "本段讲解写作结构，与该分类匹配。"
	if err := normalizeCategoryRecommendations(segments, options["writer"]); err != nil {
		t.Fatal(err)
	}
}

func TestDecompositionDiscardsOnlyUnusableRecommendations(t *testing.T) {
	options := []SegmentCategory{{ID: "writing"}, {ID: "editing"}}
	byAgent := map[string][]SegmentCategory{"writer": {options[0]}, "editor": {options[1]}}
	segments := []Segment{
		{ID: "valid", AgentID: "writer", RecommendedCategoryID: "writing", CategoryReason: "Matches writing"},
		{ID: "other-agent", AgentID: "writer", RecommendedCategoryID: "editing", CategoryReason: "Wrong scope", CategoryID: "editing"},
		{ID: "invented", AgentID: "editor", RecommendedCategoryID: "unknown", CategoryReason: "Invented"},
		{ID: "missing-reason", AgentID: "writer", RecommendedCategoryID: "writing"},
		{ID: "none", AgentID: "writer"},
	}
	if err := reconcileCategoryRecommendations(segments, options, byAgent); err != nil {
		t.Fatal(err)
	}
	if segments[0].CategoryID != "writing" || segments[0].CategoryReason != "Matches writing" {
		t.Fatal("valid recommendation was changed")
	}
	for _, segment := range segments[1:] {
		if segment.RecommendedCategoryID != "" || segment.CategoryID != "" || !strings.Contains(segment.CategoryReason, "Agent") {
			t.Fatalf("invalid recommendation retained or fallback not explained: %s", segment.ID)
		}
	}
	if err := validateAgentCategoryRecommendations(segments, byAgent); err != nil {
		t.Fatal(err)
	}
}